from __future__ import division

import math

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

from salishsea_tools import tidetools

%matplotlib inline

runname = (
    '20dec25dec', '26dec31dec', '1jan5jan', '6jan10jan', '11jan20jan', '21jan30jan',
    '31jan9feb', '10feb19feb', '20feb1mar', '2mar11mar', '12mar21mar', '22mar31mar', '1apr10apr',
    '11apr20apr', '21apr30apr', '1may10may', '11may20may', '21may30may', '31may9jun',
    '10jun19jun', '20jun29jun', '30jun9jul',
)
resultsloc = '/ocean/dlatorne/MEOPAR/SalishSea/results/spin-up/'
grid = NC.Dataset('../../nemo-forcing/grid/bathy_meter_SalishSea2.nc')
depth = grid.variables['Bathymetry'][:]

tidetools.plot_amp_phase_maps(runname, resultsloc, grid)

mod_M2_amp, mod_K1_amp, mod_M2_pha, mod_K1_pha = (
     tidetools.get_amp_phase_data(runname,resultsloc))

plt.figure(figsize=(10,5))
plt.subplot(1,2,1)
plt.pcolormesh(mod_M2_pha[840:,40:100],vmin=-105.,vmax=-75)
plt.plot(24,44,'*w')
plt.colorbar()
print mod_M2_pha[840+44,40+24], mod_M2_pha[840+44,40+24]+360.,' M2 measured = 271.8'
plt.subplot(1,2,2)
plt.pcolormesh(mod_K1_pha[840:,40:100],vmin=-125.,vmax=-95)
plt.plot(24,44,'*w')
plt.colorbar()
print mod_K1_pha[840+44,40+24], mod_K1_pha[840+44,40+24]+360.,' K1 measured = 260.'
print "Summary: M2 phases good, K1 phase low"

plt.figure(figsize=(10,5))
plt.subplot(1,2,1)
plt.pcolormesh(mod_M2_amp[840:,40:100],vmin=1,vmax=1.8)
plt.plot(24,44,'*w')
plt.colorbar()
print mod_M2_amp[840+44,40+24], ' M2 measured 1.17m'
plt.subplot(1,2,2)
plt.pcolormesh(mod_K1_amp[840:,40:100],vmin=0.4,vmax=1.)
plt.plot(24,44,'*w')
plt.colorbar()
print mod_K1_amp[840+44,40+24], ' K1 measured 0.558m'
print 'Summary: K1 amplitude okay, M2 about 50% too large'

plt.figure(figsize=(10,5))
plt.subplot(1,2,1)
plt.pcolormesh(mod_M2_pha[380:440,0:60]+360,vmin=-140+360,vmax=-105+360)
plt.plot(51,19,'*w')
plt.colorbar()
print mod_M2_pha[380+19,51]+360, 'M2 measured 241.1'
plt.subplot(1,2,2)
plt.pcolormesh(mod_K1_pha[380:440,0:60]+360,vmin=240,vmax=270)
plt.plot(51,19,'*w')
plt.colorbar()
print mod_K1_pha[380+19,51]+360, 'K1 measured 254.1'
print mod_K1_pha[380+19,3]+360
print 'Summary: M2 phase 3 degrees late, K1 phase 5 degrees late'

plt.figure(figsize=(10,5))
plt.subplot(1,2,1)
plt.pcolormesh(mod_M2_amp[380:440,0:60],vmin=.60,vmax=.80)
plt.plot(51,19,'*w')
plt.colorbar()
print mod_M2_amp[380+19,51], 'M2 measured 0.708m'
plt.subplot(1,2,2)
plt.pcolormesh(mod_K1_amp[380:440,0:60],vmin=.30,vmax=.60)
plt.plot(51,19,'*w')
plt.colorbar()
print mod_K1_amp[380+19,51], 'K1 measured 0.484m'
print 'Summary: amplitudes pretty good'

# now velocities
mod_M2_u_amp, mod_M2_u_pha, mod_M2_v_amp, mod_M2_v_pha, mod_K1_u_amp, mod_K1_u_pha, mod_K1_v_amp, mod_K1_v_pha = (
                    tidetools.get_composite_harms_uv(runname,resultsloc))

