%matplotlib inline
from matplotlib import pylab
import netCDF4 as NC
from numpy import sin, cos, pi
from numpy import arange, zeros

fT = NC.Dataset('../../input_files/SubDom_tide_elev.nc','r')
fU = NC.Dataset('../../input_files/SubDom_tide_ubar.nc','r')
fV = NC.Dataset('../../input_files/SubDom_tide_vbar.nc','r')
fB = NC.Dataset('../../input_files/SubDom_bathy_meter_NOBCchancomp.nc','r')


Depth = fB.variables['Bathymetry']
depth = Depth[:]
# masking just causes problems, set depth to 0 on land
depth[depth.mask] = 0
print depth.shape
print depth[180,:]

# Northern Boundary
iindex_N = 345-1
jlast = 398-1
# u is off-set from T and bathymetry so find values in between
depthu_N = zeros((jlast+1), dtype=float)
depthu_N[0:jlast-1] = 0.5*(depth[iindex_N,0:jlast-1]+depth[iindex_N,1:jlast])
depthu_N[jlast] = depth[iindex_N,jlast]

# plot to check        
pylab.plot (depthu_N,'x',depth[iindex_N,:],'r+')
pylab.xlim((175,205))

# now the V points, offset in the other direction
depthv_N = zeros((jlast+1), dtype=float)
depthv_N = 0.5*(depth[iindex_N,:]+depth[iindex_N-1,:])


# Western Boundary
ilast = 345-1
jindex_W = 1-1
# First the u-points
depthu_W = zeros((ilast+1), dtype=float)
depthu_W = 0.5*(depth[:,jindex_W]+depth[:,jindex_W+1])

# now V points
depthv_W = zeros((ilast+1), dtype=float)
depthv_W[0:ilast-1] = 0.5*(depth[0:ilast-1,jindex_W]+depth[1:ilast,jindex_W])
depthv_W[ilast] = depth[ilast,jindex_W]


# plot to check        
pylab.plot (depthv_W,'x',depth[:,jindex_W],'r+')
pylab.xlim((250,275))

etaelev = fT.variables['M2_amp_elev']
umax = fU.variables['M2_amp_u']
vmax = fV.variables['M2_amp_v']
print etaelev.shape, umax.shape, vmax.shape
etaphase = fT.variables['M2_phi_elev']
uphase = fU.variables['M2_phi_u']
vphase = fV.variables['M2_phi_v']

etaelev_N = etaelev[iindex_N,:]
umax_N = umax[iindex_N,:]/depthu_N   # divide by depth because CONCEPTS uses flux not current
vmax_N = vmax[iindex_N,:]/depthv_N
etaphase_N = etaphase[iindex_N,:]*pi/180.
uphase_N = uphase[iindex_N,:]*pi/180.
vphase_N = vphase[iindex_N,:]*pi/180.
etaZ1_N = etaelev_N*cos(etaphase_N)
etaZ2_N = etaelev_N*sin(etaphase_N)
uZ1_N = umax_N*cos(uphase_N)
uZ2_N = umax_N*sin(uphase_N)
vZ1_N = vmax_N*cos(vphase_N)
vZ2_N = vmax_N*sin(vphase_N)
print etaZ1_N.size

jindex_W = 1-1
etaelev_W = etaelev[:,jindex_W]
umax_W = umax[:,jindex_W]/depthu_W
vmax_W = vmax[:,jindex_W]/depthv_W
etaphase_W = etaphase[:,jindex_W]*pi/180.
uphase_W = uphase[:,jindex_W]*pi/180.
vphase_W = vphase[:,jindex_W]*pi/180.
etaZ1_W = etaelev_W*cos(etaphase_W)
etaZ2_W = etaelev_W*sin(etaphase_W)
uZ1_W = umax_W*cos(uphase_W)
uZ2_W = umax_W*sin(uphase_W)
vZ1_W = vmax_W*cos(vphase_W)
vZ2_W = vmax_W*sin(vphase_W)
# pad
uZ1_W[185:189+1] = uZ1_W[190]
            

