from __future__ import division
import netCDF4 as NC
import numpy as np


%matplotlib inline
import matplotlib.pyplot as plt
from scipy.interpolate import Rbf
from scipy import interpolate
from salishsea_tools import (nc_tools,viz_tools,stormtools)
from scipy.interpolate import griddata

fB = NC.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc','r')
lats = fB.variables['nav_lat'][:]
lons = fB.variables['nav_lon'][:]
D = fB.variables['Bathymetry']



CGRF = '/ocean/dlatorne/MEOPAR/CGRF/NEMO-atmos/slp_y2006m02d07.nc'

C = NC.Dataset(CGRF,'r')
press=C.variables['atmpres'][:]
latCGRF=C.variables['nav_lat'][:]
lonCGRF=C.variables['nav_lon'][:]
lonCGRF=lonCGRF-360




#need to truncrate CGRF domain
xs=454; xe=472
ys=518; ye=531

print lonCGRF[xs,ys], lonCGRF[xe,ye]
print latCGRF[xs,ys], latCGRF[xe,ye]

t=23
rbfi = Rbf(lonCGRF[xs:xe,ys:ye], latCGRF[xs:xe,ys:ye],press[23,xs:xe,ys:ye],function='linear')

press_i=rbfi(lons,lats)

fig, ax = plt.subplots(1, 2, figsize=(15,10))
vmin=80000; vmax=102000
vs=np.arange(vmin,vmax,20)
mesh=ax[0].contourf(lons[:],lats[:],press_i,vs,vmin=vmin,vmax=vmax)
viz_tools.plot_coastline(ax[0],fB,coords='map')
cbar=fig.colorbar(mesh,ax=ax[0])
ax[0].set_xlim([-126.5,-121.5])
ax[0].set_ylim([47,51.2])

mesh=ax[1].contourf(lonCGRF[xs:xe,ys:ye], latCGRF[xs:xe,ys:ye],press[t,xs:xe,ys:ye],vs,vmin=vmin,vmax=vmax)
viz_tools.plot_coastline(ax[1],fB,coords='map')
cbar=fig.colorbar(mesh,ax=ax[1])
ax[1].set_xlim([-126.5,-121.5])
ax[1].set_ylim([47,51.2])

points=(lonCGRF[xs:xe,ys:ye].flatten(),latCGRF[xs:xe,ys:ye].flatten())
press_2 = griddata(points,press[t,xs:xe,ys:ye].flatten(), (lons, lats), method='nearest')

fig, ax = plt.subplots(1, 2, figsize=(15,10))
vmin=80000; vmax=102000
mesh=ax[0].contourf(lons[:],lats[:],press_2,20,vmin=vmin,vmax=vmax)
viz_tools.plot_coastline(ax[0],fB,coords='map')
cbar=fig.colorbar(mesh,ax=ax[0])
ax[0].set_xlim([-126.5,-121.5])
ax[0].set_ylim([47,51.2])
ax[0].grid(True)

mesh=ax[1].contourf(lonCGRF[xs:xe,ys:ye], latCGRF[xs:xe,ys:ye],press[t,xs:xe,ys:ye],20,vmin=vmin,vmax=vmax)
viz_tools.plot_coastline(ax[1],fB,coords='map')
cbar=fig.colorbar(mesh,ax=ax[1])
ax[1].set_xlim([-126.5,-121.5])
ax[1].set_ylim([47,51.2])
ax[1].grid(True)

#define parameters (in NEMO docs)
rho0=1035
P0=101000
g=9.81

#inverse barometer effect
eta= -1/g/rho0*(press_i-P0)

fig, ax = plt.subplots(1, 1, figsize=(15,10))
vmin=-1; vmax=1
mesh=ax.contourf(lons[:],lats[:],eta,20,vmin=vmin,vmax=vmax)
viz_tools.plot_coastline(ax,fB,coords='map')
cbar=fig.colorbar(mesh,ax=ax)
ax.set_xlim([-126.5,-121.5])
ax.set_ylim([47,51.2])

#Load the data
path='/data/nsoontie/MEOPAR/SalishSea/results/storm-surges/tide_fix/feb2006/'

runs = {'all_forcing','tidesonly','weather_only'}

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

