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

fU1 = NC.Dataset('/ocean/dlatorne/MEOPAR/SalishSea/results/61d70d_nu200/SalishSea_3d_20021114_20021123_grid_U.nc','r');
fV1 = NC.Dataset('/ocean/dlatorne/MEOPAR/SalishSea/results/61d70d_nu200/SalishSea_3d_20021114_20021123_grid_V.nc','r');
fT1=  NC.Dataset('/ocean/dlatorne/MEOPAR/SalishSea/results/61d70d_nu200/SalishSea_3d_20021114_20021123_grid_T.nc','r');
nu1='nu=200'

fU2 = NC.Dataset('/ocean/dlatorne/MEOPAR/SalishSea/results/41d70d/SalishSea_3d_20021025_20021123_grid_U.nc','r');
fV2 = NC.Dataset('/ocean/dlatorne/MEOPAR/SalishSea/results/41d70d/SalishSea_3d_20021025_20021123_grid_V.nc','r');
fT2 = NC.Dataset('/ocean/dlatorne/MEOPAR/SalishSea/results/41d70d/SalishSea_3d_20021025_20021123_grid_T.nc','r');
nu2='nu=50'

#define a function to unstagger the data
#This assumes that data_array is given as (y, x)
#dim is 'Y' for unstaggering over y or 'X' for unstaggering over x
#Note that the far "Western" and far "Southern" boundaries are ignored in the u,v arrays. 
#This should be taken into account when plotting on the same grid as T data.  
def unstagger(data_array,dim):
    varin=data_array
    size = data_array.shape;
    sizeU=size[1];
    sizeV=size[0];
    if dim == 'X':
        varout = 0.5*(data_array[:,0:sizeU-1] + data_array[:,1:sizeU]);
        varout = varout[1:sizeV,:]
    else:
        varout = 0.5*(varin[0:sizeV-1,:] + varin[1:sizeV,:]);
        varout = varout[:,1:sizeU];
    return varout

times1 = fT1.variables['time_counter'] #array of output times
times2 = fT2.variables['time_counter'] #array of output times
depths = fT1.variables['deptht'] #array of depths

print times1.shape
print times2.shape

#index for output times. These should be chosen carefully since each file has a different number of times saved. 
t1=2; t2=9; #choose the last one for each.
time1=times1[t1]
time2=times2[t2]
print time1; print time2


#get the data
# get u and ugrid
u_vel1 = fU1.variables['vozocrtx']  #u currents and grid
u_lat = fU1.variables['nav_lat']
u_lon = fU1.variables['nav_lon']

#get v and v grid
v_vel1 = fV1.variables['vomecrty']  #v currents and grid
v_lat = fV1.variables['nav_lat']
v_lon = fV1.variables['nav_lon']

#get sal and temp, sear surface height
sal1=fT1.variables['vosaline']
temp1=fT1.variables['votemper']
ssh1=fT1.variables['sossheig']
T_lat = fT1.variables['nav_lat']
T_lon = fT1.variables['nav_lon']

# get u and ugrid
u_vel2 = fU2.variables['vozocrtx']  #u currents and grid

#get v and v grid
v_vel2 = fV2.variables['vomecrty']  #v currents and grid

#get sal
sal2=fT2.variables['vosaline']
temp2=fT2.variables['votemper']
ssh2=fT2.variables['sossheig']

#Look at max currents over last three ouput times of each
# for loop to print the maximum currents and location (grid points) over time. 
# columns are max current, z, y, x

print str(nu1) + ' at t= ' + str(time1)
maxes1=np.zeros((t1+1,4))
for ii in range(0,t1+1):
    U3d1 = u_vel1[ii,:,:,:]; V3d1 = v_vel1[ii,:,:,:];
    c1 = np.sqrt(U3d1**2+V3d1**2); 
    max1=c1.max();
    dim_max1 = np.unravel_index(np.argmax(c1),c1.shape)
    maxes1[ii,:] = [max1, dim_max1[0], dim_max1[1], dim_max1[2]];

print maxes1

print  str(nu2)  + ' at t= ' + str(time2)
maxes2=np.zeros((t1+1,4))
for ii in range(7,t2+1):
    U3d2 = u_vel2[ii-7,:,:,:]; V3d2 = v_vel2[ii-7,:,:,:];
    c2 = np.sqrt(U3d2**2+V3d2**2); 
    max2=c2.max();
    dim_max2 = np.unravel_index(np.argmax(c2),c2.shape)
    maxes2[ii-7,:] = [ max2, dim_max2[0], dim_max2[1], dim_max2[2]]; 

