from utils.vars import MAP_FILE_BASENAME
from utils.db_handling.db_access import get_db_object_for_dbtype_and_version
from utils.databases.superclasses.hmm_db_superclass import HmmDatabase
import gzip
import pandas as pd
import os

def get_make_foam_map_input_files(w):
    db_obj = get_db_object_for_dbtype_and_version("foam", w.db_version)
    map_resource_file = db_obj.get_map_resource_file(relative=True)
    hmmlib_gz_full_path = db_obj.get_hmmlib_file(relative=True)
    hmm_lib_basename_wo_gz = os.path.basename(os.path.splitext(hmmlib_gz_full_path)[0])
    db_dir = os.path.dirname(db_obj.get_yaml(relative=True))
    hmmlib_unzipped = os.path.join(db_dir, "processed", hmm_lib_basename_wo_gz) + ".unzipped"
    return [map_resource_file, hmmlib_unzipped]

def create_foam_map(ont_in, hmm_lib, map_out, missing_out):
    df = pd.read_csv(ont_in, sep='\t', names=['level1', 'level2', 'level3', 'level4', 'ko'], skiprows=1)
    with gzip.open(map_out, "wt") as map_handle, gzip.open(missing_out, "wt") as missing_handle:
        map_handle.write("acc\tlevel1\tlevel2\tlevel3\tlevel4\tko\n")
        for entry in HmmDatabase.iter_hmmlib(hmm_lib):
            ko = entry['NAME'].split(":")[1].split("_")[0]
            annots = df.loc[df.ko == ko]
            if len(annots.index) == 0:
                missing_handle.write(entry['ACC'] + "\n")
                continue
            for index, row in annots.iterrows():
                map_handle.write("\t".join([entry['ACC']] + [str(x) for x in row.values.tolist()]) + "\n")

rule make_foam_map:
    input:
        get_make_foam_map_input_files
    output:
        "foam/{db_version}/processed/"+MAP_FILE_BASENAME+".map.gz",
        "foam/{db_version}/processed/"+MAP_FILE_BASENAME+".unmapped.log.gz"
    message:
        "generating map {output[0]}"
    benchmark:
        "foam/{db_version}/"+MAP_FILE_BASENAME+".map.gz.benchmark.txt"
    resources:
        mem_mb=400
    version:
        "0.2"
    run:
        create_foam_map(input[0], input[1], output[0], output[1])
