from __future__ import division, print_function
from salishsea_tools import (nc_tools,viz_tools,stormtools,tidetools)
from datetime import datetime, timedelta
from pylab import *
from sklearn import linear_model
import matplotlib.pyplot as plt
import scipy.io as sio
import netCDF4 as nc
import numpy as np
import math
import glob
import os
import datetime

%matplotlib inline

saline=sio.loadmat('/ocean/jieliu/research/meopar/autodataupdate/ferrydata/SBE1920150121.mat') #####need to change each time 
a=saline['ferryData']
b=a['data']
dataa = b[0,0]
time=dataa['matlabtime'][0,0]
lonn=dataa['Longitude'][0,0]
latt=dataa['Latitude'][0,0]
salinity=dataa['Practical_Salinity'][0,0]

a=len(time)
lon1=np.zeros([a,1])
lat1=np.zeros([a,1])
salinity1=np.zeros([a,1])
for i in np.arange(0,a):
    matlab_datenum = np.float(time[i])
    python_datetime = datetime.datetime.fromordinal(int(matlab_datenum)) + timedelta(days=matlab_datenum%1) - timedelta(days = 366)
    if((python_datetime.year == 2015) & (python_datetime.month == 1) & (python_datetime.day == 22) & (python_datetime.hour >= 2))&(python_datetime.hour <= 3):
        lon1[i]=lonn[i]
        lat1[i]=latt[i]
        salinity1[i]=salinity[i]

mask=lon1[:,0]!=0
lon1_2_4=lon1[mask]
lat1_2_4=lat1[mask]
salinity1_2_4=salinity1[mask]
lon11=lon1_2_4[0:-1:50]
lat11=lat1_2_4[0:-1:50]
salinity11=salinity1_2_4[0:-1:50]

def date(year, month, day_start, day_end, period, grid):
    
    day_range = np.arange(day_start, day_end+1)
    day_len = len(day_range)
    files_all = [None] * day_len
    inds = np.arange(day_len)
    
    for i, day in zip(inds, day_range):
        run_date = datetime.datetime(year,month, day)
        results_home = '/data/dlatorne/MEOPAR/SalishSea/nowcast/'
        results_dir = os.path.join(results_home, run_date.strftime('%d%b%y').lower())
        filename = 'SalishSea_' + period + '_' + run_date.strftime('%Y%m%d').lower() + '_' + run_date.strftime('%Y%m%d').lower() + '_' + grid + '.nc'
        file_single = os.path.join(results_dir, filename)
        files_all[i] = file_single
        
    print(files_all)

    return(files_all)

bathy, X, Y = tidetools.get_SS2_bathy_data()

def find_dist_3rd (q):
    k=0
    values =0
    dist = np.zeros(9)
    weights = np.zeros(9)
    value=np.zeros(9)
    regr =linear_model.LinearRegression()
    regr.fit(lon11,lat11);
    regr.coef_

    [x1, j1] = tidetools.find_closest_model_point(lon11[q],regr.predict(lon11[q]),X,Y,bathy,lon_tol=0.0052,lat_tol=0.00210,allow_land=False)
    for i in np.arange(x1-1,x1+2):
        for j in np.arange(j1-1,j1+2):
            dist[k]=tidetools.haversine(lon11[q],lat11[q],longitude[i,j],latitude[i,j])
            weights[k]=1.0/dist[k]
            value[k]=saline_nemo_3rd[i,j]*weights[k]
            values=values+value[k]
            k+=1
    return values,weights


def find_dist_4rd (q):
    k=0
    valuess=0
    dist = np.zeros(9)
    weights = np.zeros(9)
    value=np.zeros(9)
    regr =linear_model.LinearRegression()
    regr.fit(lon11,lat11);
    regr.coef_

    [x1, j1] = tidetools.find_closest_model_point(lon11[q],regr.predict(lon11[q]),X,Y,bathy,lon_tol=0.0052,lat_tol=0.00210,allow_land=False)
    for i in np.arange(x1-1,x1+2):
        for j in np.arange(j1-1,j1+2):
            dist[k]=tidetools.haversine(lon11[q],lat11[q],longitude[i,j],latitude[i,j])
            weights[k]=1.0/dist[k]
            value[k]=saline_nemo_4rd[i,j]*weights[k]
            valuess=valuess+value[k]
            k+=1
    return valuess,weights

