import netCDF4 as nc
import numpy as np
from salishsea_tools import tidetools, nc_tools
import matplotlib.pyplot as plt

%matplotlib inline

grid=nc.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')
bathy,X,Y=tidetools.get_bathy_data(grid)

lines = np.loadtxt(
        '/data/nsoontie/MEOPAR/tools/bathymetry/thalweg_working.txt',
        delimiter=" ", unpack=False)
lines = lines.astype(int)


lines.shape

plt.plot(-bathy[lines[:,0],lines[:,1]])

ie=1100

thalweg_nonorth=-bathy[lines[0:ie,0],lines[0:ie,1]]
plt.plot(thalweg_nonorth)

imax=np.argmax(thalweg_nonorth[1000:ie])
thalweg_nonorth[1000:ie][imax]


ind=np.where(thalweg_nonorth==thalweg_nonorth[1000:ie][imax])
ind

indp = ind[0][-1]
indp

slope = (0-thalweg_nonorth[indp])/(len(thalweg_nonorth)-1-indp)
thalweg_right=thalweg_nonorth.copy()
thalweg_right[indp:] = (np.arange(indp,len(thalweg_nonorth))-(len(thalweg_nonorth)-1))*slope
plt.plot(thalweg_right)

thalweg_right[-1]

def smooth(x,window_len,window):
    s=np.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]]
    if window == 'flat': #moving average
        w=np.ones(window_len,'d')
    else:
         w=eval('np.'+window+'(window_len)')

    y=np.convolve(w/w.sum(),s,mode='valid')
    y= y[(window_len/2-1):-(window_len/2)]

    return y

fig=plt.figure(figsize=(10,5))
w_len=30
windows=['flat','hanning','hamming', 'bartlett','blackman']
plt.plot(thalweg_right,label='original')
for w in windows:
    y=smooth(thalweg_right,w_len,w)
    print y.shape
    plt.plot(y,label=w)

plt.legend(loc=0)

w_len=30
y=smooth(thalweg_right,w_len,'blackman')
y[-1]=0;#keep right end closed
thalweg_smooth=y
plt.plot(thalweg_smooth,label='smoothed')

y[-1]

dx=500
# x is variable for horiztonal space
x=dx*np.linspace(0,thalweg_smooth.shape[0]-1,thalweg_smooth.shape[0]);

plt.plot(x/1000,thalweg_smooth)


Ny=10; dy=500
y = dy*np.linspace(0,Ny-1,Ny)

thalweg_y=np.tile(thalweg_smooth,Ny)
thalweg_y=thalweg_y.reshape((y.shape[0],thalweg_smooth.shape[0]))
thalweg_y.shape


plt.plot(thalweg_y[0,:])

xx,yy=np.meshgrid(x,y)
xx.shape

plt.pcolormesh(xx,yy,thalweg_y)
plt.colorbar()

thalweg_final=-thalweg_y #positive bathy values for saving

new_bathy = nc.Dataset('../bathy2D.nc', 'w')
new_bathy.createDimension('y', thalweg_final.shape[0])
new_bathy.createDimension('x', thalweg_final.shape[1])
nc_tools.show_dimensions(new_bathy)

new_x = new_bathy.createVariable('x', float, ('y', 'x'), zlib=True)
new_x.setncattr('units', 'metres')
new_y = new_bathy.createVariable('y', float, ('y', 'x'), zlib=True)
new_y.setncattr('units', 'metres')
newdepths = new_bathy.createVariable(
    'Bathymetry', float, ('y', 'x'), 
    zlib=True, least_significant_digit=1)
newdepths.setncattr('units', 'metres')

new_x[:]=xx;
new_y[:]=yy;
newdepths[:]=thalweg_final

new_bathy.title="""Thalweg Bathymetry for 2D model"""
new_bathy.institution= """ 
Dept of Earth, Ocean & Atmospheric Sciences, University of British Columbia"""
new_bathy.comment= """
Based on bathy_meter_SalishSea2.nc along thalweg with smoothing"""
new_bathy.reference= """
/data/nsoontie/MEOPAR/2Ddomain/notebooks/Generate_2D_bathy.ipynb"""

nc_tools.show_dataset_attrs(new_bathy)


new_bathy.close()

B=nc.Dataset('bathy2D.nc')
thal=B.variables['Bathymetry']
plt.pcolormesh(thal[:]); plt.colorbar()

plt.plot(-thal[0,:])

thal[0,-1]