import gzip
from collections import defaultdict
from utils.input_output import iter_file


def get_hit_dict(task_best_hits_file, db_type):
    """
    read task_best_hits_file into dictionary
    :param task_best_hits_file: tsv-file, with 2 columns
        column 1: query accession
        column 2: list of target db hit accessions, separated by ';'
    :param db_type: type of the target database
    :return: defaultdict: key: target db hit accession; value: query accession
    """
    hit_to_accs_set = defaultdict(set)
    for line in iter_file(task_best_hits_file):
        fields = line.strip().split("\t")
        acc_query = fields[0]
        acc_hits = fields[1]
        for hit in acc_hits.split(";"):
            # disting. between nr and uniprot
            if db_type == "ncbi_nr":
                if hit.count("|") > 0:
                    continue
            elif db_type.startswith("uniprot_"):
                hit = hit.split("|")[-1]
            elif db_type == "pfams":
                hit = hit.split(".")[0]
            elif db_type == "dbcan":
                hit = hit.split('_')[0]
            hit_to_accs_set[hit].add(acc_query)
    return hit_to_accs_set


def get_ncbi_nr_taxinfo_dict(db_tax_map):
    """
    read db_tax_map into dict
    :param db_tax_map: ncbi nr taxmap, gzipped
    :return: dictionary:
        key: taxID
        value: string, containing tab-separated entries (zero based index):
            taxid, taxids_lineage, taxlevel[0], taxlevel[1], ...
    """
    taxinfo_dict = {}
    # db_tax_map = taxdump.map.gz
    with gzip.open(db_tax_map, 'rt') as handle:
        # dump headline
        _ =  handle.readline()[1:]
        for line in handle:
            line = line.strip()
            taxid = line.split("\t")[0]
            taxinfo_dict[taxid] = line
    return taxinfo_dict


def get_uniprot_taxinfo_dict(db_tax_map, taxlevels):
    """
    read db_tax_map into dict
    :param db_tax_map: uniprot_tax.txt
    :return: dictionary:
        key: taxID
        value: string, containing tab-separated entries (zero based index):
            taxid, taxids_lineage, taxlevel[0], taxlevel[1], ...
    """
    taxinfo_dict = {}
    taxlevel_count = len(taxlevels)
    # db_tax_map = uniprot_tax.txt
    ranks, parents, names = get_ranks_parents_names_dicts_for_taxids(db_tax_map)
    for taxid in names:
        taxids = ['-'] * taxlevel_count
        lineage = ['-'] * taxlevel_count
        orig_taxid = taxid
        while taxid != "":
            rank = ranks[taxid]
            if rank in taxlevels:
                taxlevel_index = taxlevels.index(rank)
                taxids[taxlevel_index] = taxid
                lineage[taxlevel_index] = names[taxid]
            taxid = parents[taxid]
        taxinfo_dict[orig_taxid] = orig_taxid + "\t" + ";".join(taxids) + "\t" + "\t".join(lineage)
    del names
    del parents
    del ranks
    return taxinfo_dict


def get_ranks_parents_names_dicts_for_taxids(db_tax_map):
    ranks = {}
    parents = {}
    names = {}
    with gzip.open(db_tax_map, 'rt') as handle:
        handle.readline()
        for line in handle:
            fields = line.strip().split("\t") + [''] * 9
            taxid = fields[0]
            name = fields[2]
            rank = fields[7].lower()
            parent = fields[9]

            names[taxid] = name
            parents[taxid] = parent
            ranks[taxid] = rank
    return ranks, parents, names


def build_ncbi_nr_taskmap_headline_stub(db_tax_map):
    with gzip.open(db_tax_map, 'rt') as handle:
        headline_stub =  handle.readline()[1:]
    return headline_stub


def build_uniprot_taskmap_headline_stub(taxlevels):
    headline_stub = "taxid\ttaxids_lineage\t" + "\t".join(taxlevels) + "\n"
    return headline_stub


