import tarfile
import re
import gzip
from collections import defaultdict

from utils.db_handling.db_access import get_db_object_for_dbtype_and_version


def get_taxdump_tar_gz(w):
    db_obj = get_db_object_for_dbtype_and_version("ncbi_taxdump", w.version)
    return db_obj.get_taxdump_tar_gz_file(relative=True)

rule make_nr_tax2annot_map:
    input:
        get_taxdump_tar_gz
    output:
        "ncbi_nr/{version}/processed/tax2annot.map.gz"
    message:
        "generating taxmap {output[0]}"
    benchmark:
        "ncbi_nr/{version}/processed/tax2annot.map.gz.benchmark.txt"
    resources:
        mem_mb=1024*3
    version:
        "0.1"
    params: taxlevel=config['taxlevel']
    run:
        taxlevel = params.taxlevel
        taxlevel_count = len(taxlevel)
        ranks = {}
        parents = {}
        names = {}
        merged = defaultdict(set)
        with tarfile.open(input[0], "r:gz") as tar:
            names_handle = tar.extractfile("names.dmp")
            merged_handle = tar.extractfile("merged.dmp")
            nodes_handle = tar.extractfile("nodes.dmp")

            for line in nodes_handle:
                fields = [x.strip() for x in line.decode(tarfile.ENCODING).split('|')]
                taxid = fields[0]
                parent_taxid = fields[1]
                rank = fields[2]    # superkingdom, kingdom, ...
                ranks[taxid] = rank
                parents[taxid] = parent_taxid

            for line in merged_handle:
                fields = [x.strip() for x in line.decode(tarfile.ENCODING).split('|')]
                old_tax_id = fields[0]
                new_tax_id = fields[1]
                merged[new_tax_id].add(old_tax_id)

            for line in names_handle:
                fields = [x.strip() for x in line.decode(tarfile.ENCODING).split('|')]
                if fields[3] == "scientific name":
                    taxid = fields[0]
                    if fields[2] != "":
                        unique_name = fields[2]
                    else:
                        unique_name = fields[1]
                    names[taxid] = unique_name
        tags_regexp = re.compile('<[^>]*>')
        taxmap_out = ["#taxid\ttaxids_lineage\t" + "\t".join(taxlevel)]
    #with gzip.open(output[0], 'wt') as handle:
        #handle.write("#taxid\ttaxids_lineage\t" + "\t".join(taxlevel) + "\n")
        for taxid in set(list(parents.keys()) + list(parents.values())):
            taxids = ['-'] * taxlevel_count
            lineage = ['-'] * taxlevel_count
            orig_taxid = taxid
            while True:
                if ranks[taxid] in taxlevel:
                    p = taxlevel.index(ranks[taxid])
                    taxids[p] = taxid
                    lineage[p] = tags_regexp.sub('', names[taxid]).strip()
                if taxid == parents[taxid]:
                    break
                taxid = parents[taxid]
            taxids = ";".join(taxids)
            lineage = "\t".join(lineage)
            for taxid in [orig_taxid] + list(merged[orig_taxid]):
                taxmap_out.append(taxid + '\t' + taxids + '\t' + lineage)
                #handle.write(taxid + '\t' + taxids + '\t' + lineage + '\n')
        with gzip.open(output[0], 'wt') as outhandle:
            outhandle.write('\n'.join(taxmap_out) + '\n')
