import os
import pandas as pd
import tarfile

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_tigrfam_map(role_name_in, role_link_in, info_tar_in, tigrfam_out):
    df_role_names = pd.read_csv(role_name_in, sep='\t', usecols=[1, 2, 3], names=['role_id', 'role_level', 'role_desc'])
    df_role_names['role_level'] = df_role_names['role_level'].str.rstrip(':')
    # index: role_id; columns: mainrole, sub1role
    df_role_names = df_role_names.pivot(index='role_id', columns='role_level', values='role_desc')
    df_role_links = pd.read_csv(role_link_in, sep='\t', usecols=[0, 1], names=['fam', 'role_id'])
    df_descr = build_df_fam_descr(info_tar_in)

    df_descr_id = df_descr.merge(df_role_links, how='left', on='fam')
    df_descr_id_roles = df_descr_id.merge(df_role_names, how='left', on='role_id')
    df_final = df_descr_id_roles.rename(
        columns={'fam': '#tigrfam', 'sub1role': 'subrole'}
    )
    df_final[['mainrole', 'subrole']] = \
        df_final[['mainrole', 'subrole']].fillna(INCONCLUSIVE_ANNOTATIONS['unclassified'])
    df_final['tigrfam'] = df_final['#tigrfam']

    df_final = df_final[['tigrfam', 'mainrole', 'subrole', 'descr' ]]
    df_final.columns = ['tigrfam ID', 'mainrole', 'subrole', 'descr' ]
    df_final.to_csv(tigrfam_out, sep='\t', index=False, compression='gzip')


def build_df_fam_descr(info_tar_in):
    """
    build dataframe with two columns:
        col1: name: 'fam', content: tigrfam names
        col2: name: 'descr', content: tigrfam descriptions
    :param info_tar_in: tigrfam_INFO.tar.gz file
    :return: df: columns: [fam, descr]
    """
    with tarfile.open(info_tar_in, "r:gz") as tar:
        tar_members = tar.getnames()
        lst_descr = []
        for member in tar_members:
            fam = ""
            descr = ""
            for line in tar.extractfile(member):
                if line.startswith(b"AC "):
                    fam = line[4:].decode().strip()
                elif line.startswith(b"DE "):
                    descr = line[4:].decode().strip()
                if fam and descr:
                    break
            lst_descr.append([fam, descr])
    df_descr = pd.DataFrame(lst_descr, columns=['fam', 'descr'])
    return df_descr


def get_tigrfam_map_resource_files(w):
    file_dict = {}
    db_obj = get_db_object_for_dbtype_and_version("tigrfams", w.version)
    file_dict["role_name_in"] = db_obj.get_role_name_file(relative=True)
    file_dict["role_link_in"] = db_obj.get_role_link_file(relative=True)
    file_dict["info_tar_in"] = db_obj.get_info_tar_file(relative=True)
    return file_dict


rule make_tigrfam_map:
    input:
        unpack(get_tigrfam_map_resource_files)
    output:
        "tigrfams/{version}/processed/"+MAP_FILE_BASENAME+".map.gz"
    message:
        "generating map {output[0]}"
    benchmark:
        "tigrfams/{version}/processed/"+MAP_FILE_BASENAME+".map.gz.benchmark.txt"
    resources:
        mem_mb=400
    version:
        "0.1"
    run:
        create_tigrfam_map(input.role_name_in, input.role_link_in, input.info_tar_in, output[0])
