import os
import re
import yaml

from utils.create_best_hits import best_hit_blastp
from utils.create_task_map import build_db_specific_headline, get_hit_dict, \
    get_ncbi_nr_taxinfo_dict, get_uniprot_taxinfo_dict, write_tax_task_map
from utils.helper_funcs_for_snakemake_workflow import get_task_best_hits_file

rule build_run_diamond_command:
  input:
    query_fasta='seqs/missing_taxa.faa',
    # db_yaml=expand("{db_dir}/{{db_type}}/{{db_version}}/db_config.yaml", db_dir=config['db_base_dir']),
    # dmnd_db=expand("{db_dir}/{{db_type}}/{{db_version}}/processed/diamond_db.dmnd", db_dir=config['db_base_dir'])
    db_yaml=db_setup_workflow("{db_type}/{db_version}/db_config.yaml"),
    dmnd_db=db_setup_workflow("{db_type}/{db_version}/processed/diamond_db.dmnd")
  output:
    cmd_file='tasks/{annot_type}_annot_by_diamond_on_{db_type}.v{db_version}.task{taskid}.result-cmd.txt',
    yaml='tasks/{annot_type}_annot_by_diamond_on_{db_type}.v{db_version}.task{taskid}.yaml'
  version:
    "0.1"
  log:
    'tasks/{annot_type}_annot_by_diamond_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.dmnd_db
    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 = ""

    else:
        task['params']['outfmt'] = '6 qseqid sallseqid qlen evalue bitscore length nident'
        task['params']['out'] = annot_result
        task['params']['query'] = query_file
        task['params']['db'] = prog_db_string
        task['params']['threads'] = threads
        task['params']['top'] = 0

        cmd = task['prog'] + " " + " ".join([f"--{x[0]} {x[1]}" for x in task['params'].items()])
        err_cmd = f"2> {log_file}"
        cmd += " " + err_cmd

    # 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_diamond:
  input:
    cmd_file='tasks/{annot_type}_annot_by_diamond_on_{db_type}.v{db_version}.task{taskid}.result-cmd.txt'
  output:
    'tasks/{annot_type}_annot_by_diamond_on_{db_type}.v{db_version}.task{taskid}.result'
  message:
    "performing task {wildcards.taskid} using diamond..."
  conda:
    "../../envs/diamond.yaml"
  log:
    'tasks/{annot_type}_annot_by_diamond_on_{db_type}.v{db_version}.task{taskid}.log'
  benchmark:
    "tasks/{annot_type}_annot_by_diamond_on_{db_type}.v{db_version}.task{taskid}.result.benchmark.txt"
  resources:
    mem_mb=1024*14
  version:
    "0.1"
  threads:
    32
  shell:
    """
    bash {input.cmd_file}
    """

rule get_besthits_diamond:
    input:
        'tasks/{annot_type}_annot_by_diamond_on_{db_type}.v{db_version}.task{taskid}.result'
    output:
        'tasks/{annot_type}_annot_by_diamond_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)]
        db_obj = get_db_object_for_task(task)
        out = []
        if "query-cover" in task:
            qc = task["query-cover"]
        elif "query-cover" in task['params']:
            qc = task['params']['query-cover']
        else:
            qc = 0
        for acc, hits in best_hit_blastp(input[0], qc, db_obj).items():
            for hit in hits:
                out.append(acc + '\t' + hit)
        with open(output[0], "w") as handle:
            handle.write("\n".join(out))

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)
    db_type = db_obj.get_type()
    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())
    # distinguish between nr and uniprot db to append necessary input files for lineage reconstruction
    # TODO this is extremely error prone if db.yaml is not generated automatically
    if db_type in ["ncbi_nr", "uniprot_sp", "uniprot_tr", "uniprot_complete"]:
        file_dct["db_taxmap"] = db_setup_workflow(db_obj.get_tax2annot_map_file())
    file_dct["job_taxmap"] = 'tax/taxmap.txt'

    return file_dct


rule map_diamond_task_result:
    input:
        unpack(input_map_task_result)
    output:
        'tasks/{annot_type}_annot_by_diamond_on_{db_type}.v{db_version}.task{taskid}.map'
    message:
        "getting lineage information for best hits"
    benchmark:
        "tasks/{annot_type}_annot_by_diamond_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_obj = get_db_object_for_task(task_dict)
        db_type = db_obj.get_type()
        taxlevels = params.taxlevels
        task_name = task_dict['shortname']
        unclassified_annotation_string = INCONCLUSIVE_ANNOTATIONS['unclassified']
        db_is_uniprot_or_ncbi_based = (db_type in ["ncbi_nr", "uniprot_sp", "uniprot_tr", "uniprot_complete"])
        task_best_hits_file = input["best_hits_file"]
        db_acc_to_x_map = input["db_map"]
        input_job_taxmap = input["job_taxmap"]
        if db_is_uniprot_or_ncbi_based:
            db_tax2annot_map = input["db_taxmap"]
            map_containing_annot_level_names = db_tax2annot_map
        else:
            db_tax2annot_map = None
            map_containing_annot_level_names = db_acc_to_x_map

        # "#acc\tsource\thit\tannotation_info[0]\tannotation_info[1]\t...\n"
        headline = build_db_specific_headline(map_containing_annot_level_names, db_type, taxlevels)
        #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)
        taxinfo_dict = None
        job_tax_map = input_job_taxmap
        if db_is_uniprot_or_ncbi_based:
            # reconstruct lineage for nr db
            if db_type == "ncbi_nr":
                taxinfo_dict = get_ncbi_nr_taxinfo_dict(db_tax2annot_map)
            # reconstruct lineage for uniprot db
            elif db_type.startswith("uniprot_"):
                taxinfo_dict = get_uniprot_taxinfo_dict(db_tax2annot_map, taxlevels)
            else:
                raise ValueError('This can not be reached')
        # build a tsv with columns:
        # [#acc, source, hit, annotation_info[0], annotation_info[1], ...]
        write_tax_task_map(output[0], headline, job_tax_map, db_acc_to_x_map, hits_dict, taxinfo_dict,
            unclassified_annotation_string, task_name, task_no, db_is_uniprot_or_ncbi_based)