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

# Load the data. Path name can be changed to look at different data.  
path = '/data/nsoontie/MEOPAR/SalishSea/results/north/tides/'

fU = NC.Dataset(path +'SalishSea_4h_20020915_20020915_grid_U.nc','r');
fV = NC.Dataset(path +'SalishSea_4h_20020915_20020915_grid_V.nc','r');
fT = NC.Dataset(path +'SalishSea_4h_20020915_20020915_grid_T.nc','r');
grid = NC.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc','r')
bathy = grid.variables['Bathymetry'][:,:]


u_vel = fU.variables['vozocrtx']  #u currents
t=fU.variables['time_counter']
ssh=fT.variables['sossheig']
v_vel=fV.variables['vomecrty'] # v currents
sal = fT.variables['vosaline']
depths = fT.variables['deptht']
temp=fT.variables['votemper']



plt.figure(figsize=(8,8))
ax=[30, 100,880,898]
cmin=0; cmax=400;
mu =bathy == 0
bathy= np.ma.array(bathy,mask=mu)
pylab.pcolor(bathy,vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.axis(ax)


#quick plot of the u field
depth=0
ts = 4

plt.figure(figsize=(10,10))
cmin=-2; cmax=2;

#plot U
U=u_vel[ts,depth,:,:]
print 'min/max U'
print U.min()
print U.max()
mu =U == 0
U= np.ma.array(U,mask=mu)
pylab.subplot(2,2,1)
pylab.pcolor(U,vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.title('U')
#pylab.axis(ax)

#plot V
pylab.subplot(2,2,2)
V=v_vel[ts,depth,:,:]
mu =V == 0
V= np.ma.array(V,mask=mu)
pylab.pcolor(V,vmin=cmin,vmax=cmax)
pylab.colorbar()
pylab.title('V')
#pylab.axis(ax)

print 'min/max V'
print V.min()
print V.max()

#plot U
mu =U == 0
U= np.ma.array(U,mask=mu)
pylab.subplot(2,2,3)
pylab.pcolor(U,vmin=-1,vmax=1)
pylab.colorbar()
pylab.title('U')
pylab.axis(ax)

#plot V
pylab.subplot(2,2,4)
V= np.ma.array(V,mask=mu)
pylab.pcolor(V,vmin=-1,vmax=1)
pylab.colorbar()
pylab.title('V')
pylab.axis(ax)

smin=20; smax=32;
tmin=7; tmax=25;

plt.figure(figsize=(16,8))

S=sal[ts,depth,:,:]
mu =S == 0
S= np.ma.array(S,mask=mu)
pylab.subplot(1,2,1)
pylab.pcolor(S,vmin=smin,vmax=smax,cmap='winter')
pylab.colorbar()
#pylab.axis(ax)

pylab.subplot(1,2,2)
T=temp[ts,depth,:,:]
mu =T == 0
T= np.ma.array(T,mask=mu)
pylab.pcolor(T,vmin=tmin,vmax=tmax,cmap='hot')
pylab.colorbar()
#pylab.axis(ax)

emin=-2;emax=2;

plt.figure(figsize=(8,8))

E=ssh[ts,:,:]
mu =E == 0
E= np.ma.array(E,mask=mu)
pylab.pcolor(E,vmin=emin,vmax=emax,cmap='jet')
pylab.colorbar()
#pylab.axis(ax)


pylab.plot(ssh[:,895,45])
pylab.xlabel('output time')
pylab.ylabel('ssh (m)')
pylab.axis([0,5,-1.5,1.5])