#!/bin/bash

pipline_repo_path=`pwd`
solvent_puffin_path=<SOLVENT-PUFFIN-PATH>
complex_puffin_path=<COMPLEX-PUFFIN-PATH>
system_name=mcl1-19G
solvent_config_file=$pipline_repo_path/qm-mm/input_files/single_points/${system_name}_solvent.yml
complex_config_file=$pipline_repo_path/qm-mm/input_files/single_points/${system_name}_complex_shrunk.yml
qmqm_solvent_config_file=$pipline_repo_path/qm-mm/input_files/qmqmmm_single_points/${system_name}_solvent.yml
qmqm_complex_config_file=$pipline_repo_path/qm-mm/input_files/qmqmmm_single_points/${system_name}_complex_shrunk.yml
n_puffins=25
# Run classical FEP calculation
yank_jid=$(sbatch $pipline_repo_path/fep/yank.sh)
yank_analyze_jid=$(sbatch --dependency=afterok:${yank_jid##* } $pipline_repo_path/fep/yank_analyze.sh)
yank_output_directory=`grep "output_dir" $pipline_repo_path/fep/yank.sh | awk '{print $2}'`

# Run extended end state sampling with yank
interacting_yank_yaml_file=$pipline_repo_path/fep/yank_interacting.yml
separated_yank_yaml_file=$pipline_repo_path/fep/yank_separated.yml
yank_endstate_interacting_jid=$(sbatch --export=yank_yaml_file=$interacting_yank_yaml_file $pipline_repo_path/fep/endstate.sh)
yank_endstate_separated_jid=$(sbatch --export=yank_yaml_file=$separated_yank_yaml_file $pipline_repo_path/fep/endstate.sh)
# Extract end state trajectories
yank_interacting_output_dir=`grep "output_dir" $interacting_yank_yaml_file | awk '{print $2}'`
yank_separated_output_dir=`grep "output_dir" $separated_yank_yaml_file | awk '{print $2}'`
trajectory_output_dir=yank_trajectories
yank_trajectories_jid=$(sbatch --dependency=afterok:${yank_endstate_interacting_jid##* },afterok:${yank_endstate_separated_jid##* } --export=interacting_output_dir=$yank_interacting_output_dir,separated_output_dir=$yank_separated_output_dir,output_directory=$trajectory_output_dir $pipline_repo_path/fep/yank_gettraj.sh)
solvent_trajectory=${trajectory_output_dir}/solvent_l1.xtc
complex_trajectory=${trajectory_output_dir}/complex_l1.xtc

# Upload solvent trajectory to the database
initial_upload_solvent_jid=$(sbatch --dependency=afterok:${yank_analyze_jid##* } --export=config_file=$solvent_config_file,trajectory_file=$solvent_trajectory $pipline_repo_path/NEQ_switching_ani2x/slurm/initialize_database.sh)
initial_upload_complex_jid=$(sbatch --dependency=afterok:${yank_analyze_jid##* } --export=config_file=$complex_config_file,trajectory_file=$complex_trajectory $pipline_repo_path/NEQ_switching_ani2x/slurm/initialize_database.sh)

# Run initial QM/MM calculations
mkdir -p $system_name
mkdir -p $system_name/solvent
mkdir -p $system_name/complex
initial_solvent_output_dir=$system_name/solvent/initial_upload
initial_complex_output_dir=$system_name/complex/initial_upload
initial_qmmm_solvent_jid=$(sbatch --dependency=afterok:${initial_upload_solvent_jid##* } --array=0-${n_puffins} -e ${initial_solvent_output_dir}/puff.err -o ${initial_solvent_output_dir}/puff.out $solvent_puffin_path)
initial_qmmm_complex_jid=$(sbatch --dependency=afterok:${initial_upload_complex_jid##* } --array=0-${n_puffins} -e ${initial_complex_output_dir}/puff.err -o ${initial_complex_output_dir}/puff.out $complex_puffin_path)

# Run active learning
cd $system_name/solvent
final_qmmm_solvent_jid=`$pipline_repo_path/NEQ_switching_ani2x/fully_automated/submit_multiple.sh ${system_name}_solvent $solvent_puffin_path $n_puffins $initial_qmmm_solvent_jid | grep "Final dependency:" | awk '{print $3}'`
cd ../complex
final_qmmm_complex_jid=`$pipline_repo_path/NEQ_switching_ani2x/fully_automated/submit_multiple.sh ${system_name}_complex_shrunk $complex_puffin_path $n_puffins $initial_qmmm_complex_jid | grep "Final dependency:" | awk '{print $3}'`
cd ../..

# Prepare QM/QM/MM calculations
mkdir -p $system_name/solvent/qmqm
mkdir -p $system_name/complex/qmqm
cd $system_name/solvent/qmqm
create_qmqm_calculations_jid=$(sbatch --dependency=afterany:${final_qmmm_solvent_jid##* } --export=config_file=$qmqm_solvent_config_file $pipline_repo_path/NEQ_switching_ani2x/slurm/create_qmqm_calculations.sh)
cd ../../complex/qmqm
create_qmqm_calculations_jid=$(sbatch --dependency=afterany:${final_qmmm_complex_jid##* } --export=config_file=$qmqm_complex_config_file $pipline_repo_path/NEQ_switching_ani2x/slurm/create_qmqm_calculations.sh)
cd ../../..

# Run QM/QM/MM calculations
qmqmmm_solvent_jid=$(sbatch --dependency=afterok:${initial_upload_solvent_jid##* } --array=0-${n_puffins} -e ${system_name}/solvent/qmqm/puff.err -o ${system_name}/solvent/qmqm/puff.out $solvent_puffin_path)
qmqmmm_complex_jid=$(sbatch --dependency=afterok:${initial_upload_complex_jid##* } --array=0-${n_puffins} -e ${system_name}/complex/qmqm/puff.err -o ${system_name}/complex/qmqm/puff.out $complex_puffin_path)

# Run transfer learning
# TODO Submit slurm jobs for transfer learning and populate the following variables.
# TODO Execution of the jobs must depend on the afterany:${qmqmmm_solvent_jid##* } and afterany:${qmqmmm_complex_jid##* }.
transfer_learning_solvent_jid=23028668
transfer_learning_complex_jid=23028668
transfer_learning_solvent_output_dir="some/directory"
transfer_learning_complex_output_dir="some/directory"

# Run final NEQ switching simulation
cd $system_name/solvent/qmqm
$pipline_repo_path/NEQ_switching_ani2x/transfer_learning_runs/submit_ensemble.sh ${system_name}_solvent transfer_learning_solvent_output_dir $transfer_learning_solvent_jid
cd ../../complex/qmqm
$pipline_repo_path/NEQ_switching_ani2x/transfer_learning_runs/submit_ensemble.sh ${system_name}_complex_shrunk transfer_learning_complex_output_dir $transfer_learning_complex_jid
cd ../../..