%matplotlib inline
from matplotlib import pylab
import matplotlib.pyplot as plt
import netCDF4 as NC
import numpy as np


path = '/data/nsoontie/MEOPAR/SalishSea/results/storm-surges/spinups/'


runs = {'restart',  'initialize'}

fUs={}; fVs={}; fTs={};

for key in runs:
    print key
    fUs[key] = NC.Dataset(path + key +'/SalishSea_4h_20021018_20021024_grid_U.nc','r');
    fVs[key] = NC.Dataset(path + key +'/SalishSea_4h_20021018_20021024_grid_V.nc','r');
    fTs[key] = NC.Dataset(path + key +'/SalishSea_4h_20021018_20021024_grid_T.nc','r');



# define plot parameters
#time, depth, colormap etc
timestamp =12#model output time 
times={}
for key in runs:
    times[key]= fVs[key].variables['time_counter'] #array of output times
    time = times[key][timestamp]  
    print key + ' ' + str(time/86400)#physical time in seconds
depthlevel = 0  #model level to plot
depths = fVs[key].variables['depthv'] #array of depths
depth = depths[depthlevel]  #physical depth in meters
cmap = plt.get_cmap('jet') 
u_lat = fUs[key].variables['nav_lat']
u_lon = fUs[key].variables['nav_lon']
T_lat = fTs[key].variables['nav_lat']
T_lon = fTs[key].variables['nav_lon']
v_lat = fVs[key].variables['nav_lat']
v_lon = fVs[key].variables['nav_lon']


def get_variables(fU,fV,fT,timestamp,depth):
    #function returnes masked u,v, T,S, SSH from NETCDF files fU,fV,fT at timestamp and depth
    
    # get u and ugrid
    u_vel = fU.variables['vozocrtx']  #u currents and grid
    U = u_vel[timestamp,depthlevel,:,:] #grab data at specified level and time.
    #mask u so that white is plotted on land points
    mu =U== 0
    U= np.ma.array(U,mask=mu)
    
    #get v and v grid
    v_vel = fV.variables['vomecrty']  #v currents and grid
    V = v_vel[timestamp,depthlevel,:,:] #grab data at specified level and time.

    #mask v so that white is plotted on land points
    mu = V == 0
    V= np.ma.array(V,mask=mu)

    #grid for T points
    eta = fT.variables['sossheig']
    E = eta[timestamp,:,:]
    mu=E==0; E = np.ma.array(E,mask=mu)

    sal = fT.variables['vosaline']
    S= sal[timestamp,depthlevel,:,:]
    mu=S==0; S= np.ma.array(S,mask=mu)

    temp=fT.variables['votemper']
    T = temp[timestamp,depthlevel,:,:]
    mu=T==0; T = np.ma.array(T,mask=mu)
    
    return U, V, E, S, T


Us={}; Vs={}; Es ={}; Ss={}; Ts={};

for key in runs:
    [Us[key], Vs[key], Es[key], Ss[key], Ts[key]] = get_variables(fUs[key],fVs[key],fTs[key],timestamp,depth)

def combine_data(data_list):
    #combines ouput files so that they are easier to analyze
    #data_list is a dict object that contains the netcdf handles for the files of importance. 
    us={}; vs={}; lats={}; lons={}; sals={}; tmps={}; sshs={};
    for k in data_list:
        net=data_list.get(k)
        us[k]=net.variables['vozocrtx']
        vs[k]=net.variables['vomecrty']
        lats[k]=net.variables['nav_lat']        
        lons[k]=net.variables['nav_lon']
        tmps[k]=net.variables['votemper']
        sals[k]=net.variables['vosaline']
        sshs[k]=net.variables['sossheig']
    return us, vs, lats, lons, tmps, sals, sshs


allstations = {'PointAtkinson': 7795, 'Victoria': 7120, 'PatriciaBay': 7277, 'CampbellRiver': 8074, 'Plume': 1,
               'Thalweg1': [], 'Thalweg2': [], 'Thalweg3': [], 'Thalweg4': [], 'Thalweg5': [], 'Thalweg6': []} 
run_stations={}
us={}; vs={}; lats={}; lons={}; tmps={}; sals={}; sshs={}; ts={};

for run in runs:
    for key in allstations:
        string = path + run + '/1h_' + key + '.nc'
        run_stations[key] = NC.Dataset(string,'r');
    [us[run], vs[run], lats[run], lons[run], tmps[run], sals[run], sshs[run]] = combine_data(run_stations)
    run_stations={};

