import numpy as np
import pdb
from batch_modelcore import *
from batch_io import *
import datetime


# To run 'from the command line'
# module load anaconda
# source activate mypy3
# cd <relevant code directory>
# python3 <file>.py

# To run as a 'batch' job, see 
# https://research.reading.ac.uk/act/knowledgebase/racc-batch-jobs/
# sbatch <script.sh>
# squeue -u <userid>
# scancel <jobid>

# Experiment details
exptname = 'forLin0000' # the 'name' of this experiment
exptdesc = 'A test experiment' # some notes to describe this experiment
code = 'forLin0000_nens.py' # name of this file
path='/storage/silver/metstudent/ug/yd804409/' # output will appear here

# Select 'type' of simulation
# Options: 
#  'debug': for testing quickly, small and fast run with on-screen output to track progress
#  'singlehindcast': like ECMWF hindcast, 20 'winters' of 11 members, 12 weekly forecasts / winter
#  'singleforecast': like ECMWF forecast, 5 'winters' of 51 members, 12 weekly forecasts / winter
#  'multihindcast': as ECMWF hindcast but 5 times the ensemble size
#  'multiforecast': as ECMWF forecast but 5 times the ensemble size
#  'verylarge': large ensemble and long duration 
runtype = 'debug'

# Select 'type' of decision-making to use
# Options: 
#  10: perfect foresight (deterministic)
#  20: ensemble mean forecast (deterministic)
#  30: max-min optimisation
#  40: max-mean optimisation
#  50: max-max optimisation 
# Note you can choose multiple, e.g., [10,20,30,40], model will loop over choices 
# with one optimisation method per output file
opttype = [10,20]


# Parameters to investigate, model will loop over your choices with all choices
# in a single output file

# Meteorological elements
alpha = np.arange(0.33,0.67,0.33) # Weigel's alpha
beta = 0.0 # Weigel's beta
arp = [0.05,0.50,0.95] # AR1 (autocorrelation/persistence) of predictable component
aru = [0.05,0.50,0.95] # AR1 (autocorrelation/persistence) of unpredictable component

# Bucket model elements
maxcap=[0.00,4.00] # max capacity limit on reservoir before overtop	
drawprice=[0.5,1.0,2.0] # income gained per unit of drawdown


# Random number seed
# Each time you run the model, the 'weather' timeseries will be identically produced
# If you want a different 'weather' change this number
seed = 1 


# Perform some initial set up actions 
stime = datetime.datetime.today()
cv.solvers.options['show_progress'] = False # Switch off the solver verbosity
mdi = -999.999 # missing data indicator

# Met model parameters
# nt = number of timesteps in forecast period
# nens = number of ensemble members in meteorological forecast
baserate = 5. # background value to 'add' to forecast (mean state)
arb = 0.0 # AR1 (autocorrelation/persistence) of Weigel's beta term, don't change!
if (runtype == 'debug'):
	# additional on-screen output to track progress
	debug = True
nt = 100
nens = [5,15,25,51,101,201,401] 


# Decision model parameters
maxdraw=10. # maximum drawdown of reservoir in one timestep
mindraw=0. # minimum drawdown of reservoir in one timestep
mincap=0. # min capacity limit on reservoir before shortfall
shortprice=1. # cost per unit of shortfall
overtopprice=1. # cost per unit of overtopping
initres=0. # initial content of reservior

# Set up dimensions of arrays
d0 = len(alpha)
d0k = 'alpha'
d0v = alpha

d1 = len(nens)
d1k = 'nens'
d1v = nens

d2 = len(arp)
d2k = 'arp'
d2v = arp

d3 = len(aru)
d3k = 'aru'
d3v = aru

d4 = len(maxcap)
d4k = 'maxcap'
d4v = maxcap

d5 = len(drawprice)
d5k = 'drawprice'
d5v = drawprice


