from utils.vars import MAP_FILE_BASENAME
from utils.db_handling.db_access import get_db_object_for_dbtype_and_version

import pandas as pd
import gzip

def get_make_resfam_map_input_files(w):
    db_obj = get_db_object_for_dbtype_and_version(f"resfams_{w.branch}", w.db_version)
    return db_obj.get_map_resource_file(relative=True)

def create_resfam_map(xls_in, map_out):
    df = pd.read_excel(xls_in, skiprows=3, usecols=['ResfamID', 'Resfam Family Name', 'Description', 'Mechanism Classification', 'Antibiotic Classification (Resfam Only)'],
        engine='openpyxl',)
    with gzip.open(map_out, "wt") as map_handle:
        map_handle.write("resfam id\tmechanistic class\tantibiotic class\tresfam family\tdescr\n")
        for index, row in df.iterrows():
            data = [str(x) for x in row.values.tolist()]
            map_handle.write("\t".join([data[0], data[3], data[4], data[1], data[2]]) + "\n")

rule make_resfam_map:
    input:
        get_make_resfam_map_input_files
    output:
        "resfams_{branch,[^/]+}/{db_version,[^/]+}/processed/"+MAP_FILE_BASENAME+".map.gz",
    message:
        "generating map {output[0]}"
    benchmark:
        "resfams_{branch}/{db_version}/"+MAP_FILE_BASENAME+".map.gz.benchmark.txt"
    resources:
        mem_mb=400
    version:
        "0.1"
    run:
        create_resfam_map(input[0], output[0])