pylab.plot(umax_N[177:197],'b+',depth[344,177:197]/100.,'rx',umax_N[177:197]/depthu_N[177:197]*20,'g+')

phi = arange(0,6.2,0.1)
eta_C = etaelev_W[200]*cos(phi-etaphase_W[200])
eta_N = etaZ1_W[200]*cos(phi)+etaZ2_W[200]*sin(phi)
pylab.plot(phi,eta_C,'+g',phi,eta_N,'-r')

nemo = NC.Dataset('JPP_bdytide_M2_grid_T.nc', 'w', format='NETCDF3_CLASSIC')
nemo.description = 'Tide data from WebTide via JPP'

# dimensions
nemo.createDimension('xb', etaZ1_N.size+etaZ1_W.size)
nemo.createDimension('yb', 1)

# variables
xb = nemo.createVariable('xb', 'int32', ('xb',))
xb.units = 'non dim'
xb.longname = 'counter around boundary'
yb = nemo.createVariable('yb', 'int32', ('yb',))
yb.units = 'non dim'
nbidta = nemo.createVariable('nbidta', 'int32' , ('yb','xb'))
nbidta.units = 'non dim'
nbidta.longname = 'i grid position'
nbjdta = nemo.createVariable('nbjdta', 'int32' , ('yb','xb'))
nbjdta.units = 'non dim'
nbjdta.longname = 'j grid position'
nbrdta = nemo.createVariable('nbrdta', 'int32' , ('yb','xb'))
nbrdta.units = 'non dim'
z1 = nemo.createVariable('z1','float32',('yb','xb'))
z1.units = 'm'
z1.longname = 'tidal elevation: cosine'
z2 = nemo.createVariable('z2','float32',('yb','xb'))
z2.units = 'm'
z2.longname = 'tidal elevation: sine'
    
# data
total_length = etaZ1_N.size+etaZ1_W.size
xb[:] = arange(1, total_length+1)
yb[0] = 1

etaZ1_W[etaZ1_W.mask] = 0.
etaZ1_N[etaZ1_N.mask] = 0.
etaZ2_W[etaZ2_W.mask] = 0.
etaZ2_N[etaZ2_N.mask] = 0.

   # West Boundary
nbidta[0,0:etaZ1_W.size] = 1
nbjdta[0,0:etaZ1_W.size] = arange(1,etaZ1_W.size+1)
z1[0,0:etaZ1_W.size] = etaZ1_W
z2[0,0:etaZ1_W.size] = etaZ2_W
   # North Boundary
nbidta[0,etaZ1_W.size:] = arange(1,etaZ1_N.size+1)
nbjdta[0,etaZ1_W.size:] = iindex_N+1
z1[0,etaZ1_W.size:] = etaZ1_N
z2[0,etaZ2_W.size:] = etaZ2_N    

    # and both
nbrdta[:] = 1

pylab.plot(nbidta[0,:etaZ1_W.size],nbjdta[0,:etaZ1_W.size],'bo',
           nbidta[0,etaZ1_W.size:],nbjdta[0,etaZ1_W.size:],'gx')

nemo.close()

u = NC.Dataset('JPP_bdytide_M2_grid_U.nc', 'w', format='NETCDF3_CLASSIC')
u.description = 'Tide data from WebTide via JPP'

#dimensions
u.createDimension('xb', etaZ1_N.size+etaZ1_W.size)
u.createDimension('yb', 1)

# variables
xb = u.createVariable('xb', 'int32', ('xb',))
xb.units = 'non dim'
xb.longname = 'counter around boundary'
yb = u.createVariable('yb', 'int32', ('yb',))
yb.units = 'non dim'
nbidta = u.createVariable('nbidta', 'int32' , ('yb','xb'))
nbidta.units = 'non dim'
nbidta.longname = 'i grid position'
nbjdta = u.createVariable('nbjdta', 'int32' , ('yb','xb'))
nbjdta.units = 'non dim'
nbjdta.longname = 'j grid position'
nbrdta = u.createVariable('nbrdta', 'int32' , ('yb','xb'))
nbrdta.units = 'non dim'
u1 = u.createVariable('u1','float32',('yb','xb'))
u1.units = 'm'
u1.longname = 'tidal x-velocity: cosine'
u2 = u.createVariable('u2','float32',('yb','xb'))
u2.units = 'm'
u2.longname = 'tidal x-velocity: sine'