for key in runs:
    fUs[key] = NC.Dataset(path + key +'/SalishSea_4h_20060201_20060207_grid_U.nc','r');
    fVs[key] = NC.Dataset(path + key +'/SalishSea_4h_20060201_20060207_grid_V.nc','r');
    fTs[key] = NC.Dataset(path + key +'/SalishSea_4h_20060201_20060207_grid_T.nc','r');

fig, axs = plt.subplots(1, 3, figsize=(15,8),sharey=True)
depthlevel=0
t = 41
vmin=-.5; vmax=.5
cmap='jet'
Us={}; Vs={}; Es ={}; Ss={}; Ts={};
for key in runs:
        [Us[key], Vs[key], Es[key], Ss[key], Ts[key]] = stormtools.get_variables(fUs[key],fVs[key],fTs[key],t,depthlevel)

ax=axs[0]
mesh=ax.pcolormesh(Es['all_forcing']-Es['tidesonly'],cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Last time of all_forcing-tidesonly')

ax=axs[1]
mesh=ax.pcolormesh(eta,cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Correction from removing \ninverse barometer')
    
ax=axs[2]
mesh=ax.pcolormesh(Es['all_forcing']-Es['tidesonly']-eta,cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Inverse barometer removed \nfrom weather_only-tidesonly')

print np.min(Es['all_forcing']-Es['tidesonly']-eta)
print np.max(Es['all_forcing']-Es['tidesonly']-eta)
print np.mean(Es['all_forcing']-Es['tidesonly']-eta)

fig, axs = plt.subplots(1, 3, figsize=(15,8),sharey=True)
depthlevel=0
t = 41
vmin=-.5; vmax=.5
cmap='jet'
Us={}; Vs={}; Es ={}; Ss={}; Ts={};
for key in runs:
        [Us[key], Vs[key], Es[key], Ss[key], Ts[key]] = stormtools.get_variables(fUs[key],fVs[key],fTs[key],t,depthlevel)

ax=axs[0]
mesh=ax.pcolormesh(Es['weather_only']-Es['tidesonly'],cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Last time of weather_only-tidesonly')

ax=axs[1]
mesh=ax.pcolormesh(eta,cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Correction from removing \ninverse barometer')
    
ax=axs[2]
mesh=ax.pcolormesh(Es['weather_only']-Es['tidesonly']-eta,cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Inverse barometer removed \nfrom weather_only-tidesonly')

print np.min(Es['weather_only']-Es['tidesonly']-eta)
print np.max(Es['weather_only']-Es['tidesonly']-eta)
print np.mean(Es['weather_only']-Es['tidesonly']-eta)

print press[20:24,xs:xe,ys:ye].shape
press4 = np.mean(press[20:24,xs:xe,ys:ye],axis=0)
t=23
rbfi = Rbf(lonCGRF[xs:xe,ys:ye], latCGRF[xs:xe,ys:ye],press4,function='linear')
press_i4=rbfi(lons,lats)

#inverse barometer effect
eta4= -1/g/rho0*(press_i4-P0)

fig, axs = plt.subplots(1, 3, figsize=(15,8),sharey=True)
depthlevel=0
t = 41
vmin=-.5; vmax=.5
cmap='jet'
Us={}; Vs={}; Es ={}; Ss={}; Ts={};
for key in runs:
        [Us[key], Vs[key], Es[key], Ss[key], Ts[key]] = stormtools.get_variables(fUs[key],fVs[key],fTs[key],t,depthlevel)

ax=axs[0]
mesh=ax.pcolormesh(Es['weather_only']-Es['tidesonly'],cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Last time of weather_only-tidesonly')

ax=axs[1]
mesh=ax.pcolormesh(eta4,cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Correction from removing \ninverse barometer')
    
ax=axs[2]
mesh=ax.pcolormesh(Es['weather_only']-Es['tidesonly']-eta4,cmap=cmap,vmin=vmin,vmax=vmax)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Surge (m)')
viz_tools.plot_coastline(ax,fB)
ax.set_title('Inverse barometer removed \nfrom weather_only-tidesonly')

print np.min(Es['weather_only']-Es['tidesonly']-eta4)
print np.max(Es['weather_only']-Es['tidesonly']-eta4)
print np.mean(Es['weather_only']-Es['tidesonly']-eta4)