from collections import defaultdict
from utils.config import INCONCLUSIVE_ANNOTATIONS
from utils.lca_determination.weighted_lca_class import WeightedLCA
from utils.input_output import get_headline
from utils.protein_group_sql_db import ProteinGroupSqlDbInteractor

heterogeneous_annotation_string = INCONCLUSIVE_ANNOTATIONS['heterogenous']


def add_unclassified_lineages_for_acc_without_result(acc_to_lineages_dict,  n_lca_levels, pgid_to_accs_part_dict):
    """
    :param acc_to_lineages_dict
    :param n_lca_levels: int, number of levels
    :param pgid_to_accs_dict:
    :return acc_to_lineages_dict
    """
    for pgid in pgid_to_accs_part_dict.keys():
        for acc in pgid_to_accs_part_dict[pgid]:
            if acc not in acc_to_lineages_dict:
                acc_to_lineages_dict[acc] = ([INCONCLUSIVE_ANNOTATIONS['unclassified']] * n_lca_levels)
    return acc_to_lineages_dict


def read_task_map_into_dict(task_map, n_lca_levels, pgid_to_accs_part_dict, part_accs=None):
    """
    :param task_map: result task map file
    :param n_lca_levels: number of levels
    :param pgid_to_accs_part_dict
    :param part_accs: list of protein accessions
    :return acc_to_lineages_dict keys = accession, values [[lineage1][lineage2]..]
    """
    acc_to_lineages_dict = defaultdict(list)
    with open(task_map, "r") as handle:
        for line in handle:
            line = line.strip(' \n\r')
            if len(line) == 0:
                continue
            if line[0] == "#":
                continue
            fields = line.split("\t")
            query_accession = fields[0]
            lca_level_annotations = fields[-n_lca_levels:]
            if part_accs is None:
                acc_to_lineages_dict[query_accession].append(lca_level_annotations)
            else:
                if query_accession in part_accs:
                    acc_to_lineages_dict[query_accession].append(lca_level_annotations)
    acc_to_lineages_dict = add_unclassified_lineages_for_acc_without_result(acc_to_lineages_dict, n_lca_levels, pgid_to_accs_part_dict)
    return acc_to_lineages_dict


def flatten_to_set(lst):
    return {item for sublist in lst for item in sublist}


def get_chunks(lst_of_pgs, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst_of_pgs), n):
        yield lst_of_pgs[i:i + n]


def create_lca(task_map, pg_db, output_lca, output_lca_support, method, is_taxonomic_annotation, n_lca_levels, path_to_taxmap,
               threshold=None, ignore_unclassified=False, minimum_number_of_annotations=0, sep='\t'):
    """
    :param task_map: file containing results of task
    :param pg_db: path to protein group sql database, containing protein_group information (protein accs, spectra)
    :param output_lca: path to lca output
    :param output_lca_support:  path to lca_support output
    :param method: lca method ('lca' or 'democratic lca' given in config)
    :param is_taxonomic_annotation: true if taxonomic task else False
    :param n_lca_levels: number of levels for task
    :param path_to_taxmap: path to taxmap
    :param threshold: threshold for group_lca, float value given in config
    :param ignore_unclassified: True -> remove all complete unclassified lineages before analysis
    :param minimum_number_of_annotations: int minimal number of classified entries contained in lineage, remove lineages with less
    :param sep: separator (e.g.'\t') between columns in task_map
    # 2 files are written, the first contains the results of the lca determination for the specified task,
     e.g. #pg	superkingdom	phylum	class	order	family	genus	species
         0	Bacteria	Firmicutes	Bacilli	Lactobacillales	Streptococcaceae	Lactococcus	Lactococcus lactis
         1	various	various	various	various	various	various	various
         ...
    # the second contains the number of spectra or protein IDs (if spectra information are not available), which support
    the lca (in brackets if below threshold, or 0 if unclassified) in relation to possible support per group
    e.g. #pg	superkingdom	phylum	class	order	family	genus	species
            0	2/2	2/2	2/2	2/2	2/2	2/2	2/2
            1	(23)/28	(18)/28	(18)/28	(18)/28	(18)/28	(18)/28	(18)/28
            ...
    """
    # write header
    lca_headline_elems = get_headline(task_map).split(sep)[-n_lca_levels:]
    lca_headline = '\t'.join(lca_headline_elems)
    WeightedLCA.write_headers(output_lca, output_lca_support, lca_headline)
    pg_db_obj = ProteinGroupSqlDbInteractor(pg_db)
    with_spectra = pg_db_obj.is_with_spectra()

    if method=='lca':
        for pgid_range in list(get_chunks(range(0, pg_db_obj.get_number_of_pgs()), 1000)):
            pgid_to_accs_dict = pg_db_obj.get_acc_list_per_pg(pgids=list(pgid_range))
            part_accs = flatten_to_set(pgid_to_accs_dict.values())
            acc_to_annotations_dict = read_task_map_into_dict(task_map, n_lca_levels, pgid_to_accs_dict, part_accs)
            if with_spectra:
                accs_to_spectra_dict = pg_db_obj.get_acc_to_spectra_dict(pgids=list(pgid_range))
            else:
                accs_to_spectra_dict = {}
            lca = WeightedLCA(lca_headline_elems, acc_to_annotations_dict, pgid_to_accs_dict, accs_to_spectra_dict,
                              method, path_to_taxmap, is_taxonomic_annotation, threshold, ignore_unclassified,
                              minimum_number_of_annotations)
            pgid_to_lcas_dict, pgid_to_supports_dict = lca.weight_lca(list(pgid_range))
            WeightedLCA.write_part_lca_result(pgid_to_lcas_dict, output_lca)
            WeightedLCA.write_part_lca_result(pgid_to_supports_dict, output_lca_support)
    elif method == 'democratic_lca':
        pgid_to_accs_dict = pg_db_obj.get_acc_list_per_pg()
        if with_spectra:
            accs_to_spectra_dict = pg_db_obj.get_acc_to_spectra_dict()
        else:
            accs_to_spectra_dict = {}
        acc_to_annotations_dict = read_task_map_into_dict(task_map, n_lca_levels, pgid_to_accs_dict)
        lca = WeightedLCA(lca_headline_elems, acc_to_annotations_dict, pgid_to_accs_dict, accs_to_spectra_dict,
                          method, path_to_taxmap, is_taxonomic_annotation, threshold, ignore_unclassified,
                          minimum_number_of_annotations)
        for pgid_range in list(get_chunks(range(0, len(pgid_to_accs_dict)), 1000)):
            pgid_to_lcas_dict, pgid_to_supports_dict = lca.weight_lca(list(pgid_range))
            WeightedLCA.write_part_lca_result(pgid_to_lcas_dict, output_lca)
            WeightedLCA.write_part_lca_result(pgid_to_supports_dict, output_lca_support)