print maxes2

print 'difference'
print maxes1[:,0]-maxes2[:,0]

#Quiver plot.
plt.figure(figsize=(14,8))
st=4 #step size for quiver arrows
sc=20
axq = [-123.5,-122.75,48.25,49];

dl=0 #depth level

#A mask for the land
lmin=3; lmax=10;
tempPlot = temp1[t1,0,:,:]; land=tempPlot; mu =land != 0
land= np.ma.array(land,mask=mu)
time=times1[t1];

#data
U1 = u_vel1[t1,dl,:,:] #3D data at this time
mu =U1 == 0; U1 = np.ma.array(U1,mask=mu) 
U1_unstag = unstagger(U1,'X')

V1 = v_vel1[t1,dl,:,:] #grab data at specified level and time.
mu = V1 == 0; V1= np.ma.array(V1,mask=mu)
V1_unstag = unstagger(V1,'Y');

U2 = u_vel2[t2,dl,:,:] #grab data at specified level and time.
mu =U2 == 0; U2 = np.ma.array(U2,mask=mu)
U2_unstag = unstagger(U2,'X')

V2 = v_vel2[t2,dl,:,:] #grab data at specified level and time.
mu = V2 == 0; V2= np.ma.array(V2,mask=mu)
V2_unstag = unstagger(V2,'Y')

# rotate the velocity vectors:
theta=29 * np.pi / 180 #rotation angle in radians. 
U1_true = U1_unstag * np.cos(theta) - V1_unstag * np.sin(theta) #E velocity
V1_true = U1_unstag * np.sin(theta) + V1_unstag * np.cos(theta) #N velocity

U2_true = U2_unstag * np.cos(theta) - V2_unstag * np.sin(theta) #E velocity
V2_true = U2_unstag * np.sin(theta) + V2_unstag * np.cos(theta) #N velocity

#modify the T grid
Tsize = T_lat.shape; end_lat=Tsize[0]; end_lon=Tsize[1];
T_lat_plot = T_lat[1:end_lat,1:end_lon];
T_lon_plot = T_lon[1:end_lat,1:end_lon];
print T_lat_plot.shape
print U1_true.shape

#quiver plots
#plotting N and E velocities (U_true, V_true)
plt.subplot(1,2,1)
Q = pylab.quiver(T_lon_plot[::st, ::st],T_lat_plot[::st, ::st],U1_true[::st, ::st],V1_true[::st, ::st],scale=sc,color='b')
qk = pylab.quiverkey(Q,.1, .3, 2, r'$2 m/s$', labelpos='N', color='w',
               fontproperties={'weight': 'bold','size': 16},labelcolor='white')
#overlay with zero-contour of sea surface height
pylab.pcolor(T_lon,T_lat,land,vmin=lmin,vmax=lmax,cmap='hot')
pylab.contour(T_lon,T_lat,ssh1[0,:,:],[0],colors='k',aspect=(1 / np.cos(np.median(T_lon) * np.pi / 180)))
pylab.axis(axq); pylab.title(str(nu1) +' at t=' + str(time1/86400) + 'd') 

#NW and NE velocities. Just checking to see that the tilting was reasonable.  
#In this plot the currents shouldn't really match up with the map
plt.subplot(1,2,2)
Q = pylab.quiver(T_lon_plot[::st,::st],T_lat_plot[::st,::st],U2_true[::st,::st],V2_true[::st,::st],scale=sc,color='b')
qk = pylab.quiverkey(Q,.1, .3, 2, r'$2 m/s$', labelpos='N', color='w',
               fontproperties={'weight': 'bold','size': 16},labelcolor='white')
#overlay with zero-contour of sea surface height
pylab.pcolor(T_lon,T_lat,land,vmin=lmin,vmax=lmax,cmap='hot')
pylab.contour(T_lon,T_lat,ssh2[0,:,:],[0],colors='k',aspect=(1 / np.cos(np.median(T_lon) * np.pi / 180)))
pylab.axis(axq);  pylab.title(str(nu2) +' at t=' + str(time2/86400) + 'd') 

minv=-1
maxv=1
cmap = plt.get_cmap('jet') 

fig =plt.figure(figsize=(14,14))
pylab.subplot(2,2,1)
pylab.pcolormesh(U1,cmap=cmap,vmin=minv,vmax=maxv)
pylab.colorbar()
pylab.title('U for ' +str(nu1)+' at t='+ str(time1/86400) + 'd')

