import matplotlib.pyplot as plt
import numpy as np
import pdb
from batch_modelcore import *
from batch_io import *
import json



# Set output directory (where the plot files will appear)
opath = '../figs/'

# Load in the data from the experiment
ipath='/storage/silver/metstudent/ug/yd804409/' # location of output from the model
ensmean = loadexpt(ipath+'forLin0000_opttype20.json')




# You can type <variable>.keys for a list of the fields stored
# Also, this helpful little function displays the structural information of the file
exptinfo(ensmean)
exptname = ensmean['exptname']


# Do plot
fig = plt.figure(figsize=(8, 18))

# Title to go at the 'top of the page'
plt.suptitle('my title')

# Note: here choosing a particular alpha, arp, aru etc to plot
ax1 = fig.add_subplot(3,1,1)
xvals = ensmean['nens']
minx = np.min(xvals)
maxx = np.max(xvals)
stepx = 50
miny = 0.
maxy = 1.01
stepy = .25
ax1.axis([minx,maxx,miny,maxy])
ax1.tick_params(direction='out', which='both', labelsize='x-large')
ax1.set_xlabel('nens', fontsize='x-large')
ax1.set_ylabel('skill (crps)... lower is better', fontsize='x-large')
ax1.set_xticks(np.arange(minx,maxx,stepx))
ax1.set_yticks(np.arange(miny,maxy,stepy))
ax1.plot(xvals,ensmean['crps'][0,:,0,0,0,0],'or-',label='ensmean') # See https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.plot.html for info on lines/colours
ax1.legend(loc='upper right', fontsize='medium')



ax3 = fig.add_subplot(3,1,3)
xvals = ensmean['nens']
minx = np.min(xvals)
maxx = np.max(xvals)
stepx = 50
miny = 0.
maxy = 1.01
stepy = .25
ax3.axis([minx,maxx,miny,maxy])
ax3.tick_params(direction='out', which='both', labelsize='x-large')
ax3.set_xlabel('nens', fontsize='x-large')
ax3.set_ylabel('skill (mae)... lower is better', fontsize='x-large')
ax3.set_xticks(np.arange(minx,maxx,stepx))
ax3.set_yticks(np.arange(miny,maxy,stepy))
ax3.plot(xvals,ensmean['mae'][0,:,0,0,0,0],'or-',label='ensmean') # See https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.plot.html for info on lines/colours
ax3.legend(loc='upper right', fontsize='medium')

plt.savefig(opath+exptname+'_testfig.png')
plt.close(fig)

	
#import pdb; pdb.set_trace() # uncomment if you wish to stop the programme for debugging