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

# Load data

fU1 = NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/viscosity/nu60_long/SalishSea_4h_20020915_20020921_grid_U.nc','r');
fV1 = NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/viscosity/nu60_long/SalishSea_4h_20020915_20020921_grid_V.nc','r');
fT1= NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/viscosity/nu60_long/SalishSea_4h_20020915_20020921_grid_T.nc','r');
nu1='60,evd10'

fU2 = NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/viscosity/apr60_nu40_long/SalishSea_4h_20020915_20020921_grid_U.nc','r');
fV2 = NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/viscosity/apr60_nu40_long/SalishSea_4h_20020915_20020921_grid_V.nc','r');
fT2 = NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/viscosity/apr60_nu40_long/SalishSea_4h_20020915_20020921_grid_T.nc','r');
nu2='40,evd60'

# define plot parameters
#time, depth, colormap etc
timestamp =41#model output time
times = fT1.variables['time_counter'] #array of output times
time = times[timestamp]  #physical time in seconds
depthlevel = 0 #model level to plot
depths = fT1.variables['deptht'] #array of depths
depth = depths[depthlevel]  #physical depth in meters
cmap = plt.get_cmap('jet')      #colormap

# oint near instabilities
ibad=264; #ibad=305;
jbad=347; #jbad=343;
#point near plume
iplume=310
jplume=440

# get u and ugrid
u_vel1 = fU1.variables['vozocrtx']  #u currents and grid
U1 = u_vel1[timestamp,:,:,:] #3D data at this time
#U1 = u_vel1[timestamp,depthlevel,:,:] #grab data at specified level and time.
u_lat = fU1.variables['nav_lat']
u_lon = fU1.variables['nav_lon']
#mask u so that white is plotted on land points
mu =U1 == 0
U1 = np.ma.array(U1,mask=mu)

#get v and v grid
v_vel1 = fV1.variables['vomecrty']  #v currents and grid
V1 = v_vel1[timestamp,:,:,:] #grab data at specified level and time.
v_lat = fV1.variables['nav_lat']
v_lon = fV1.variables['nav_lon']
#mask v so that white is plotted on land points
mu = V1 == 0
V1= np.ma.array(V1,mask=mu)

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

# get u and ugrid
u_vel2 = fU2.variables['vozocrtx']  #u currents and grid
U2 = u_vel2[timestamp,:,:,:] #grab data at specified level and time.
#mask u so that white is plotted on land points
mu =U2 == 0
U2 = np.ma.array(U2,mask=mu)

#get v and v grid
v_vel2 = fV2.variables['vomecrty']  #v currents and grid
V2 = v_vel2[timestamp,:,:,:] #grab data at specified level and time.
#mask v so that white is plotted on land points
mu = V2 == 0
V2= np.ma.array(V2,mask=mu)
#get sal
sal2=fT2.variables['vosaline']
temp2=fT2.variables['votemper']
ssh2=fT2.variables['sossheig']

#max currents at this time
currents1 = np.sqrt(U1**2+V1**2)
dims=currents1.shape
max1=currents1.max(); ind_max1 = np.argmax(currents1)
dim_max1 = np.unravel_index(ind_max1,dims)

print 'max currents at t=' + str(time/86400) + 'd' 
print 'nu=' +str(nu1) +':  max and index of max'
print max1, dim_max1

currents2 = np.sqrt(U2**2+V2**2)
dims=currents2.shape
max2=currents2.max(); ind_max2 = np.argmax(currents2)
dim_max2 = np.unravel_index(ind_max2,dims)

print 'nu=' +str(nu2) +':  max and index of max'
print max2, dim_max2

#plot U,V for both visocities
cmin=-2; cmax=2; #range on colourbar
ax=[125,350,200,450]