pylab.subplot(2,2,2)
pylab.pcolormesh(U2,cmap=cmap,vmin=minv,vmax=maxv)
pylab.colorbar()
pylab.title('U for ' +str(nu2)+' at t='+ str(time2/86400) + 'd')

pylab.subplot(2,2,3)
pylab.pcolormesh(V1,cmap=cmap,vmin=minv,vmax=maxv)
pylab.colorbar()
pylab.title('V for ' +str(nu1)+' at t='+ str(time1/86400) + 'd')

pylab.subplot(2,2,4)
pylab.pcolormesh(V2,cmap=cmap,vmin=minv,vmax=maxv)
pylab.colorbar()
pylab.title('V for ' +str(nu2)+' at t='+ str(time2/86400) + 'd')

sshP1 = ssh1[t1,:,:]
mu =sshP1 == 0; sshP1= np.ma.array(sshP1,mask=mu)
sshP2 = ssh2[t2,:,:]
mu =sshP2 == 0; sshP2= np.ma.array(sshP2,mask=mu)

minS=-1
maxS=1
cmap = plt.get_cmap('jet') 

fig =plt.figure(figsize=(16,8))
pylab.subplot(1,3,1)
pylab.pcolormesh(sshP1,cmap=cmap,vmin=minS,vmax=maxS)
pylab.colorbar()
pylab.title('ssh for ' +str(nu1)+' at t='+ str(time1/86400) + 'd')

pylab.subplot(1,3,2)
pylab.pcolormesh(sshP2,cmap=cmap,vmin=minS,vmax=maxS)
pylab.colorbar()
pylab.title('ssh for ' +str(nu2)+' at t='+ str(time2/86400) + 'd')

pylab.subplot(1,3,3)
pylab.pcolormesh(sshP1-sshP2,cmap=cmap,vmin=-0.1,vmax=0.1)
pylab.colorbar()
pylab.title('difference in ssh: ' +str(nu1) +' - ' +str(nu2))

salP1 = sal1[t1,dl,:,:]
mu =salP1 == 0; salP1= np.ma.array(salP1,mask=mu)
salP2 = sal2[t2,dl,:,:]
mu =salP2 == 0; salP2= np.ma.array(salP2,mask=mu)

minS=20
maxS=31
axS = [minS, maxS, -100, 0]
cmap = plt.get_cmap('winter') 

fig =plt.figure(figsize=(14,8))
pylab.subplot(1,2,1)
pylab.pcolormesh(salP1,cmap=cmap,vmin=minS,vmax=maxS)
pylab.colorbar()
pylab.title('S for ' +str(nu1)+' at t='+ str(time1/86400) + 'd')

pylab.subplot(1,2,2)
pylab.pcolormesh(salP2,cmap=cmap,vmin=minS,vmax=maxS)
pylab.colorbar()
pylab.title('S for ' +str(nu2)+' at t='+ str(time2/86400) + 'd')

smin=28; smax=34

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

thalweg_lon = T_lon[lines[:,0],lines[:,1]]
thalweg_lat = T_lat[lines[:,0],lines[:,1]]

ds=np.arange(0,lines.shape[0],1);

vs=np.arange(34,27.5,-0.5);

salT1=sal1[t1,:,lines[:,0],lines[:,1]];
salT2=sal2[t2,:,lines[:,0],lines[:,1]];

#masking
mu =salT1 == 0; salT1= np.ma.array(salT1,mask=mu)
mu =salT2 == 0; salT2= np.ma.array(salT2,mask=mu)
XX,ZZ = np.meshgrid(ds,-depths[:])

plt.figure(figsize=(19,12))
pylab.subplot(2,1,1)
pylab.pcolormesh(XX,ZZ,salT1,vmin=smin,vmax=smax,cmap='rainbow')
pylab.colorbar()
CS=pylab.contour(XX,ZZ,salT1,vs, colors='black')
pylab.clabel(CS, fontsize=9, inline=1)
pylab.title('Sal ' + str(nu1))

pylab.subplot(2,1,2)
pylab.pcolormesh(XX,ZZ,salT2,vmin=smin,vmax=smax,cmap='rainbow')
pylab.colorbar()
CS=pylab.contour(XX,ZZ,salT2,vs, colors='black')
pylab.clabel(CS, fontsize=9, inline=1)
pylab.title('Sal ' + str(nu2))