from utils.lca_determination.determine_lca import create_lca
from utils.db_handling.db_access import DbAccessor, get_db_object_for_dbtype_and_version


def get_lca(data, heterog_label):
    data = list(zip(*data))
    heterog = False
    for t in range(len(data)):
        if heterog:
            data[t] = heterog_label
        elif len(set(data[t])) == 1:
            data[t] = data[t][0]
        else:
            data[t] = heterog_label
            heterog = True
    return "\t".join(data)


def get_n_lca_levels(db_type, taskid, result_map):
    return get_db_object_for_task(config['tasks'][int(taskid)]).get_no_of_lca_levels() if db_type != 'custom_map' else None


def get_lca_levels(result_map, db_type):
    """
    :param result_map: result maps columns [#acc, source, and the task specific annotation result columns] for all db_types
      with exception of the result map of the database "custom_map":
      this includes [#acc, source, hit, and the task specific annotation result columns]
    """
    with open(result_map, 'r') as custom_map_handle:
        column_names = custom_map_handle.readline().strip(" \n\r").split("\t")
    if db_type == "custom_map":
        return len(column_names) - 3
    else:
        return len(column_names) -2


def get_lca_task_input(w):
    dct = {
        "pg_db": 'pgs/protein_groups_db.sql',
        "result_map": f'tasks/{w.annot_type}_annot_by_{w.tool}_on_{w.db_type}.v{w.db_version}.task{w.taskid}.map'
    }
    if w.annot_type == "tax":
        dct["tax_map"] = db_setup_workflow(get_db_object_for_dbtype_and_version(db_type='ncbi_nr',db_version='newest',
            db_base_dir=config['db_base_dir']).get_tax2annot_map_file())
    return dct


rule lca_task:
    input:
        unpack(get_lca_task_input)
    output:
        'tasks/{annot_type}_annot_by_{tool}_on_{db_type}.v{db_version}.task{taskid}.lca',
        'tasks/{annot_type}_annot_by_{tool}_on_{db_type}.v{db_version}.task{taskid}.lcasupport'
    message:
        "inferring lowest common ancestor ..."
    params:
        threshold = float(config['lca']['threshold']),
        method = config['lca']['method'],
        ignore_unclassified = config['lca']['ignore_unclassified'],
        min_nb_annotations = config['lca']['minimum_number_of_annotations'],
        db_type = "{db_type}",
        db_base_dir = config['db_base_dir'],
        is_taxonomic_annotation = lambda wildcards: True if wildcards.annot_type == 'tax' else False,
        n_lca_levels = lambda wildcards, input: get_n_lca_levels(wildcards.db_type, wildcards.taskid, input.result_map)
    run:
        if params.n_lca_levels is None:
            params.n_lca_levels = get_lca_levels(input.result_map, params.db_type)

        create_lca(task_map=input.result_map, pg_db=input.pg_db, output_lca=output[0], output_lca_support=output[1],
            method=params.method, is_taxonomic_annotation=params.is_taxonomic_annotation,
            n_lca_levels=params.n_lca_levels, path_to_taxmap=input.get("tax_map"), threshold=params.threshold,
            ignore_unclassified=params.ignore_unclassified, minimum_number_of_annotations=params.min_nb_annotations)
