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


def offtarget_prediction(args, jobtag):
	# == Load Run Parameters values ==
	user_jobtag = args.user_jobtag
	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 ==
	ncores = args.ncores
	maxtime = args.maxtime
	parallel_processes = args.parallel_processes
	dry_run = args.dry_run

	# == 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")
	config_cluster_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 parallel_processes:
		# --> Upon SLURM run request, the guide_prediction.smk is split in two separate runs
		# --> That's because samtool's conda package crashes on a libcrypto error when
		#       it's deployed by snakemake on a SLURM node
		cluster_smk_setup = ('--cluster "sbatch -t {cluster.time} -n {cluster.cores}" '
							 f'--cluster-config  {config_cluster_path}')
	# == Define dynamic SMK call variable ==
	allowed_rules = ['']

	# ->=== CONFIG FILES IMPORT ===<-
	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)

	# === Import Variables from Configuration File ===
	run_name = str(dynamic_config_guidepred['run_name'])
	mode = str(dynamic_config_guidepred['processing_mode'])
	query_index = list(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
	offtarget_genomes = offtarget_mode_formatting(mode, dynamic_config_guidepred)
	reference_genome = str(dynamic_config_guidepred['sequence_id'][0])

	# === Adjust SMK run based on processing_mode
	if mode == 'fast':
		allowed_rules = ['--omit-from symlink_genomes']

	# == Setup core off-target variables
	editors_list = []
	guides_per_editor_path = ''
	for index in query_index:
		for offtarget_genome, genome_type in offtarget_genomes:
			# == Set output paths ==
			guides_per_editor_path = str(
				f"{root_dir}/{mode}/jobs/{run_name}/guide_prediction-{reference_genome}/offtarget_prediction/{offtarget_genome}/{index}_")

			if genome_type == 'main_ref':
				# == Recover Guide Prediction filepath ==
				guides_report_path = Path(f"{root_dir}/{mode}/jobs/{run_name}/"
										  f"guide_prediction-{reference_genome}/guides_report_ref/{index}_Guides_found.csv")
				# TODO: Add BE report
				# == Group guides by editor and export DF as pickles by editor
				grouped_guide_dict = group_guide_table(guides_report_path, editing_tool_request)
				editors_list.extend(export_guides_by_editor(grouped_guide_dict, guides_per_editor_path))

			# Account for alternate genomes
			if genome_type == 'extended':
				guides_diff_path = Path(f"{root_dir}/{mode}/jobs/{run_name}/"
										f"guide_prediction-{reference_genome}/guides_report_{offtarget_genome}/{index}_Guide_differences.csv")
				grouped_diff_guide_dict = group_guide_table(guides_diff_path, editing_tool_request)
				editors_list.extend(export_guides_by_editor(grouped_diff_guide_dict, guides_per_editor_path))

	# === Export Variables to Configuration File ===
	# NOTE: This whole block (until the end) was inside the index loop
	config_template['guides_per_editor_path'] = guides_per_editor_path
	config_template['editors_list'] = list(set(editors_list))
	config_template['tmp_processing_casoff'] = f"{root_dir}/{config_template['tmp_processing_casoff']}"
	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[0] == 'extended'}
	config_template['reference_id'] = reference_genome

	# == Set the temporary directory up for CasOffinder ==
	set_export(config_template['tmp_casoff_path'])

	# === 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,
			launch_shell_cmd(f"snakemake "
							 f"--snakefile {project_file_path('smk.pipelines', 'offtarget_prediction.smk')} "
							 f"{smk_run_triggers} "
							 f"{allowed_rules[smk_setup_idx]} "
							 f"-j {ncores} "
							 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} ",
							 smk_verbosity[smk_setup_idx]
							 )
		except subprocess.CalledProcessError as e:
			print(f"Error: {e}")
		except ValueError:
			print(f"Process completed in a previous run. Moving to the next one...")
