import os
import re
import yaml

from utils.create_best_hits import best_hit_hmmer3
from utils.create_task_map import build_db_specific_headline, get_hit_dict, write_func_task_map
from utils.helper_funcs_for_snakemake_workflow import get_task_best_hits_file

rule build_run_hmmer_command:
  input:
    query_fasta='seqs/all.faa',
    db_yaml=db_setup_workflow("{db_type}/{db_version}/db_config.yaml"),
    hmm_lib=db_setup_workflow("{db_type}/{db_version}/processed/hmmlib-{hmm_ver}"),
    preprocessed_dbfiles=db_setup_workflow(expand("{{db_type}}/{{db_version}}/processed/hmmlib-{{hmm_ver}}.{ext}",
                                ext=["h3f", "h3i", "h3m", "h3p"]))
  output:
    cmd_file='tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.result-cmd.txt',
    yaml='tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.yaml'
  version:
    "0.1"
  log:
    'tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.log'
  threads:
    32
  run:
    """
    Build the command for annotation of provided query file and write it to text file.
    Write out parameters to task specific yaml file.
    """
    task_no = int(wildcards.taskid)
    task = config['tasks'][task_no]
    query_file = input['query_fasta']
    prog_db_string = input.hmm_lib
    annot_result = re.sub(r"-cmd\.txt$", "", output.cmd_file)
    log_file = log[0]

    #no params
    if 'params' not in task:
        task['params'] = {}

    #empty query
    if os.path.getsize(query_file) == 0:
        for o in output:
            open(o, "w").close()
        cmd = ""

    #hmmer
    else:
        single_args = {'Z', 'T', 'E'}
        cmd = f"{task['prog']} --noali --notextw --cpu {threads} --tblout {annot_result} -o {log_file} "
        cmd += " ".join([f"--{x[0]} {x[1]}" if x[0] not in single_args else f"-{x[0]} {x[1]}" for x in task['params'].items()])
        cmd += f" {prog_db_string} {query_file}"

    # write cmd_file
    with open(output.cmd_file, "w") as handle:
        handle.write(cmd)
        handle.write("\n")

    #write yaml
    with open(output.yaml, "w") as handle:
        task = dict(task)
        if "params" in task:
            task['params'] = dict(task['params'])
        handle.write(yaml.dump(task))


rule run_hmmer:
  input:
    cmd_file='tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.result-cmd.txt'
  output:
    'tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.result'
  message:
    "performing task {wildcards.taskid} using hmmer ..."
  conda:
    "../../envs/hmmer-{hmm_ver}.yaml"
  log:
    'tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.log'
  benchmark:
    "tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.result.benchmark.txt"
  resources:
    mem_mb=800
  version:
    "0.1"
  threads:
    32
  shell:
    """
    bash {input.cmd_file}
    """


rule get_hmmer_besthits:
    input:
        'tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.result'
    output:
        'tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.best_hits'
    message:
        "collecting best hits of task {wildcards.taskid} ..."
    version:
        "0.1"
    threads:
        1
    run:
        task = config['tasks'][int(wildcards.taskid)]
        with open(output[0], "w") as handle:
            handle.write("\n".join([x[0] + "\t" + x[1] for x in best_hit_hmmer3(input[0], task['prog'],
                dbtype=wildcards.db_type).items()]))

def input_map_task_result(wildcards):
    file_dct = {}
    task = config['tasks'][int(wildcards.taskid)]
    db_base_dir = get_db_base_dir()
    db_obj = get_db_object_for_task(task)
    file_dct["best_hits_file"] = get_task_best_hits_file(wildcards.taskid, task, db_base_dir)
    file_dct["db_map"] = db_setup_workflow(db_obj.get_db_map_file())
    return file_dct


rule map_hmm_task_result:
    input:
        unpack(input_map_task_result)
    output:
        'tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.map'
    message:
        "getting lineage information for best hits"
    benchmark:
        "tasks/{annot_type}_annot_by_hmmer-{hmm_ver}_on_{db_type}.v{db_version}.task{taskid}.map.benchmark.txt"
    resources:
        mem_mb=1024*2
    version:
        "0.1"
    params: taxlevels = config['taxlevel']
    run:
        from utils.config import INCONCLUSIVE_ANNOTATIONS
        #get task and db infos
        task_no = int(wildcards.taskid)
        task_dict = config['tasks'][task_no]
        db_type = wildcards.db_type
        task_name = task_dict['shortname']
        unclassified_annotation_string = INCONCLUSIVE_ANNOTATIONS['unclassified']
        task_best_hits_file = input["best_hits_file"]
        map_containing_annot_level_names = input["db_map"]

        # "#acc\tsource\thit\tannotation_info[0]\tannotation_info[1]\t...\n"
        headline = build_db_specific_headline(map_containing_annot_level_names, db_type)
        #get hits; key: accession of target db; values: accessions from query fasta
        hits_dict = get_hit_dict(task_best_hits_file, db_type=db_type)
        # build a tsv with columns:
        # [#acc, source, hit, annotation_info[0], annotation_info[1], ...]
        write_func_task_map(output[0], headline, map_containing_annot_level_names, hits_dict, unclassified_annotation_string, task_name,
            task_no)
