# == Native Modules
from os.path import abspath
import pickle
import subprocess
from pathlib import Path
import pandas as pd
# == Installed Modules
import yaml
# == Project Modules
from prog.medit_lib import (
    expand_pam,
    export_guides_by_editor,
    file_exists,
    group_guide_table,
    handle_offtarget_request,
    handle_shell_exception,
    launch_shell_cmd,
    offtarget_mode_formatting,
    project_file_path,
    set_export,
    write_yaml_to_file,
    combine_guide_tables
)


def offtarget_prediction(args, jobtag):
    # == Load Run Parameters values ==
    user_jobtag = args.user_jobtag
    dna_bulge = args.dna_bulge
    rna_bulge = args.rna_bulge
    max_mismatch = args.max_mismatch
    root_dir = abspath(args.output)
    db_path_full = f"{abspath(args.db_path)}/medit_database"
    editing_tool_request = args.select_editors
    # == Load SLURM-related values ==
    cluster_request = args.cluster_request
    ncores = args.ncores
    maxtime = args.maxtime
    parallel_processes = int(args.parallel_processes)
    dry_run = args.dry_run
    
    # === Mode alias Switch===
    #   == This takes into account that whenever 'fast' mode is activated, 
    #       this is actually stored in the standard file structure. 
    #       This happens because 'fast' is a variation of 'standard'
    mode_alias = {
        "fast": "standard",
        "standard": "standard",
        "vcf": "vcf"
    }

    # == Set export paths tied to the SMK pipeline ==
    config_dir_path = f"{root_dir}/config"

    if editing_tool_request:
        editing_tool_request = list(editing_tool_request.split(","))

    # == Set import/export paths for dynamic YAML files ==
    config_db_path = f"{db_path_full}/config_db/config_db.yaml"
    dynamic_config_off_path = f"{config_dir_path}/config_off_{jobtag}.yaml"  # OffTarget config generated in this program
    dynamic_config_guidepred_path = f"{config_dir_path}/config_{jobtag}.yaml"  # Main config generated by guide_prediction.py

    # == Check existence of the config file from guide_prediciton ==
    if not file_exists(dynamic_config_guidepred_path):
        raise (f"This program depends on the configuration file associated with the JOBTAG {jobtag}. "
               f"Such file wasn't found on this parent folder: {root_dir}. "
               f"Make sure to provide the same path as the OUTPUT argument given to 'medit guide_prediction' ")

    #   == Load template configuration files ==
    config_template_path = project_file_path("smk.config", "medit_offtarget.yaml")
    cluster_template_path = project_file_path("smk.config", "medit_cluster.yaml")

    # == Define dynamic SMK call variables ==
    cluster_smk_setup = ''
    smk_verbosity = [True]
    smk_run_triggers = ''
    dryrun_setup = ''

    #   == Check the dry run request
    if dry_run:
        dryrun_setup = '-n'
    if user_jobtag:
        smk_run_triggers = '--rerun-triggers "mtime"'
    #   == Check the request to run on a cluster
    if cluster_request:
        cluster_smk_setup = ('--cluster "sbatch -t {cluster.time} -n {cluster.cores}" '
                             f'--cluster-config  {cluster_template_path}')

    # == Define dynamic SMK call variable ==
    allowed_rules = ['']

    # ->=== CONFIG FILES IMPORT ===<-
    with open(config_db_path, 'r') as config_db_handle:
        config_db = yaml.safe_load(config_db_handle)
    with open(dynamic_config_guidepred_path, 'r') as dynamic_config_handle:
        dynamic_config_guidepred = yaml.safe_load(dynamic_config_handle)
    with open(config_template_path, 'r') as config_handle:
        config_template = yaml.safe_load(config_handle)
    with open(cluster_template_path, 'r') as cluster_handle:
        cluster_template = yaml.safe_load(cluster_handle)

    # === Import Variables from Configuration File ===
    run_name = str(dynamic_config_guidepred['run_name'])
    mode = str(dynamic_config_guidepred['processing_mode'])
    query_index = (dynamic_config_guidepred['query_index'])
    root_dir = str(dynamic_config_guidepred['output_directory'])

    # === mEdit offtarget prediction will process reference genome and alternate genomes ===
    # == Set internal variables for Off-target processing
    reference_genome = str(config_db['sequence_id'])
    offtarget_genomes = offtarget_mode_formatting(mode, reference_genome, dynamic_config_guidepred)

    # === Assess the run 'mode' in light of available resources
    if not dryrun_setup:
        mode_checkpoint, mode = handle_offtarget_request(mode, db_path_full, offtarget_genomes)
        if mode == 'fast':
            offtarget_genomes = offtarget_mode_formatting(mode, reference_genome, dynamic_config_guidepred)
            allowed_rules = ['--omit-from "set_consensus_fasta" '
                             '--omit-from "set_gscan_indices" '
                             '--omit-from "casoff_run_extended" ']
        if not mode_checkpoint:
            # == The user was consulted and decided NOT to proceed
            exit(0)

    # == Get available options of already run editors or base editors
    guide_search_params_path = f"{root_dir}/{mode_alias[mode]}/jobs/{run_name}/guide_prediction-{reference_genome}/dynamic_params/{query_index[0]}_guide_search_params.pkl"
    be_search_params_path = f"{root_dir}/{mode_alias[mode]}/jobs/{run_name}/guide_prediction-{reference_genome}/dynamic_params/{query_index[0]}_guide_be_search_params.pkl"

    with open(guide_search_params_path, 'rb') as file:
        guide_search_params= pickle.load(file)
    with open(be_search_params_path, 'rb') as file:
        be_search_params = pickle.load(file)

    for editor, stats in be_search_params.items():
        guide_search_params[editor] = stats[0]

    # == Setup core off-target variables
    pam_per_editor_dict = {}
    alt_pam_per_editor_dict = {}
    pam_is_first_per_editor_dict = {}
    guides_per_editor_path = ""
    genome_type_dict = {}

    for index in query_index:
        combined_guide_report = pd.DataFrame()
        for offtarget_genome, genome_type in offtarget_genomes:

            # set one time directories for inputs and summary reports
            guides_per_editor_path = str(f"{root_dir}/{mode_alias[mode]}/jobs/{run_name}/guide_prediction-{reference_genome}/offtarget_prediction/dynamic_params")
            # == Recover Guide Prediction filepath ==
            if genome_type == 'main_ref':
                # == Define path to Report tables
                guides_report_path = Path(f"{root_dir}/{mode_alias[mode]}/jobs/{run_name}/"
                                          f"guide_prediction-{reference_genome}/guides_report_ref/{index}_Guides_found.csv")
                be_report_path = Path(f"{root_dir}/{mode_alias[mode]}/jobs/{run_name}/"
                                      f"guide_prediction-{reference_genome}/guides_report_ref/{index}_BaseEditors_found.csv")
                # == Combine guides reports for endonucleases and BEs
                combined_guide_report = combine_guide_tables(combined_guide_report, guides_report_path, genome_type)
                combined_guide_report = combine_guide_tables(combined_guide_report, be_report_path, genome_type)

            # Account for alternate genomes
            if genome_type == 'extended':
                guides_diff_path = Path(f"{root_dir}/{mode_alias[mode]}/jobs/{run_name}/"
                                        f"guide_prediction-{reference_genome}/guides_report_{offtarget_genome}/{index}_Guide_differences.csv")
                # Check if diff guides is empty and concatenates if it is not
                combined_guide_report = combine_guide_tables(combined_guide_report, guides_diff_path, genome_type)
            genome_type_dict.setdefault(str(offtarget_genome), str(genome_type))
            # Group and export guides per editing tool
            grouped_diff_guide_dict = group_guide_table(combined_guide_report, editing_tool_request)
            editors_extracted_from_guides = export_guides_by_editor(grouped_diff_guide_dict, guides_per_editor_path+f"/{index}_")

            for editor in editors_extracted_from_guides:
                # configure pams in extended forms if there's an IUPAC that is not N
                pam_per_editor_dict.setdefault(editor, expand_pam(guide_search_params[editor][0])[0])
                alt_pam_per_editor_dict.setdefault(offtarget_genome, {}).setdefault(editor, expand_pam(guide_search_params[editor][0])[1])
                pam_is_first_per_editor_dict.setdefault(offtarget_genome, {}).setdefault(editor, "--start" if guide_search_params[editor][1] is True else " ")

    # === Export Variables to Configuration File ===
    config_template['guides_per_editor_path'] = guides_per_editor_path
    config_template['pam_per_editor_dict'] = pam_per_editor_dict
    config_template['alt_pam_per_editor_dict'] = alt_pam_per_editor_dict
    config_template['pam_is_first_per_editor_dict'] = pam_is_first_per_editor_dict
    config_template['offtarget_genomes'] = {str(tup[0]): str(tup[1]) for tup in offtarget_genomes}
    config_template['offtarget_extended'] = {str(tup[0]): str(tup[1]) for tup in offtarget_genomes if
                                             tup[1] == 'extended'}
    config_template['genome_types'] = genome_type_dict
    config_template['DNAbb'] = dna_bulge
    config_template['RNAbb'] = rna_bulge
    config_template['max_mismatch'] = max_mismatch
    # == Assign cluster options ==
    cluster_template['__default__']['cores'] = ncores
    cluster_template['__default__']['time'] = maxtime

    # === Write YAML configs to mEdit Root Directory ===
    write_yaml_to_file(config_template, dynamic_config_off_path)

    # === Invoke SMK Pipelines ===
    print("# == Calling Off-Target Prediction pipeline == #")
    for smk_setup_idx in range(len(allowed_rules)):
        try:
            # --> When cluster submission is switched on,
            smk_command = (f"snakemake "
                           f"--snakefile {project_file_path('smk.pipelines', 'offtarget_prediction.smk')} "
                           f"{smk_run_triggers} "
                           f"{allowed_rules[smk_setup_idx]} "
                           f"-j {parallel_processes} "
                           f"--conda-frontend 'mamba' "
                           f"{cluster_smk_setup} "
                           f"--configfile {config_db_path} "
                           f"{dynamic_config_guidepred_path} {dynamic_config_off_path} "
                           f"--use-conda "
                           f"--rerun-incomplete "
                           f"{dryrun_setup} ")
            shell_result = launch_shell_cmd(smk_command, smk_verbosity[smk_setup_idx])
            handle_shell_exception(shell_result, smk_command, verbose=True)
        except subprocess.CalledProcessError as e:
            print(f"Error: {e}")
        except ValueError:
            print(f"Process completed in a previous run. Moving to the next one...")
