# imports
import numpy as np
import netCDF4 as NC

#get thalwag co-ordinates (provided by Nancy)
lines = np.loadtxt('/ocean/klesouef/meopar/tools/compare_tides/thalweg-coords.txt', delimiter=" ", unpack=False)
lines = lines.astype(int)

#get the model grid data
grid = NC.Dataset('/ocean/klesouef/meopar/nemo-forcing/grid/bathy_meter_SalishSea2.nc','r')
lon = grid.variables['nav_lon'][:]
lat = grid.variables['nav_lat'][:]
bathy = grid.variables['Bathymetry'][:]
grid.close()

#get the lat and lon from the model grid at each of the thalweg points
thalweg_lon = lon[lines[:,0],lines[:,1]]
thalweg_lat = lat[lines[:,0],lines[:,1]]

print thalweg_lon.shape
print thalweg_lat.shape

output = np.zeros((len(thalweg_lat),2))
output[:,0] = thalweg_lon
output[:,1] = thalweg_lat

#save these points for loading into MATLAB
np.savetxt('thalweg-lonlat.txt', output)

#get Foreman data
#data in /ocean/klesouef/meopar/tools/compare_tides
consts = np.loadtxt('Data_Files/foreman_m2_thalweg.txt')
m2_amp_foreman = consts[:,0]
m2_pha_foreman = consts[:,1]
m2_pha_foreman[m2_pha_foreman>180]=m2_pha_foreman[m2_pha_foreman>180]-360

consts = np.loadtxt('Data_Files/foreman_k1_thalweg.txt')
k1_amp_foreman = consts[:,0]
k1_pha_foreman = consts[:,1]
k1_pha_foreman[k1_pha_foreman>180]=k1_pha_foreman[k1_pha_foreman>180]-360

#get model data
from salishsea_tools import tidetools
runname = '40d'
resultsloc = '/ocean/dlatorne/MEOPAR/SalishSea/results/'
mod_M2_amp_40d, mod_K1_amp_40d, mod_M2_pha_40d, mod_K1_pha_40d = tidetools.get_amp_phase_data(runname,resultsloc)

runname = '40d','41d70d'
resultsloc = '/ocean/dlatorne/MEOPAR/SalishSea/results/'
mod_M2_amp_40d70d, mod_K1_amp_40d70d, mod_M2_pha_40d70d, mod_K1_pha_40d70d = tidetools.get_amp_phase_data(runname,resultsloc)

runname = '41d50d_nu200','51d60d_nu200','61d70d_nu200'
resultsloc = '/ocean/dlatorne/MEOPAR/SalishSea/results/'
mod_M2_amp, mod_K1_amp, mod_M2_pha, mod_K1_pha = tidetools.get_amp_phase_data(runname,resultsloc)

#rename our model results
m2_amp_model_thalweg_40d = mod_M2_amp_40d[lines[:,0],lines[:,1]]*100
m2_pha_model_thalweg_40d = mod_M2_pha_40d[lines[:,0],lines[:,1]]

m2_amp_model_thalweg_40d70d = mod_M2_amp_40d70d[lines[:,0],lines[:,1]]*100
m2_pha_model_thalweg_40d70d = mod_M2_pha_40d70d[lines[:,0],lines[:,1]]

m2_amp_model_thalweg = mod_M2_amp[lines[:,0],lines[:,1]]*100
m2_pha_model_thalweg = mod_M2_pha[lines[:,0],lines[:,1]]


k1_amp_model_thalweg_40d = mod_K1_amp_40d[lines[:,0],lines[:,1]]*100
k1_pha_model_thalweg_40d = mod_K1_pha_40d[lines[:,0],lines[:,1]]

k1_amp_model_thalweg_40d70d = mod_K1_amp_40d70d[lines[:,0],lines[:,1]]*100
k1_pha_model_thalweg_40d70d = mod_K1_pha_40d70d[lines[:,0],lines[:,1]]

k1_amp_model_thalweg = mod_K1_amp[lines[:,0],lines[:,1]]*100
k1_pha_model_thalweg = mod_K1_pha[lines[:,0],lines[:,1]]

#plot everything
thalnum = range(0,len(m2_amp_model_thalweg))

%matplotlib inline
import matplotlib.pylab as plt

#M2
plt.figure(figsize=(12,5))
#the Foreman model doesn't quite cover our domain, so ignore some points
plt.plot(thalnum[60:620],m2_amp_foreman[60:620],'b',label='Foreman')
plt.plot(thalnum,m2_amp_model_thalweg,'m',label=runname)
plt.plot(thalnum,m2_amp_model_thalweg_40d,'r',label='40d')
plt.plot(thalnum,m2_amp_model_thalweg_40d70d,'g',label='40d_41d70d')
plt.legend(loc=4)
plt.xlabel('thalweg point')
plt.ylabel('m2 amp [cm]')
plt.ylim([0,100])
plt.xlim([0,max(thalnum)])
plt.title('Foreman model and our model: M2 amplitude along thalweg')

plt.figure(figsize=(12,5))
plt.plot(thalnum[60:620],m2_pha_foreman[60:620],'b',label='Foreman')
plt.plot(thalnum,m2_pha_model_thalweg,'m',label=runname)
plt.plot(thalnum,m2_pha_model_thalweg_40d,'r',label='40d')
plt.plot(thalnum,m2_pha_model_thalweg_40d70d,'g',label='40d_41d70d')
plt.ylim([-150,150])
plt.legend(loc=4)
plt.xlim([0,max(thalnum)])
plt.xlabel('thalweg point')
plt.ylabel('m2 pha [deg]')
plt.title('Foreman model and our model: M2 phase along thalweg')


#K1
plt.figure(figsize=(12,5))
#the Foreman model doesn't quite cover our domain, so ignore some points
plt.plot(thalnum[60:620],k1_amp_foreman[60:620],'b',label='Foreman')
plt.plot(thalnum,k1_amp_model_thalweg,'m',label=runname)
plt.plot(thalnum,k1_amp_model_thalweg_40d,'r',label='40d')
plt.plot(thalnum,k1_amp_model_thalweg_40d70d,'g',label='40d_41d70d')
plt.legend(loc=4)
plt.xlabel('thalweg point')
plt.ylabel('k1 amp [cm]')
plt.ylim([0,100])
plt.xlim([0,max(thalnum)])
plt.title('Foreman model and our model: K1 amplitude along thalweg')

plt.figure(figsize=(12,5))
plt.plot(thalnum[60:620],k1_pha_foreman[60:620],'b',label='Foreman')
plt.plot(thalnum,k1_pha_model_thalweg,'m',label=runname)
plt.plot(thalnum,k1_pha_model_thalweg_40d,'r',label='40d')
plt.plot(thalnum,k1_pha_model_thalweg_40d70d,'g',label='40d_41d70d')
plt.ylim([-150,0])
plt.legend(loc=4)
plt.xlim([0,max(thalnum)])
plt.xlabel('thalweg point')
plt.ylabel('k1 pha [deg]')
plt.title('Foreman model and our model: K1 phase along thalweg')