from __future__ import division

import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np

from salishsea_tools import viz_tools
%matplotlib inline

grid = nc.Dataset('../../NEMO-forcing/grid/bathy_meter_SalishSea2.nc')
bathy = grid.variables['Bathymetry'][:]

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
mesh = ax.pcolormesh(bathy, cmap='winter_r')
fig.colorbar(mesh)
ax.plot(173,696,'ok')

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('winter_r')
cmap.set_bad('burlywood')
mesh = ax.pcolormesh(bathy, cmap=cmap)
fig.colorbar(mesh)
plt.axis((150, 250, 650, 750))
ax.plot(173,696,'ok')

data = nc.Dataset('/data/dlatorne/MEOPAR/SalishSea/nowcast/03apr15/SalishSea_1h_20150403_20150403_grid_T.nc')
salinity = data.variables['vosaline'][:]
surfaceheight = data.variables['sossheig'][:]

print salinity.shape
# mask salinity
m = salinity == 0
sal = np.ma.array(salinity, mask=m)
m = surfaceheight == 0
eta = np.ma.array(surfaceheight, mask=m)

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
mesh = ax.pcolormesh(sal[17,0,650:750,150:250], cmap=cmap)
fig.colorbar(mesh)

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
mesh = ax.pcolormesh(eta[17,650:750,150:250], cmap=cmap)
fig.colorbar(mesh)

atmos = nc.Dataset('/ocean/sallen/allen/research/MEOPAR/Operational/ops_y2015m04d03.nc')
precip = atmos.variables['precip'][:]
u_wind = atmos.variables['u_wind'][:]
v_wind = atmos.variables['v_wind'][:]
print precip.shape

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
mesh = ax.pcolormesh(precip[19])
ax.set_xlim((0,256))
ax.set_ylim((0,266))
fig.colorbar(mesh)

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
mesh = ax.pcolormesh(u_wind[19])
ax.set_xlim((0,256))
ax.set_ylim((0,266))
fig.colorbar(mesh)

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
mesh = ax.pcolormesh(v_wind[19])
ax.set_xlim((0,256))
ax.set_ylim((0,266))
fig.colorbar(mesh)

print u_wind[19,10,10], v_wind[19,10,10]
print u_wind[19].max(), u_wind[19].min()
print v_wind[19].max(), v_wind[19].min()

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
mesh = ax.pcolormesh(u_wind[18])
ax.set_xlim((0,256))
ax.set_ylim((0,266))
fig.colorbar(mesh)

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
mesh = ax.pcolormesh(u_wind[20])
ax.set_xlim((0,256))
ax.set_ylim((0,266))
fig.colorbar(mesh)