%matplotlib inline
from matplotlib import pylab
import matplotlib.pyplot as plt
import netCDF4 as NC
import numpy as np

fT = NC.Dataset('../../../nemo-forcing/open_boundaries/source_files/bc4SA_extractedFromBCCBaseline.nc','r')


print fT.file_format

for dim in fT.dimensions.values():
     print dim

for var in fT.variables.values():
    print var

latR = fT.variables['lat_rho'][:]
lonR = fT.variables['lon_rho'][:]

import sys
sys.path.append('../../bathymetry')
import bathy_tools

fB = NC.Dataset('../../../nemo-forcing/grid/bathy_meter_SalishSea.nc','r')
D = fB.variables['Bathymetry'][:]
lat = fB.variables['nav_lat'][:]
lon = fB.variables['nav_lon'][:]

import datetime
test = fT.variables['ocean_time'][:]
start_time = datetime.datetime(1990,1,1,0,0,0) 
time = np.array([start_time + datetime.timedelta(seconds=i) for i in test])    

t=15
fig = bathy_tools.plot_colourmesh(
    fB, 'Salish Sea Bathymetry, '+str(time[t]), 
    axis_limits=(-126.3, -122.15, 47.05, 51))
#plt.plot(lonR,latR,'kx')
plt.contourf(lonR,latR,fT.variables['salt'][t,29,:,:],np.arange(28,33,0.1),cmap='RdPu')
plt.xlim((-125.5,-124.0))
plt.ylim((48.0,49.0))

zr = fT.variables['zr'][:]
plt.pcolor(zr[0,:,:14,2])
plt.colorbar()

#plot the depth of the bottom cell from south (x=0) to north (x=14) at time=0
for tt in range(0,29):
    bottom_depth = zr[tt,0,:14,2]
    depth1 = zr[tt,10,:14,2]
    depth2 = zr[tt,20,:14,2]
    depth3 = zr[tt,28,:14,2]
    plt.plot(range(0,14),bottom_depth,'b')
    plt.plot(range(0,14),depth1,'r')
    plt.plot(range(0,14),depth2,'g')
    plt.plot(range(0,14),depth3,'m')
    plt.ylabel('depth of cell [m]')
    plt.xlabel('eta_rho [number of grid cells]')

salt = fT.variables['salt'][:]
plt.contourf(np.transpose(salt[:,:,4,2]))
plt.colorbar()

plt.plot(time,fT.variables['salt'][:,29,7,1],'r')
plt.plot(time,fT.variables['salt'][:,15,7,1],'b')
plt.plot(time,fT.variables['salt'][:,4,7,1],'g')
plt.plot(time,fT.variables['salt'][:,29,14,1],'r:')
plt.plot(time,fT.variables['salt'][:,15,14,1],'b:')
plt.plot(time,fT.variables['salt'][:,4,14,1],'g:')

#get the min and max for the colorbar
def get_cbar_lims(salt):
    cbarmin = round(np.min(salt))
    cbarmax = round(np.max(salt[salt<9e35]))+1
    return cbarmin, cbarmax

#define a function to find the index of the model cell closest to the required depth
def find_vert_cell_for_req_depth(req_depth,time,eta_rho_ind,xi_rho_ind):
    all_depths = zr[time,:,eta_rho_ind,xi_rho_ind]
    if np.logical_and(np.min(all_depths)<0,np.min(all_depths)<req_depth):
        bottom_depth = np.min(all_depths)
        #a range of tolerances is needed because depths do not increase evenly
        tols = np.arange(0,-10,-0.01)
        tol_count = 0
        vert_cell_index = np.array([])
        while np.size(vert_cell_index)==0:
            tol = tols[tol_count]
            vert_cell_index = np.where(np.logical_and(all_depths>(req_depth+tol),all_depths<(req_depth-tol)))
            tol_count = tol_count+1
        vert_cell_index = vert_cell_index[0].astype(int)
    else:
        vert_cell_index = np.ma.zeros((1,))
        vert_cell_index[0] = np.ma.masked
    return vert_cell_index[0]

a, b = fT.variables['lat_rho'].shape
req_depths = np.arange(0,-300,-20)
depth_ind = np.ma.zeros((a,b,len(req_depths)))
time = 0

for dep in range(0,len(req_depths)):
    for i in range(0,a):
        for j in range(0,b):
            depth_ind[i,j,dep] = find_vert_cell_for_req_depth(req_depths[dep],time,i,j)
depth_ind = depth_ind.astype(int)

print(depth_ind[:,:,0])
print a,b

#define a function to plot all salinity data at a given depth, with 24 subplots per year
def salinity_contours_all_years(depth,salt,time):
    cbarmin, cbarmax = get_cbar_lims(salt[:,depth,:,:])
    numplots = 24
    for t in range(0,len(time)/numplots):
        plt.figure(figsize=(20,9))
        for k in range(1,numplots+1):
            fig = plt.subplot(1,numplots,k)
            plt.contourf(salt[k+t*numplots-1,depth,:,:],np.arange(cbarmin,cbarmax,0.5))
            plt.title(time[k+t*numplots-1].strftime('%m/%y'))
            fig.axes.get_xaxis().set_ticks([])
            fig.axes.get_yaxis().set_ticks([])
        plt.colorbar()





salinity_contours_all_years(29,salt,time)


salinity_contours_all_years(20,salt,time)
    

salinity_contours_all_years(0,salt,time)


salinity_contours_all_years(0,temp,time)


salinity_contours_all_years(15,temp,time)