import pandas as pd

"""
Create a taxonomically restricted taxdump.tar.gz

Requires: 
- list of taxids, that shall be covered by new taxdump.tar
- extracted names.dmp, nodes.dmp, merged.dmp
Output:
- files names.dmp.new, nodes.dmp.new, merged.dmp.new 

to create a working taxdump.tar.gz, these need to be renamed to match input file names (names.dmp, nodes.dmp, merged.dmp)
 and archived (tar -czvf taxdump.tar.gz names.dmp, nodes.dmp, merged.dmp)
"""


def build_dict_from_two_dmp_columns(dmp_file_path, key_column_id, value_column_id):
    columns_dict = {}
    with open(dmp_file_path) as in_handle:
        for line in in_handle:
            fields = [x.strip() for x in line.split('|')]
            key = int(fields[key_column_id])
            value = int(fields[value_column_id])
            columns_dict[key] = value
    return columns_dict


def build_child_to_parent_taxid_dict(nodes_dmp_file_path):
    taxid_column = 0
    parent_taxid_column = 1
    return build_dict_from_two_dmp_columns(nodes_dmp_file_path, taxid_column, parent_taxid_column)


def get_new_taxid(taxid, dict_oldid_to_newid):
    new_id = dict_oldid_to_newid[taxid]
    if new_id in dict_oldid_to_newid:
        get_new_taxid(new_id, dict_oldid_to_newid)
    else:
        return new_id


def replace_merged_taxids(set_taxids, merged_dmp_file_path):
    old_taxid_column = 0
    new_taxid_column = 1
    dict_oldid_to_newid = build_dict_from_two_dmp_columns(merged_dmp_file_path, old_taxid_column, new_taxid_column)
    new_set = set()
    for taxid in set_taxids:
        if taxid in dict_oldid_to_newid:
            new_taxid = get_new_taxid(taxid, dict_oldid_to_newid)
            new_set.add(new_taxid)
        else:
            new_set.add(taxid)
    return new_set


def get_parent_of_taxid(taxid, parents_dict):
    if taxid in parents_dict:
        return parents_dict[taxid]
    elif taxid in parents_dict.values():
        raise KeyError("taxid has no parent")
    else:
        raise ValueError("taxid '{taxid}' is neither present as child nor as parent in parents_dict."
                         .format(taxid=taxid))


def build_set_of_taxids_incl_parents(query_taxids, parents_dict):
    set_taxids_incl_parents = set()
    for orig_query in query_taxids:
        query = orig_query
        set_taxids_incl_parents.add(query)
        while True:
            try:
                parent = get_parent_of_taxid(query, parents_dict)
                if parent not in set_taxids_incl_parents:
                    set_taxids_incl_parents.add(parent)
                    query = parent
                else:
                    break
            except KeyError:
                break
    return set_taxids_incl_parents


def filter_file_for_taxids(file, out_file, taxids, field_ids_to_check):
    with open(file, 'r') as in_handle:
        with open(out_file, 'w') as out_handle:
            for line in in_handle:
                fields = [x.strip() for x in line.split('|')]
                fields_to_check = [int(fields[i]) for i in field_ids_to_check]
                line_has_taxid_of_hit = any([taxid_i in taxids for taxid_i in fields_to_check])
                if line_has_taxid_of_hit:
                    out_handle.write(line)


fields_ids_of_taxids = {
    'nodes.dmp': [0],
    'merged.dmp': [0, 1],
    'names.dmp': [0],
}


def build_me_a_new_taxdump():
    set_hit_taxids = set(pd.read_csv("nr_testdata_fix/nr_hits_taxids", header=None)[0])
    set_hit_taxids = replace_merged_taxids(set_hit_taxids, "nr_testdata_fix/taxdump_extracted/merged.dmp")
    parents_dict = build_child_to_parent_taxid_dict("nr_testdata_fix/taxdump_extracted/nodes.dmp")
    set_hit_taxids_incl_parents = build_set_of_taxids_incl_parents(set_hit_taxids, parents_dict)

    for file_to_filter in fields_ids_of_taxids:
        filter_file_for_taxids(
            file="nr_testdata_fix/taxdump_extracted/" + file_to_filter,
            out_file="nr_testdata_fix/taxdump_extracted/" + file_to_filter + '.new',
            taxids=set_hit_taxids_incl_parents,
            field_ids_to_check=fields_ids_of_taxids[file_to_filter]
        )
