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, get_accs_per_pg_from_yaml, get_spectra_per_acc_from_yaml, is_with_spectra)

heterogeneous_annotation_string = INCONCLUSIVE_ANNOTATIONS['heterogenous']


def add_unclassified_lineages_for_acc_without_result(acc_to_lineages_dict,  n_lca_levels, pgid_to_accs_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_dict.keys():
        for acc in pgid_to_accs_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_dict):
    """
    :param task_map: result task map file
    :param n_lca_levels: number of levels
    :param pgid_to_accs_dict: key: int protein group number, value = list of associated 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:]
            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_dict)
    return acc_to_lineages_dict


def create_lca(task_map, pg_yaml, 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_yaml: protein group yaml, 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
            ...
    """
    lca_headline_elems = get_headline(task_map).split(sep)[-n_lca_levels:]
    with_spectra = is_with_spectra(pg_yaml)
    pgid_to_accs_dict = get_accs_per_pg_from_yaml(pg_yaml)
    accs_to_spectra_dict = get_spectra_per_acc_from_yaml(pg_yaml)
    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,
                      with_spectra, method, path_to_taxmap, is_taxonomic_annotation, threshold, ignore_unclassified,
                      minimum_number_of_annotations)
    lca.weight_and_write_lca(output_lca, output_lca_support)