def build_db_specific_headline(db_map_containing_annot_level_names, db_type, taxlevels=None):
    """
    Build headline of the form:
        '#acc\tsource\thit\t(tab-separated-annotation-level-names)\n'
    :param db_map_containing_annot_level_names:
    :param db_type:
    :param taxlevels:
    :return:
    """
    if db_type == "ncbi_nr":
        headline_stub = build_ncbi_nr_taskmap_headline_stub(db_map_containing_annot_level_names)
    elif db_type.startswith("uniprot_"):
        headline_stub = build_uniprot_taskmap_headline_stub(taxlevels)
    else:
        with gzip.open(db_map_containing_annot_level_names, 'rt') as db_map_handle:
            db_map_column_names = db_map_handle.readline().lstrip("#").strip().split("\t")
            headline_stub = "\t".join(db_map_column_names[1:]) + "\t" + db_map_column_names[0] + "\n"
    headline = "#acc\tsource\t" + headline_stub
    return headline


def write_into_task_map(acc, lst_annotations_raw, unclassified_annotation_string, hits, outhandle, task_name, task_no):
    lst_annotations = [unclassified_annotation_string if x == "-" else x for x in lst_annotations_raw]
    source = f"{task_name} [task {task_no}]"
    annot_results = '\t'.join(lst_annotations)
    for query in sorted(hits):
        outhandle.write(
            f"{query}\t{source}\t{annot_results}\t{acc}\n"
        )


def read_db_map_line(line):
    fields = line.strip().split("\t")
    target_db_accession = fields[0]
    target_db_annotations = fields[1:]
    return target_db_accession, target_db_annotations


def write_func_task_map_for_emapper_v1(task_map, headline, db_acc_to_x_map, hits_dict, unclassified_annotation_string,
                                       task_name, task_no):
    na = set()
    with open(task_map, 'w') as outhandle:
        outhandle.write(headline)
        with gzip.open(db_acc_to_x_map, 'rt') as db_map_handle:
            for line in db_map_handle:
                target_db_accession, target_db_annotations = read_db_map_line(line)
                search_keys = [f"{target_db_annotations[2]}.{target_db_accession}"]
                if target_db_accession not in na:
                    search_keys.append(f"NA.{target_db_accession}")
                    na.add(target_db_accession)
                for target in search_keys:
                    if target in hits_dict:
                        write_into_task_map(target_db_accession, target_db_annotations, unclassified_annotation_string,
                                            hits_dict[target], outhandle, task_name, task_no)


def write_func_task_map(task_map, headline, db_acc_to_x_map, hits_dict, unclassified_annotation_string, task_name,
                        task_no):
    with open(task_map, 'w') as outhandle:
        outhandle.write(headline)
        with gzip.open(db_acc_to_x_map, 'rt') as db_map_handle:
            for line in db_map_handle:
                target_db_accession, target_db_annotations = read_db_map_line(line)
                # accession number from map is among best hits
                if target_db_accession in hits_dict:
                    write_into_task_map(target_db_accession, target_db_annotations, unclassified_annotation_string, hits_dict[target_db_accession],
                                        outhandle, task_name, task_no)


def write_tax_task_map(task_map, headline, job_tax_map, db_acc_to_x_map, hits_dict, taxinfo_dict,
                       unclassified_annotation_string, task_name, task_no, db_is_uniprot_or_ncbi_based):
    # build a tsv with columns:
    # [#acc, source, hit, annotation_info[0], annotation_info[1], ...]
    with open(task_map, 'w') as outhandle:
        outhandle.write(headline)
        for line in open(job_tax_map):
            outhandle.write(line)
        with gzip.open(db_acc_to_x_map, 'rt') as db_map_handle:
            for line in db_map_handle:
                fields = line.strip().split("\t")
                target_db_accession = fields[0]
                # accession number from map is among best hits
                if target_db_accession in hits_dict:
                    if db_is_uniprot_or_ncbi_based:
                        taxid = fields[1]
                        lst_annotations_raw = taxinfo_dict[taxid].split("\t")
                    else:
                        lst_annotations_raw = fields[1:]
                    lst_annotations = [unclassified_annotation_string if x == "-" else x for x in lst_annotations_raw]
                    for query in hits_dict[target_db_accession]:
                        outhandle.write(
                            query + "\t" + task_name + " [task " + str(task_no) + "]\t" + "\t".join(lst_annotations) + "\n"
                        )
