import os

from utils.db_handling.db_access import get_db_object_for_dbtype_and_version


def input_filter_idmapping_dat(wildcards):
    db_obj = get_db_object_for_dbtype_and_version("uniprot_tax", wildcards.db_version)
    idmap = db_obj.get_idmap()
    return os.path.relpath(idmap, os.getcwd())


rule filter_idmapping_dat:
    input:
        idmap=input_filter_idmapping_dat
    output:
        temp("uniprot_tax/{db_version}/processed/idmapping_filtered_for_ncbi_taxids.dat.gz")
    message:
        "filtering file for NCBI TaxID entries:\n\t{input.idmap}"
    shell:
        '''
        zcat {input.idmap} | grep -P '\tNCBI_TaxID\t' | cut -f1,3 | gzip > {output[0]}
        '''


def input_make_uniprot_taxmap(wildcards):
    db_obj = get_db_object_for_dbtype_and_version("uniprot_tax", wildcards.db_version)
    taxmap = db_obj.get_taxmap_txt(relative=True)
    idmap_filtered = os.path.join(
        "uniprot_tax",
        wildcards.db_version,
        "processed",
        "idmapping_filtered_for_ncbi_taxids.dat.gz")
    return [taxmap, idmap_filtered]


rule make_uniprot_acc2tax_map:
    input:
        input_make_uniprot_taxmap
    output:
        "uniprot_tax/{db_version}/processed/acc2tax.map.gz",
        "uniprot_tax/{db_version}/processed/acc2tax.unmapped.log.gz"
    message:
        "generating taxmap {output[0]}"
    benchmark:
        "uniprot_tax/{db_version}/processed/acc2tax.map.gz.benchmark.txt"
    resources:
        mem_mb=1024*3
    version:
        "0.1"
    params:
        db_obj=lambda wildcards: get_db_object_for_dbtype_and_version("uniprot_tax", wildcards.db_version)
    run:
        params.db_obj.create_acc2tax_map(input[0], input[1], output[0], output[1])
