# == 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,
	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}')

	# ->=== 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'])
	root_dir = str(dynamic_config_guidepred['output_directory'])
	# == mEdit offtarget prediction will only process one reference genome in this version ==
	sequence_id = str(dynamic_config_guidepred['sequence_id'][0])

	# == Set output paths ==
	guides_per_editor_path = str(
		f"{root_dir}/{mode}/jobs/{run_name}/guide_prediction-{sequence_id}/offtarget_prediction")

	# == Recover Guide Prediction filepath ==
	guides_report_path = Path(f"{root_dir}/{mode}/jobs/{run_name}/"
							  f"guide_prediction-{sequence_id}/guides_report_ref/Guides_found.csv")

	# == 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 = export_guides_by_editor(grouped_guide_dict, guides_per_editor_path)

	# === Export Variables to Configuration File ===
	config_template['guides_per_editor_path'] = guides_per_editor_path
	config_template['editors_list'] = editors_list
	config_template['tmp_processing_casoff'] = f"{root_dir}/{config_template['tmp_processing_casoff']}"

	# == 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 == #")
	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"-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
						 )
	except subprocess.CalledProcessError as e:
		print(f"Error: {e}")
	except ValueError:
		print(f"Process completed in a previous run. Moving to the next one...")