plt.figure(figsize=(15,7))
umin=-1.5; umax=1.5; #range on colourbar
smin=10; smax=32;
Tmin=6; Tmax=14
ax=[-126,-122,47,51]
i=1
#plot u
for key in runs:
    pylab.subplot(1,3,i)
    pylab.pcolor(u_lon,u_lat,Us[key],vmin=umin,vmax=umax,cmap=cmap)
    pylab.title(key)
    pylab.axis(ax)
    i=i+1
    pylab.colorbar()
    
pylab.subplot(1,3,3)
pylab.pcolor(u_lon,u_lat,Us['restart']-Us['initialize'],vmin=-.2,vmax=.2,cmap=cmap)
pylab.colorbar()
for key in allstations:
    pylab.plot(lons['restart'][key],lats['restart'][key],'o')
pylab.axis(ax)
pylab.title('difference')

# Plot the u and v currents at a specified depth (depthlevel)
plt.figure(figsize=(15,7))

i=1
#plot u
for key in runs:
    pylab.subplot(1,3,i)
    pylab.pcolor(T_lon,T_lat,Ss[key],vmin=smin,vmax=smax,cmap='winter')
    pylab.title(key)
    pylab.axis(ax)
    pylab.colorbar()
    i=i+1
    
pylab.subplot(1,3,3)
pylab.pcolor(T_lon,T_lat,Ss['restart']-Ss['initialize'],vmin=-1,vmax=1,cmap='jet')
pylab.colorbar()
for key in allstations:
    pylab.plot(lons['restart'][key],lats['restart'][key],'o')
pylab.axis(ax)
pylab.title('difference')
 

#load thalweg points
lines = np.loadtxt('../../tools/analysis_tools/thalweg.txt', delimiter=" ", unpack=False)
lines = lines.astype(int)

#load data to pot

dep = fTs['restart'].variables['deptht']
#array of thalweg indices
ds=np.arange(0,lines.shape[0],1);
thalweg_lon = T_lon[lines[:,0],lines[:,1]]
thalweg_lat = T_lat[lines[:,0],lines[:,1]]


#plotting parameters
smin=28; smax=34;
lns=np.arange(34,27.5,-0.5);

#masking
XX,ZZ = np.meshgrid(ds,-dep[:])

plt.figure(figsize=(19,12))
i=1;
for key in runs:
    net=fTs[key].variables['vosaline']
    salThal=net[timestamp,:,lines[:,0],lines[:,1]];
    mu =salThal == 0; salThal= np.ma.array(salThal,mask=mu)
    pylab.subplot(2,1,i)
    pylab.pcolormesh(XX,ZZ,salThal,vmin=smin,vmax=smax,cmap='rainbow')
    pylab.colorbar()
    CS=pylab.contour(XX,ZZ,salThal,lns, colors='black')
    pylab.clabel(CS, fontsize=9, inline=1)
    i=i+1
    pylab.title(key)

stations = {'PointAtkinson': 7795, 'Victoria': 7120, 'PatriciaBay': 7277, 'CampbellRiver': 8074, 'Plume': 1}
thalwegs = ['Thalweg1','Thalweg2', 'Thalweg3', 'Thalweg4', 'Thalweg5', 'Thalweg6']

vmin=-1.5
vmax=1.5
plt.figure(figsize=(15,20))
depth=0

i=1;
for key in stations:
    for run in runs:
        ax=pylab.subplot(6,2,i)
        if run == 'restart':
            cols='b'
        else:
            cols='g'
        pylab.plot(us[run][key][:,depth,0,0],'-',color=cols,label=run + ' u' )
        pylab.plot(vs[run][key][:,depth,0,0],'--',color=cols,label=run + ' v' )
        pylab.title(key)
        ax.set_ylim([vmin,vmax])
    pylab.grid()
    i=i+2

    

i=2;
for key in thalwegs:
    for run in runs:
        ax=pylab.subplot(6,2,i)
        if run == 'restart':
            cols='b'
        else:
            cols='g'
        pylab.plot(us[run][key][:,depth,0,0],'-',color=cols,label=run + ' u')
        pylab.plot(vs[run][key][:,depth,0,0],'--',color=cols,label=run + ' v' )
        pylab.title(key)
        ax.set_ylim([vmin,vmax])
    pylab.grid()
    i=i+2
    
pylab.legend(bbox_to_anchor=(-1.05, 1), loc=2)

print 'Surface velocties'
print 'Horizontal axis is number of hours past Oct 18, 2002 at 12am.'
print 'Dashed lines are v. Solid lines are u'

depth=10;
print depths[depth]