# Roll over the different optimization methods chosen
for oo in opttype:

	# Storage space.  Note only storing summary statistics not full timeseries
	met_mae = np.ones([d0,d1,d2,d3,d4,d5])*mdi
	met_rmse = np.ones([d0,d1,d2,d3,d4,d5])*mdi
	met_crps = np.ones([d0,d1,d2,d3,d4,d5])*mdi

	dec_income_mean = np.ones([d0,d1,d2,d3,d4,d5])*mdi
	dec_income_std = np.ones([d0,d1,d2,d3,d4,d5])*mdi
	dec_store_mean = np.ones([d0,d1,d2,d3,d4,d5])*mdi
	dec_store_std = np.ones([d0,d1,d2,d3,d4,d5])*mdi
	dec_over_mean = np.ones([d0,d1,d2,d3,d4,d5])*mdi
	dec_draw_mean = np.ones([d0,d1,d2,d3,d4,d5])*mdi
	dec_short_mean = np.ones([d0,d1,d2,d3,d4,d5])*mdi

	# Iterate over experiments
	# (This probably isn't "good" python but it works!)
	for ii in np.arange(0,len(alpha),1):

		for jj in np.arange(0,len(nens),1):

			for kk in np.arange(0,len(arp),1):

				for ll in np.arange(0,len(aru),1):
				
					for mm in np.arange(0,len(maxcap),1):

						for nn in np.arange(0,len(drawprice),1):

							if (debug): print('Experiment: ',ii,' - ',jj,' - ',kk,' - ',ll, ' - ',mm,' - ',nn,' (',oo,')')

							# Not all values of beta are sensible, only calculate if valid
							# else, this should leave MDIs
							if (beta <= np.sqrt(1.-alpha[ii]**2)):

								# Generate the meteorological forecast then store copy of summary
								# variables		
								metfc = gen_metfc(nt,nens[jj],alpha[ii],beta,baserate,arp[kk],aru[ll],arb,seed)
								met_mae[ii,jj,kk,ll,mm,nn] = metfc['mae']
								met_rmse[ii,jj,kk,ll,mm,nn] = metfc['rmse']
								met_crps[ii,jj,kk,ll,mm,nn] = metfc['crps']

								# Run the decision-making algorithm using the perfect, ensmean (det), 
								# maxmin, maxmean, maxmax (stoch) methods as appropriate

								# Deterministic - perfect
								if (oo == 10):

									decfc = gen_decfc_ens(maxdraw,
													 							mindraw,
																				maxcap[mm],
																				mincap,
																				drawprice[nn],
																				shortprice,
																				overtopprice,
																				initres,
																				metfc['xx'],
																				np.resize(metfc['xx'],(1,nt)),
																				oo,
																				True)

								# Deterministic - ensmean
								if (oo == 20):

									decfc = gen_decfc_ens(maxdraw,
													 							mindraw,
																				maxcap[mm],
																				mincap,
																				drawprice[nn],
																				shortprice,
																				overtopprice,
																				initres,
																				metfc['xx'],
																				np.resize(metfc['mu_ff'],(1,nt)),
																				oo,
																				True)


								# Stoch - maxmin
								elif (oo == 30):

									decfc = gen_decfc_ens(maxdraw,
													 							mindraw,
																				maxcap[mm],
																				mincap,
																				drawprice[nn],
																				shortprice,
																				overtopprice,
																				initres,
																				metfc['xx'],
																				metfc['ff_xx'],
																				oo,
																				True)


								# Stoch - maxmean
								elif (oo == 40):

									decfc = gen_decfc_ens(maxdraw,
													 							mindraw,
																				maxcap[mm],
																				mincap,
																				drawprice[nn],
																				shortprice,
																				overtopprice,
																				initres,
																				metfc['xx'],
																				metfc['ff_xx'],
																				oo,
																				True)


								# Stoch - maxmax
								elif (oo == 50):

									decfc = gen_decfc_ens(maxdraw,
													 							mindraw,
																				maxcap[mm],
																				mincap,
																				drawprice[nn],
																				shortprice,
																				overtopprice,
																				initres,
																				metfc['xx'],
																				metfc['ff_xx'],
																				oo,
																				True)

								# In each case calculate then store a copy of the summary variables
								dec_income_mean[ii,jj,kk,ll,mm,nn] = decfc['income_mean']
								dec_income_std[ii,jj,kk,ll,mm,nn] = decfc['income_std']
								dec_store_mean[ii,jj,kk,ll,mm,nn] = decfc['store_mean']
								dec_store_std[ii,jj,kk,ll,mm,nn] = decfc['store_std']
								dec_over_mean[ii,jj,kk,ll,mm,nn] = decfc['over_mean']
								dec_draw_mean[ii,jj,kk,ll,mm,nn] = decfc['draw_mean']
								dec_short_mean[ii,jj,kk,ll,mm,nn] = decfc['short_mean']




	# Save the data out
	filename=path+'/'+exptname+'_opttype'+str(oo)+'.json'
	data = {'exptname':exptname,
					'desc':exptdesc,
					'filename':filename,
					'path':path,
					'code':code,
					'seed':np.asarray(seed).tolist(),
					'mdi':mdi,
					'nt':np.asarray(nt).tolist(),
					'nens':np.asarray(nens).tolist(),
					'baserate':np.asarray(baserate).tolist(),
					'alpha':np.asarray(alpha).tolist(),
					'beta':np.asarray(beta).tolist(),
					'arp':np.asarray(arp).tolist(),
					'aru':np.asarray(aru).tolist(),
					'arb':np.asarray(arb).tolist(),
					'opttype':np.asarray(oo).tolist(),
					'maxdraw':np.asarray(maxdraw).tolist(),
					'mindraw':np.asarray(mindraw).tolist(),
					'maxcap':np.asarray(maxcap).tolist(),
					'mincap':np.asarray(mincap).tolist(),
					'drawprice':np.asarray(drawprice).tolist(),
					'shortprice':np.asarray(shortprice).tolist(),
					'overtopprice':np.asarray(overtopprice).tolist(),
					'initres':np.asarray(initres).tolist(),
					'd0':d0,
					'd0k':d0k,
					'd0v':np.asarray(d0v).tolist(),
					'd1':d1,
					'd1k':d1k,
					'd1v':np.asarray(d1v).tolist(),
					'd2':d2,
					'd2k':d2k,
					'd2v':np.asarray(d2v).tolist(),
					'd3':d3,
					'd3k':d3k,
					'd3v':np.asarray(d3v).tolist(),
					'd4':d4,
					'd4k':d4k,
					'd4v':np.asarray(d4v).tolist(),
					'd5':d5,
					'd5k':d5k,
					'd5v':np.asarray(d5v).tolist(),
					'mae':np.asarray(met_mae).tolist(),
					'rmse':np.asarray(met_rmse).tolist(),
					'crps':np.asarray(met_crps).tolist(),
					'income_mean':np.asarray(dec_income_mean).tolist(),
					'income_std':np.asarray(dec_income_std).tolist(),
					'store_mean':np.asarray(dec_store_mean).tolist(),
					'store_std':np.asarray(dec_store_std).tolist(),
					'over_mean':np.asarray(dec_over_mean).tolist(),
					'draw_mean':np.asarray(dec_draw_mean).tolist(),
					'short_mean':np.asarray(dec_short_mean).tolist()}
	saveexpt(filename,data)




etime = datetime.datetime.today()
print('Start:',stime)
print('End:',etime)