import pandas as pd

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

def create_pfam_map(clan_in, pfam_in, pfam_out):
    df_clans = pd.read_csv(clan_in, sep='\t', usecols=[0, 1, 3], names=['clan', 'symbol', 'cdescr'])
    df_pfams = pd.read_csv(pfam_in, sep='\t', usecols=[0, 1, 4], names=['pfam', 'clan', 'pdescr'])

    # merge clan and pfam description on shared clan accession
    df_merged = df_pfams.merge(df_clans, on='clan', how='outer')
    # replace nan from missing clan accessions w/ default 'unclassified' values
    clan_unclassified = INCONCLUSIVE_ANNOTATIONS['unclassified']
    nan_values = {'clan': clan_unclassified, 'symbol': '-', 'cdescr': '-'}
    df_merged = df_merged.fillna(value=nan_values)
    # deleting trailing whiespaces in descr
    df_merged['pdescr'] = df_merged['pdescr'].str.strip()
    df_merged['cdescr'] = df_merged['cdescr'].str.strip()
    # sort by pfam accession and reorder columns
    df_merged = df_merged.sort_values(by='pfam')
    df_merged = df_merged[['pfam', 'clan', 'symbol', 'cdescr', 'pdescr']]
    # write out pfam.map.gz
    pfam_header = ['pfam ID', 'clan', 'clan_symbol', 'clan_descr', 'pfam_descr']
    df_merged.to_csv(pfam_out, sep='\t', header=pfam_header, index=False, compression='gzip')


def get_pfam_map_resource_files(w):
    file_dict = {}
    db_obj = get_db_object_for_dbtype_and_version("pfams", w.version)
    file_dict["clan_in"] = db_obj.get_clan_txt_file(relative=True)
    file_dict["pfam_in"] = db_obj.get_pfam_clans_file(relative=True)
    return file_dict


rule make_pfam_map:
    input:
        unpack(get_pfam_map_resource_files)
    output:
        pfam_out = "pfams/{version}/processed/"+MAP_FILE_BASENAME+".map.gz"
    # pfam_out = "{pfamsdb}.map.gz"
    message:
        "generating map {output[0]}"
    benchmark:
        "pfams/{version}/processed/"+MAP_FILE_BASENAME+".map.gz.benchmark.txt"
    resources:
        mem_mb=400
    run:
        create_pfam_map(input.clan_in, input.pfam_in, output.pfam_out)