stations = {'PointAtkinson': 7795, 'Victoria': 7120, 'PatriciaBay': 7277, 'CampbellRiver': 8074, 'Plume': 1}
thalwegs = ['Thalweg1','Thalweg2', 'Thalweg3', 'Thalweg4', 'Thalweg5', 'Thalweg6']

vmin=-1.5
vmax=1.5
plt.figure(figsize=(15,20))

i=1;
for key in stations:
    for run in runs:
        ax=pylab.subplot(6,2,i)
        if run == 'restart':
            cols='b'
        else:
            cols='g'
        pylab.plot(us[run][key][:,depth,0,0],'-',color=cols,label=run + ' u' )
        pylab.plot(vs[run][key][:,depth,0,0],'--',color=cols,label=run + ' v' )
        pylab.title(key)
        ax.set_ylim([vmin,vmax])
    pylab.grid()
    i=i+2

    

i=2;
for key in thalwegs:
    for run in runs:
        ax=pylab.subplot(6,2,i)
        if run == 'restart':
            cols='b'
        else:
            cols='g'
        pylab.plot(us[run][key][:,depth,0,0],'-',color=cols,label=run + ' u')
        pylab.plot(vs[run][key][:,depth,0,0],'--',color=cols,label=run + ' v' )
        pylab.title(key)
        ax.set_ylim([vmin,vmax])
    pylab.grid()
    i=i+2
    
pylab.legend(bbox_to_anchor=(-1.05, 1), loc=2)

print 'Velocties at ' + str(depths[depth]) + ' m'
print 'Horizontal axis is number of hours past Oct 18, 2002 at 12am.'
print 'Dashed lines are v. Solid lines are u'

depth=25;
print depths[depth]

stations = {'PointAtkinson': 7795, 'Victoria': 7120, 'PatriciaBay': 7277, 'CampbellRiver': 8074, 'Plume': 1}
thalwegs = ['Thalweg1','Thalweg2', 'Thalweg3', 'Thalweg4', 'Thalweg5', 'Thalweg6']

vmin=-1.5
vmax=1.5
plt.figure(figsize=(15,20))

i=1;
for key in stations:
    for run in runs:
        ax=pylab.subplot(6,2,i)
        if run == 'restart':
            cols='b'
        else:
            cols='g'
        pylab.plot(us[run][key][:,depth,0,0],'-',color=cols,label=run + ' u' )
        pylab.plot(vs[run][key][:,depth,0,0],'--',color=cols,label=run + ' v' )
        pylab.title(key)
        ax.set_ylim([vmin,vmax])
    pylab.grid()
    i=i+2

    

i=2;
for key in thalwegs:
    for run in runs:
        ax=pylab.subplot(6,2,i)
        if run == 'restart':
            cols='b'
        else:
            cols='g'
        pylab.plot(us[run][key][:,depth,0,0],'-',color=cols,label=run + ' u')
        pylab.plot(vs[run][key][:,depth,0,0],'--',color=cols,label=run + ' v' )
        pylab.title(key)
        ax.set_ylim([vmin,vmax])
    pylab.grid()
    i=i+2
    
pylab.legend(bbox_to_anchor=(-1.05, 1), loc=2)

print 'Velocties at ' + str(depths[depth]) + ' m'
print 'Horizontal axis is number of hours past Oct 18, 2002 at 12am.'
print 'Dashed lines are v. Solid lines are u'


emin=-2
emax=2

stations = {'PointAtkinson': 7795, 'Victoria': 7120, 'PatriciaBay': 7277, 'CampbellRiver': 8074, 'Plume': 1}
thalwegs = ['Thalweg1','Thalweg2', 'Thalweg3', 'Thalweg4', 'Thalweg5', 'Thalweg6']

vmin=-1.5
vmax=1.5
plt.figure(figsize=(15,20))

i=1;
for key in stations:
    for run in runs:
        ax=pylab.subplot(6,2,i)
        pylab.plot(sshs[run][key][:,0,0],label=run )
        pylab.title(key)
        ax.set_ylim([emin,emax])
    pylab.grid()
    i=i+2

    

i=2;
for key in thalwegs:
    for run in runs:
        ax=pylab.subplot(6,2,i)
        pylab.plot(sshs[run][key][:,0,0],label=run )
        pylab.title(key)
        ax.set_ylim([emin,emax])
    pylab.grid()
    i=i+2
    
pylab.legend(bbox_to_anchor=(-1.05, 1), loc=2)

print 'Sea surface height'
print 'Horizontal axis is number of hours past Oct 18, 2002 at 12am.'