# phases
plt.figure(figsize=(10,10))
plt.subplot(2,2,1)
plt.pcolormesh(mod_M2_u_pha[840:,40:100]+360.,vmin=240.,vmax=280)
plt.plot(3,55,'ow')
plt.colorbar()
print "M2 u", mod_M2_u_pha[840+55,40+3]+360.,' M2 u measured about 244-273'
plt.subplot(2,2,2)
plt.pcolormesh(mod_K1_u_pha[840:,40:100],vmin=100.,vmax=190)
plt.plot(3,55,'ok')
plt.colorbar()
print "K1 u", mod_K1_u_pha[840+55,40+3],' K1 u measured = 130-160.'
plt.subplot(2,2,3)
plt.pcolormesh(mod_M2_v_pha[840:,40:100],vmin=60.,vmax=100)
plt.plot(3,55,'ow')
plt.colorbar()
print "M2 v", mod_M2_v_pha[840+55,40+3],' M2 v measured = 64-93'
plt.subplot(2,2,4)
plt.pcolormesh(mod_K1_v_pha[840:,40:100]+360,vmin=270.,vmax=400)
plt.plot(3,55,'ok')
plt.colorbar()
print "K1 v", mod_K1_v_pha[840+55,40+3]+360,' K1 v measured = 310-355'
print "Summary : K1 u late, Rest okay within errors"

# Velocity speeds.  For some reason these have to be multiplied by depth!
plt.figure(figsize=(10,15))
plt.subplot(3,2,1)
plt.pcolormesh(mod_M2_u_amp[840:,40:100]*depth[840:,40:100],vmin=0.0001,vmax=0.6)
plt.plot(3,55,'ow')
plt.colorbar()
print mod_M2_u_amp[840+55,40+3]*depth[840+55,40+3], ' M2 measured 0.22-0.32'
plt.subplot(3,2,2)
plt.pcolormesh(mod_K1_u_amp[840:,40:100]*depth[840:,40:100],vmin=0.0001,vmax=0.2)
plt.plot(3,55,'ow')
plt.colorbar()
print mod_K1_u_amp[840+55,40+3]*depth[840+55,40+3], ' K1 measured 0.04-0.10'
plt.subplot(3,2,3)
plt.pcolormesh(mod_M2_v_amp[840:,40:100]*depth[840:,40:100],vmin=0.0001,vmax=0.6)
plt.plot(3,55,'ow')
plt.colorbar()
print mod_M2_v_amp[840+55,40+3]*depth[840+55,40+3], ' M2 measured 0.18-0.26'
plt.subplot(3,2,4)
plt.pcolormesh(mod_K1_v_amp[840:,40:100]*depth[840:,40:100],vmin=0.0001,vmax=0.2)
plt.plot(3,55,'ow')
plt.colorbar()
print mod_K1_v_amp[840+55,40+3]*depth[840+55,40+3], ' K1 measured 0.03-0.08'
plt.subplot(3,2,5)
plt.pcolormesh( np.sqrt(mod_M2_u_amp[860:,38:70]**2+mod_M2_v_amp[860:,38:70]**2)*depth[860:,38:70],
               vmax = 0.414)
plt.colorbar()
plt.subplot(3,2,6)
plt.pcolormesh( np.sqrt(mod_K1_u_amp[860:,38:70]**2+mod_K1_v_amp[860:,38:70]**2)*depth[860:,38:70],
               vmax = 0.119)
plt.colorbar()
print 'Summary: velocities are similar to those input'

# Velocites at west boundary
plt.figure(figsize=(10,5))
plt.subplot(1,2,1)
plt.pcolormesh(mod_M2_u_amp[380:440,0:60]*depth[380:440,0:60],vmin=.0,vmax=.6)
plt.plot(51,19,'*w')
plt.colorbar()
print mod_M2_u_amp[380+30,20]*depth[380+30,20]
plt.subplot(1,2,2)
plt.pcolormesh(mod_K1_u_amp[380:440,0:60]*depth[380:440,0:60],vmin=.0,vmax=.5)
plt.plot(51,19,'*w')
plt.colorbar()
print mod_K1_u_amp[380+30,20]*depth[380+30,20]