# data
total_length = etaZ1_N.size+etaZ1_W.size
xb[:] = arange(1, total_length+1)
yb[0] = 1

uZ1_W[uZ1_W.mask] = 0.
uZ1_N[uZ1_N.mask] = 0.
uZ2_W[uZ2_W.mask] = 0.
uZ2_N[uZ2_N.mask] = 0.

   # West Boundary
nbidta[0,0:etaZ1_W.size] = 1
nbjdta[0,0:etaZ1_W.size] = arange(1,etaZ1_W.size+1)
u1[0,0:etaZ1_W.size] = uZ1_W
u2[0,0:etaZ1_W.size] = uZ2_W
   # North Boundary
nbidta[0,etaZ1_W.size:] = arange(1,etaZ1_N.size+1)
nbjdta[0,etaZ1_W.size:] = iindex_N+1
u1[0,etaZ1_W.size:] = uZ1_N
u2[0,etaZ2_W.size:] = uZ2_N    
    # and both
nbrdta[:] = 1

pylab.plot(nbidta[0,:etaZ1_W.size],nbjdta[0,:etaZ1_W.size],'bo',
           nbidta[0,etaZ1_W.size:],nbjdta[0,etaZ1_W.size:],'gx')

u.close()

v = NC.Dataset('JPP_bdytide_M2_grid_V.nc', 'w', format='NETCDF3_CLASSIC')
v.description = 'Tide data from WebTide via JPP'

# dimensions
v.createDimension('xb', etaZ1_N.size+etaZ1_W.size)
v.createDimension('yb', 1)

# variables
xb = v.createVariable('xb', 'int32', ('xb',))
xb.units = 'non dim'
xb.longname = 'counter around boundary'
yb = v.createVariable('yb', 'int32', ('yb',))
yb.units = 'non dim'
nbidta = v.createVariable('nbidta', 'int32' , ('yb','xb'))
nbidta.units = 'non dim'
nbidta.longname = 'i grid position'
nbjdta = v.createVariable('nbjdta', 'int32' , ('yb','xb'))
nbjdta.units = 'non dim'
nbjdta.longname = 'j grid position'
nbrdta = v.createVariable('nbrdta', 'int32' , ('yb','xb'))
nbrdta.units = 'non dim'
v1 = v.createVariable('v1','float32',('yb','xb'))
v1.units = 'm'
v1.longname = 'tidal y-velocity: cosine'
v2 = v.createVariable('v2','float32',('yb','xb'))
v2.units = 'm'
v2.longname = 'tidal y-velocity: sine'

# data
total_length = etaZ1_N.size+etaZ1_W.size
xb[:] = arange(1, total_length+1)
yb[0] = 1

vZ1_W[vZ1_W.mask] = 0.
vZ1_N[vZ1_N.mask] = 0.
vZ2_W[vZ2_W.mask] = 0.
vZ2_N[vZ2_N.mask] = 0.

   # West Boundary
nbidta[0,0:etaZ1_W.size] = 1
nbjdta[0,0:etaZ1_W.size] = arange(1,etaZ1_W.size+1)
v1[0,0:etaZ1_W.size] = vZ1_W
v2[0,0:etaZ1_W.size] = vZ2_W
   # North Boundary
nbidta[0,etaZ1_W.size:] = arange(1,etaZ1_N.size+1)
nbjdta[0,etaZ1_W.size:] = iindex_N+1
v1[0,etaZ1_W.size:] = vZ1_N
v2[0,etaZ2_W.size:] = vZ2_N    

    # and both
nbrdta[:] = 1

pylab.plot(nbidta[0,:etaZ1_W.size],nbjdta[0,:etaZ1_W.size],'bo',
           nbidta[0,etaZ1_W.size:],nbjdta[0,etaZ1_W.size:],'gx')

v.close()