import os
import sys

# check the python version
if float(sys.version[:3]) <= 3.6:
	print("This script require Python 3. !")
	print("Try:")
	print("\n\tmodule load python3\n\n")
	print("and restart this script again.")
	print("Aborting...")
	sys.exit(1)

if len(sys.argv) < 4:
	print('Bad arguments:', sys.argv)
	print('Usage:')
	print('\n\t python3 make_mpi_ensemble_algorithm_experiment.py where experiment ensemble_size run_folder_name\n')
	print('Arguments:\n')
	print('\twhere :\t\t\tWhere the experiment ensemble folders are located.')
	print('\texperiment :\t\tName of the experiment.')
	print('\tensemble_size :\t\tSize of the ensemble.')
	print('\trun_folder_name :\tOptional. Addtional name added to the "run_" folder name in plasim folder. Used to selected experiment executable. If not provided, assumes run folder is "run".')
	sys.exit(0)

mail = 'jerome.sauer@uclouvain.be'
walltime = '48:00:00'
memory_per_task = '500mb'

home_dir = os.getenv("HOME")
scratch_dir = os.getenv("SCRATCH")
plasim_dir = home_dir + "/PLASIM_ALGORITHM/"

where = sys.argv[1]
experiment = sys.argv[2]
ensemble_size = int(sys.argv[3])

run_folder = plasim_dir+"plasim/run"

if len(sys.argv) == 5:
	run_folder_ext = sys.argv[4]
	run_folder += '_'+run_folder_ext

if not os.path.isdir(run_folder):
	print("PLASIM run directory missing!")
	print("Run "+plasim_dir+"/most.x first.")
	print("Aborting...")
	sys.exit(1)

print('Creating the directory holding the experiment ensemble...')
experiment_folder = where+"/"+experiment+"/"
os.system('mkdir -p '+experiment_folder)

print('Creating the ensemble folders...')
for i in range(1, ensemble_size+1):
	member_number = str(i).rjust(2, '0')
	ensemble_member_folder = experiment_folder+'run_'+experiment+'_'+member_number
	os.system('cp -r '+run_folder+' '+ensemble_member_folder)
	job_name = 'plasim_'+experiment+'_'+member_number
	f = open(ensemble_member_folder+'/PBS_'+job_name, 'w')
	f.write("""#!/bin/ksh
#PBS -S /usr/bin/ksh
#PBS -q nf
#PBS -l EC_total_tasks=16
#PBS -l EC_hyperthreads=1

		   
""")

	f.write('#PBS -N '+job_name+'\n')
	f.write('#PBS -o '+job_name+'.o\n')
	f.write('#PBS -e '+job_name+'.e\n\n')

	f.write('#PBS -m abe\n')
	f.write('#PBS -M '+mail+'\n\n')
	f.write('#PBS -l walltime='+walltime+'\n\n')
	f.write('#PBS -l EC_memory_per_task='+memory_per_task+'\n\n')

	f.write("""module switch PrgEnv-cray PrgEnv-gnu
module load cray-snplauncher

module load netcdf4
module load netcdf
module load cdo

""")
	f.write('cd '+ensemble_member_folder+'\n')
	f.write('./most_plasim_run\n')
	f.close()

print('Experiment ensemble created and available in the folder:')
print(experiment_folder)