import gzip
import io
import sys

from utils.vars import MAP_FILE_BASENAME


def input_make_eggnog_map(wildcards):
    db_obj = get_db_object_for_dbtype_and_version("eggnog_v4", wildcards.db_version)
    funcat_file = db_obj.get_fun_cats_file(relative=True)
    annotation_file = db_obj.get_og_annotations_file(relative=True)
    return [annotation_file, funcat_file]

rule make_eggnog_map:
    input:
        input_make_eggnog_map
    output:
        "eggnog/{db_version, 4.*}/processed/"+MAP_FILE_BASENAME+".map.gz"
    message:
        "generating eggnog version 4 database map ..."
    benchmark:
        "eggnog/{db_version}/processed/"+MAP_FILE_BASENAME+".map.gz.benchmark.txt"
    resources:
        mem_mb=1500
    run:
        annotation_file = input[0]
        funcat_file = input[1]
        #read eggnog funcats
        funcats = {}
        with open(funcat_file, "r") as handle:
            for line in handle:
                line = line.strip()
                if len(line) == 0:
                    continue
                elif line[0] != "[":
                    mainrole = line.lower()
                else:
                    key = line[1]
                    subrole = line[4:]
                    if key in funcats:
                        sys.exit("error: multiple functional roles for key ''" + key + "' in " + funcat_file)
                    funcats[key] = (mainrole, subrole)

        #read eggnog og annotations
        na = set()
        with gzip.open(annotation_file, 'rb') as gzhandle, \
                gzip.open(output[0], 'wt') as outhandle:
            inhandle = io.BufferedReader(gzhandle)
            outhandle.write("prot\tmain_role\tsub_role\tog\tdescr\n")
            for line in inhandle:
                line = line.decode().strip()
                if len(line) == 0:
                    continue
                fields = line.split("\t")
                og = fields[0]
                prots = fields[-1].split(",")
                descr = fields[3]
                fun_key = fields[4].split("'")[1]
                mainrole = funcats[fun_key][0]
                subrole = funcats[fun_key][1]
                for prot in prots:
                    outhandle.write(
                        prot + "\t" + mainrole + "\t" + subrole + "\t" + og + "\t" + descr + "\n")
