import re
from pathlib import Path

from utils.emapper_utils import create_emapper_command
from utils.create_best_hits import best_hit_emapper
from utils.create_task_map import build_db_specific_headline, get_hit_dict, write_func_task_map_for_emapper_v1
from utils.helper_funcs_for_snakemake_workflow import get_task_best_hits_file
from utils.config import INCONCLUSIVE_ANNOTATIONS

rule build_run_emapper_command:
  input:
    query_fasta='seqs/all.faa',
    db_yaml=db_setup_workflow("eggnog/{db_version}/db_config.yaml"),
    emapper_data=db_setup_workflow("eggnog/{db_version}/emapper_dl/data"),
  output:
    cmd_file='tasks/{annot_type}_annot_by_emapper_v{ver}_on_{db_type}.v{db_version}.task{taskid}.result-cmd.txt',
    yaml='tasks/{annot_type}_annot_by_emapper_v{ver}_on_{db_type}.v{db_version}.task{taskid}.yaml'
  version:
    "0.1"
  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_dict = config['tasks'][task_no]
    db_obj = get_db_object_for_task(task_dict)
    task_dict['db_type'] = db_obj.get_db_type()
    query_file = input['query_fasta']
    prog_db_string = input.emapper_data
    result_file_name = Path(output.yaml).stem + '.result'
    output_directory = Path(output.yaml).parent
    annot_result = re.sub(r"-cmd\.txt$", "", output.cmd_file)
    create_emapper_command(output.cmd_file, output.yaml, task_dict, query_file, prog_db_string, result_file_name,
      output_directory, threads)

rule run_emapper:
  input:
    cmd_file='tasks/{annot_type}_annot_by_emapper_v{ver}_on_{db_type}.v{db_version}.task{taskid}.result-cmd.txt'
  output:
    'tasks/{annot_type}_annot_by_emapper_v{ver}_on_{db_type}.v{db_version}.task{taskid}.result'
  message:
    "performing task {wildcards.taskid} using emapper..."
  conda:
    "../../envs/emapper_v{ver}.yaml"
  log:
    'tasks/{annot_type}_annot_by_emapper_v{ver}_on_{db_type}.v{db_version}.task{taskid}.log'
  benchmark:
    "tasks/{annot_type}_annot_by_emapper_v{ver}_on_{db_type}.v{db_version}.task{taskid}.result.benchmark.txt"
  resources:
    mem_mb=1024*14
  version:
    "0.1"
  threads:
    32
  shell:
    """
    # execute command file
    bash {input.cmd_file} > {log}
    # cp eggnog result: output[0] + ".emapper.annotations" -> output[0]
    cp "{output[0]}.emapper.annotations" {output[0]}
    """


rule get_emapper_v1_besthits:
  input:
    'tasks/{annot_type}_annot_by_emapper_v1_on_eggnog_v4.v{db_version}.task{taskid}.result'
  output:
    'tasks/{annot_type}_annot_by_emapper_v1_on_eggnog_v4.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_emapper(input[0], task['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_emapper_v1_task_result:
  input:
    unpack(input_map_task_result)
  output:
    'tasks/{annot_type}_annot_by_emapper_v1_on_eggnog_v4.v{db_version}.task{taskid}.map'
  message:
    "getting lineage information for best hits"
  benchmark:
    "tasks/{annot_type}_annot_by_emapper_v1_on_eggnog_v4.v{db_version}.task{taskid}.map.benchmark.txt"
  resources:
    mem_mb=1024*2
  version:
    "0.1"
  run:
    #get task and db infos
    task_no = int(wildcards.taskid)
    task_dict = config['tasks'][task_no]
    db_obj = get_db_object_for_task(task_dict)
    db_type = db_obj.get_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"]
    db_tax2annot_map = None

    # "#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_for_emapper_v1(output[0], headline, map_containing_annot_level_names, hits_dict,
      unclassified_annotation_string, task_name, task_no)


rule map_emapper_v2_task_result:
  input:
    emapper_result="tasks/{annot_type}_annot_by_emapper_v2_on_eggnog_v5.v{db_version}.task{taskid}.result"
  output:
    'tasks/{annot_type}_annot_by_emapper_v2_on_eggnog_v5.v{db_version}.task{taskid}.map'
  message:
    "getting lineage information for best hits"
  benchmark:
    "tasks/{annot_type}_annot_by_emapper_v2_on_eggnog_v5.v{db_version}.task{taskid}.map.benchmark.txt"
  resources:
    mem_mb=1024*2
  version:
    "0.1"
  run:
    unclassified_annotation_string = INCONCLUSIVE_ANNOTATIONS['unclassified']
    #get task and db infos
    task_no = int(wildcards.taskid)
    task_dict = config['tasks'][task_no]
    db_obj = get_db_object_for_task(task_dict)
    task_name = task_dict['shortname']
    source_task_string = f"{task_name} [task {task_no}]"
    # build a tsv with headline columns:
    db_obj.write_task_map(output[0],input["emapper_result"],source_task_string,
      unclassified_annotation_string)