def find_dist_5rd (q):
    k=0
    valuesss=0
    dist = np.zeros(9)
    weights = np.zeros(9)
    value=np.zeros(9)
    regr =linear_model.LinearRegression()
    regr.fit(lon11,lat11);
    regr.coef_

    [x1, j1] =tidetools.find_closest_model_point(lon11[q],regr.predict(lon11[q]),X,Y,bathy,lon_tol=0.0052,lat_tol=0.00210,allow_land=False)
    for i in np.arange(x1-1,x1+2):
        for j in np.arange(j1-1,j1+2):
            dist[k]=tidetools.haversine(lon11[q],lat11[q],longitude[i,j],latitude[i,j])
            weights[k]=1.0/dist[k]
            value[k]=saline_nemo_5rd[i,j]*weights[k]
            valuesss=valuesss+value[k]
            k+=1
    return valuesss,weights

aa=date(2015,1,22,22,'1h','grid_T') ####need to change for different daily model results
print(len(aa))
tracers=nc.Dataset(aa[0])
j=int(aa[0][65:67])
jj=int(aa[0][67:69])
latitude=tracers.variables['nav_lat'][:] 
longitude=tracers.variables['nav_lon'][:] 
saline_nemo = tracers.variables['vosaline']
saline_nemo_3rd = saline_nemo[2,0, 0:898, 0:398] ####if change cell(12), then this should also change
saline_nemo_4rd = saline_nemo[3,0, 0:898, 0:398] 
saline_nemo_5rd = saline_nemo[4,0, 0:898, 0:398] 

matrix=np.zeros([15,9])
values=np.zeros([15,1])
value_mean_3rd_hour=np.zeros([15,1])                      
for q in np.arange(0,15):
   values[q],matrix[q,:]=find_dist_3rd(q)
   value_mean_3rd_hour[q]=values[q]/sum(matrix[q])

matrixx=np.zeros([15,9])
valuess=np.zeros([15,1])
value_mean_4rd_hour=np.zeros([15,1])
                            
for q in np.arange(0,15):
   valuess[q],matrixx[q,:]=find_dist_4rd(q)
   value_mean_4rd_hour[q]=valuess[q]/sum(matrixx[q])

matrixxx=np.zeros([15,9])
valuesss=np.zeros([15,1])
value_mean_5rd_hour=np.zeros([15,1]) 
                            
for q in np.arange(0,15):
   valuesss[q],matrixxx[q,:]=find_dist_5rd(q)
   value_mean_5rd_hour[q]=valuesss[q]/sum(matrixxx[q])

fig, ax = plt.subplots(1, 1, figsize=(8, 8))
model_salinity_3rd_hour=plt.plot(lon11,value_mean_3rd_hour,'r',label="model salinity 2nd hour" )###sometimes need to change
model_salinity_4rd_hour=plt.plot(lon11,value_mean_4rd_hour,'k',label="model salinity 3rd hour" )
model_salinity_5rd_hour=plt.plot(lon11,value_mean_5rd_hour,'g',label="model salinity 4th hour" )
observation_salinity=plt.plot(lon11,salinity11,'b',label="observation salinity")

ax.set_title(str(j)+'_'+str(jj)+'_2-4_surface_salinity_comparison')
ax.set_xlabel('Longitude[degrees east]')
ax.set_ylabel('Practical Salinity')
ax.legend()
locs,labels = xticks()
xticks(locs, map(lambda x: "%g" % x, locs))
plt.annotate('Observations from Ocean Networks Canada', (0.3,-0.1), (0, 10), xycoords='axes fraction', textcoords='offset points', va='top')
fig.savefig(str(j)+'_'+str(jj)+'_2-4_surface_salinity_comparison.png')