plt.figure(figsize=(13,14))
#u1
pylab.subplot(2,2,1) 
pylab.pcolormesh(U1[depthlevel,:,:],vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.title('U nu=' +str(nu1) + ' at z=' + str(depth) + 'm and t=' + str(time/86400) + 'd')
#pylab.axis(ax)
#v1
pylab.subplot(2,2,2)
pylab.pcolormesh(V1[depthlevel,:,:],vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.title('V nu=' + str(nu1) +' at z=' + str(depth) + 'm and t=' + str(time/86400) + 'd')
#pylab.axis(ax)
#u2
pylab.subplot(2,2,3)
pylab.pcolormesh(U2[depthlevel,:,:],vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.title('U nu=' + str(nu2)+' at z=' + str(depth) + 'm and t=' + str(time/86400) + 'd')
#pylab.axis(ax)
#v2
pylab.subplot(2,2,4)
pylab.pcolormesh(V2[depthlevel,:,:],vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.title('V nu=' + str(nu2)+' at z=' + str(depth) + 'm and t=' + str(time/86400) + 'd')
#pylab.axis(ax)

#salinty
minS=20
maxS=30
plt.figure(figsize=(14,8))
salP1 = sal1[timestamp,depthlevel,:,:]
mu =salP1 == 0
salP1= np.ma.array(salP1,mask=mu)

pylab.subplot(1,2,1) 
cmap = plt.get_cmap('winter') 
pylab.pcolormesh(salP1,cmap=cmap,vmin=minS,vmax=maxS)
pylab.colorbar()
pylab.plot(ibad,jbad,marker='o',color='blue')
pylab.plot(iplume,jplume,marker='o',color='red')
pylab.title('nu='+str(nu1))
#pylab.axis([260, 270, 340,350])

salP2 = sal2[timestamp,depthlevel,:,:]
mu =salP2 == 0
salP2= np.ma.array(salP2,mask=mu)
pylab.subplot(1,2,2)
cmap = plt.get_cmap('winter') 
pylab.pcolormesh(salP2,cmap=cmap,vmin=minS,vmax=maxS)
pylab.colorbar()
pylab.plot(ibad,jbad,marker='o',color='blue')
pylab.plot(iplume,jplume,marker='o',color='red')
pylab.title('nu='+str(nu2))

#temp
minT=8
maxT=20

plt.figure(figsize=(14,8))
tempP1 = temp1[timestamp,depthlevel,:,:]
mu =tempP1 == 0
tempP1= np.ma.array(tempP1,mask=mu)

pylab.subplot(1,2,1) 
cmap = plt.get_cmap('hot') 
pylab.pcolormesh(tempP1,cmap=cmap,vmin=minT,vmax=maxT)
pylab.colorbar()
pylab.plot(ibad,jbad,marker='o',color='blue')
pylab.plot(iplume,jplume,marker='o',color='red')
pylab.title('nu='+str(nu1))

tempP2 = temp2[timestamp,depthlevel,:,:]
mu =tempP2 == 0
tempP2= np.ma.array(tempP2,mask=mu)
pylab.subplot(1,2,2)
cmap = plt.get_cmap('hot') 
pylab.pcolormesh(tempP2,cmap=cmap,vmin=minT,vmax=maxT)
pylab.plot(ibad,jbad,marker='o',color='blue')
pylab.plot(iplume,jplume,marker='o',color='red')
pylab.title('nu='+str(nu2))
pylab.colorbar()

#Sea Surface Height
minSSH=-2
maxSSH=2

plt.figure(figsize=(14,8))
sshP1 = ssh1[timestamp,:,:]
mu =sshP1 == 0
sshP1= np.ma.array(sshP1,mask=mu)

pylab.subplot(1,2,1) 
cmap = plt.get_cmap('jet') 
pylab.pcolormesh(sshP1,cmap=cmap,vmin=minSSH,vmax=maxSSH)
pylab.colorbar()
pylab.plot(ibad,jbad,marker='o',color='blue')
pylab.plot(iplume,jplume,marker='o',color='red')
pylab.title('nu='+str(nu1))

sshP2= ssh2[timestamp,:,:]
mu =sshP2 == 0
sshP2= np.ma.array(sshP2,mask=mu)
pylab.subplot(1,2,2)
cmap = plt.get_cmap('jet') 
pylab.pcolormesh(sshP2,cmap=cmap,vmin=minSSH,vmax=maxSSH)
pylab.colorbar()
pylab.plot(ibad,jbad,marker='o',color='blue')
pylab.plot(iplume,jplume,marker='o',color='red')
pylab.title('nu='+str(nu2))

#Profiles of temp and salinity
print 't=' +str(times[timestamp]/86400)
plt.figure(figsize=(14,8))
axsal1 = [27,31,-100,0]
axsal2 = [0,35,-20,0]
#plot the salinity field at specified depth. 
ds=-depths[:]
#vertical profile at "bad point" from other runs
pylab.subplot(1,2,1) 
pylab.plot(sal1[timestamp,:,jbad-1,ibad-1],ds,label='nu='+str(nu1), marker='o')
pylab.plot(sal2[timestamp,:,jbad-1,ibad-1],ds,label='nu='+str(nu2), marker='o')
pylab.axis(axsal1)
pylab.title('Salinity near unstable/mixing region (blue dot)')
pylab.legend(loc=0)
#vertical profile at the bad point compared to output before blow up.
pylab.subplot(1,2,2) 
pylab.plot(sal1[timestamp,:,jplume-1,iplume-1],ds,label='nu='+str(nu1), marker='o')
pylab.plot(sal2[timestamp,:,jplume-1,iplume-1],ds,label='nu='+str(nu2), marker='o')
pylab.axis(axsal2)
pylab.title('Salinity near plume (red dot)')
pylab.legend(loc=0)

# Adjustment of the salinity field in the unstable region over one day
print 'Adjustment of the salinity field in the unstable region over one day'
plt.figure(figsize=(14,8))
for ii in range(0,7):
    pylab.subplot(1,2,1)
    tp=ii
    pylab.plot(sal1[tp,:,jbad-1,ibad-1],ds,label='t= '+str(times[tp]/86400) +'d')
pylab.title('nu ='+str(nu1))
pylab.axis(axsal1)
pylab.legend(loc=0)
for ii in range(0,7):
    tp=ii
    pylab.subplot(1,2,2)
    pylab.plot(sal2[tp,:,jbad-1,ibad-1],ds,label='t= '+str(times[tp]/86400) +'d')
pylab.title('nu ='+str(nu2))
pylab.axis(axsal1)
pylab.legend(loc=0)

#AAjustment of the salinity field in the plume region over the full week.
print 'Adjustment of the salinity field in the plume region over the full week.'
plt.figure(figsize=(14,8))
for ii in range(0,6):
    pylab.subplot(1,2,1)
    tp=6*ii+5
    pylab.plot(sal1[tp,:,jplume-1,iplume-1],ds,label='t= '+str(times[tp]/86400) +'d')
pylab.title('nu ='+str(nu1))
pylab.axis(axsal2)
pylab.legend(loc=0)
for ii in range(0,6):
    tp=6*ii+5
    pylab.subplot(1,2,2)
    pylab.plot(sal2[tp,:,jplume-1,iplume-1],ds,label='t= '+str(times[tp]/86400) +'d')
pylab.title('nu ='+str(nu2))
pylab.axis(axsal2)
pylab.legend(loc=0)

# for loop to print the maximum currents and location (grid points) over time. 
# First six ouput times are displayed.
# columns are max current, z, y, x
print 'nu=' + str(nu1)
maxes1=np.zeros((6,4))
for ii in range(0,6):
    U3d1 = u_vel1[ii,:,:,:]; V3d1 = v_vel1[ii,:,:,:];
    c1 = np.sqrt(U3d1**2+V3d1**2); dims=c1.shape
    max1=c1.max(); ind_max1 = np.argmax(c1)
    dim_max1 = np.unravel_index(ind_max1,dims)
    maxes1[ii,:] = [max1, dim_max1[0], dim_max1[1], dim_max1[2]];

print maxes1

print 'nu=' + str(nu2)
maxes2=np.zeros((6,4))
for ii in range(0,6):
    U3d2 = u_vel2[ii,:,:,:]; V3d2 = v_vel2[ii,:,:,:];
    c2 = np.sqrt(U3d2**2+V3d2**2); dims=c2.shape
    max2=c2.max(); ind_max2 = np.argmax(c2)
    dim_max2 = np.unravel_index(ind_max2,dims)
    maxes2[ii,:] = [ max2, dim_max2[0], dim_max2[1], dim_max2[2]];

print maxes2