import mit #Package that loads from the netcdf files to produce the structured array `dgl`
import plmit #Package that plots the dgl structured arrays

tims = np.arange(0,70) #Set the time points you want to load.  This is given in days,
#you can change to seconds or years by adjusting set_scales.py

levs = np.arange(0,1) #Set the model leves you want to load.  This is given in
#model level index where 0 is the surface level, 1 is the level below this etc.

dfl = 'pv' #This is a pointer to the diagnostic file name in the mnc_tile_xxxx files.
#It needs to be a string that would identify the relevant files is you ran a command like
#`ls pv*' in the file in linux.  The exception is for the state files, where this
#field must be 'st', so the code knows to use the different diagnostic names in the state files.

diags = ['u','v','rvf'] #List of strings giving the diagnostic names.  Diagnostics
#that are calculated, such as relative vorticity, should be last.

dgl=mit.ldgl(tims, levs, dfl, diags) #Load the diagnostics into the structured array dgl
print dgl[0]['tims'] #To check that appropriate times have been loaded

print dgl.dtype #To check the size of the loaded files

#Now plot this output in the x-y plane
clf='rv' #Set the colorfield of the plot

clims=[-1,1] #Set the colourbar limits

plmit.pxy(dgl,clf,ti='a',clims = clims) #Do the plotting. ti = 'a' means that
#the entire time series will be animated automatically.
import matplotlib.pylab as plt
import matplotlib.animation as animation

import mit
import plmit

import os
os.chdir('/work/n01/n01/liamb/pig/tracer/expts/mnc_test_0018_useconv_hydro')

tims = np.arange(0,0.5,1e-2)
levs = np.arange(0,100)
dfile = 'h'
diags = ['t','s']
clims = [1, 3]

dgl = mit.ldgl(tims, levs, dfile, diags)
print dgl.dtype
print dgl[0][0]

time = len(dgl[0]['tims'])
x = dgl[0]['x']
y = dgl[0]['y']
z = dgl[0]['z']

###########
fig = plt.figure(figsize=(14,14),facecolor='white')

ax = plt.subplot()
quad1 = ax.pcolormesh(y,z,np.transpose(np.squeeze(dgl[0]['t'][:,:,:,0])),shading='gouraud',cmap = plt.cm.bwr)
ax.set_xlabel('time')
ax.set_ylabel('amplitude')