meas_wl_harm, Am_M2_all, Ao_M2_all, gm_M2_all, go_M2_all, \
D_F95_M2_all, D_M04_M2_all, Am_K1_all, Ao_K1_all, gm_K1_all, \
go_K1_all, D_F95_K1_all, D_M04_K1_all = tidetools.calc_diffs_meas_mod(runname, resultsloc, grid)

tidetools.plot_scatter_pha_amp(Am_M2_all, Ao_M2_all, gm_M2_all, go_M2_all, 'M2')
tidetools.plot_scatter_pha_amp(Am_K1_all, Ao_K1_all, gm_K1_all, go_K1_all, 'K1')
print 'Summary: K1 amplitude good, K1 phase high, M2 amplitude low, M2 phase high'

#Plot the differences as circles of different areas on a map of the domain
tidetools.plot_diffs_on_domain(D_F95_M2_all, meas_wl_harm, 'F95', 'M2', grid)
tidetools.plot_diffs_on_domain(D_M04_M2_all, meas_wl_harm, 'M04', 'M2', grid)
tidetools.plot_diffs_on_domain(D_F95_K1_all, meas_wl_harm, 'F95', 'K1', grid)
tidetools.plot_diffs_on_domain(D_M04_K1_all, meas_wl_harm, 'M04', 'K1', grid)

fig = tidetools.plot_wlev_transect_map(grid, statnums)

runname1= '30jun9jul'
loc1 = resultsloc
runname2= ('10feb19feb','20feb1mar','2mar11mar','12mar21mar','22mar31mar', '1apr10apr', '11apr20apr')
loc2 = resultsloc
runname3= ('20dec25dec','26dec31dec','1jan5jan','6jan10jan','11jan20jan', '21jan30jan', '31jan9feb')
loc3 = loc2
runname4 = ('5nov14nov','15nov19nov','20nov29nov','30nov9dec','10dec19dec')
loc4 = loc2
runname5 = runname
loc5 = loc2statnums = np.array([54, 37, 0, 2, 3, 5, 4, 6, 8, 11, 13, 14, 15, 19, 20, 21, 22, 24, 25, 26, 63, 64,
                     56, 57,67,65,58,66,61,75,62])
tidetools.plot_wlev_const_transect('SS', statnums, runname1, loc1, grid, runname5, loc5)
#plot map that shows the location
tidetools.plot_wlev_transect_map(grid, statnums)


#plot the transect of some stations in approximate line through Salish Sea ('SS')
#Station numbers (1 based): 38, 1, 3, 4, 6, 5, 7, 9, 12, 14, 15, 16, 20, 21, 22, 23, 25, 26, 27
#Station numbers (0 based): 37, 0, 2, 3, 5, 4, 6, 8, 11, 13, 14, 15, 19, 20, 21, 22, 24, 25, 26
statnums = np.array([54, 37, 0, 2, 3, 5, 4, 6, 8, 11, 13, 14, 15, 19, 20, 21, 22, 24, 25, 26, 63, 64,
                     56, 57,67,65,58,66,61,75,62])
tidetools.plot_wlev_const_transect('SS', statnums, runname1, loc1, grid, runname5, loc5)
#plot map that shows the location
tidetools.plot_wlev_transect_map(grid, statnums)

#add a line through the US points into Puget Sound ('PS')
#Neah Bay (39), Sekiu (40), Port Angeles (41), Port Townsend (42), Foulweather Bluff (47), Everett (48), Seattle (45), Tacoma (44), Budd Inlet (43)
#Station numbers (1 based): 39,40,41,42,47,48,45,44,43
#Station numbers (0 based): 38,39,40,41,46,47,44,43,42
statnums = np.array([38,39,40,41,45,46,43,42])
tidetools.plot_wlev_const_transect('PS',statnums,runname1,loc1,grid,runname5,loc5)
#plot map that shows the location
tidetools.plot_wlev_transect_map(grid,statnums)

#plot all measured stations together, just to see 'em!
statnums = np.arange(0,76)
tidetools.plot_wlev_const_transect('all',statnums,runname1,loc1,grid,runname5,loc5)