import os
import sys
import glob
import subprocess

# List of variable to save for output

PLASIM_VAR = ['tas','hfls','hfss']
OCEAN_VAR = ['heata','fldoa','sst']
LSG_VAR = ['fluxhea', 'tbound']

PLASIM_VAR = ['time','lat','lon'] + PLASIM_VAR
OCEAN_VAR = ['time','lat','lon'] + OCEAN_VAR
LSG_VAR = ['time','lat','lon','lev'] + LSG_VAR

# check the python version
if '3.6' not in sys.version:
	print("This script require Python 3.6 !")
	print("Try:")
	print("\n\tmodule load python3\n\n")
	print("and restart this script again.")
	print("Aborting...")
	sys.exit(1)
	
#check if ecfs utils are loaded
user = os.getenv('USER')
try:
	dummy = subprocess.run(['els', 'ectmp:/'+user+'/'], check=True, stdout=subprocess.PIPE)
except:
	print("This script require the ECFS toolchain to be loaded !")
	print("Try:")
	print("\n\tmodule load ecfs\n\n")
	print("and restart this script again.")
	print("Aborting...")
	sys.exit(1)

#check if netcdf utils are loaded
try:
	dummy = subprocess.run(['ncdump'], check=True, stderr=subprocess.PIPE)
except:
	print("This script require NetCDF4 toolchain to be loaded !")
	print("Try:")
	print("\n\tmodule load netcdf4\n\n")
	print("and restart this script again.")
	print("Aborting...")
	sys.exit(1)

home_dir = os.getenv("HOME")
scratch_dir = os.getenv("SCRATCH")
perm_dir = os.getenv("PERM")
plasim_dir = home_dir + "/models/PLASIM/"

if len(sys.argv) < 5:
	print('Bad arguments:', sys.argv)
	print('Usage:')
	print('\n\t python3 restart_ensemble_experiment.py where experiment ensemble_size number_of_years where_to_save\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('\tnumber_of_years :\tNumber of years simulated by one run of the experiment.')
	print('\twhere_to_save :\t\tOptional. Where to backup the previous run. If not provide, uses $PERM.')
	sys.exit(0)

where = sys.argv[1]
basedir = where.split('/')[-1]
experiment = sys.argv[2]
ensemble_size = int(sys.argv[3])
restart_year_to_save = sys.argv[4]
try:
	save = sys.argv[5]
except:
	save = perm_dir

save_light = save + "/"+experiment+"_light/"
save += "/"+experiment+"/"