import numpy as np
import netCDF4 as nc
import matplotlib.mlab as mlab
from matplotlib import pyplot as plt
from matplotlib import animation
import matplotlib.cm as cm
import glob
import os
import matplotlib.gridspec as gridspec
from salishsea_tools import (nc_tools,viz_tools,stormtools)

%matplotlib inline

bathy = nc.Dataset('/ocean/imachuca/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')

path = '/ocean/dlatorne/MEOPAR/SalishSea/results/spin-up/'
files = glob.glob(path+'*/SalishSea_1d_*grid_T.nc')
files.sort(key=os.path.basename)
#38 files, 3 are size 5 and 35 are 10. estimate 350+15 = 365 days
files

tracers = nc.Dataset(files[0])
nc_tools.show_variables(tracers)

tracers = nc.MFDataset(files)

timesteps = tracers.variables['time_counter']
sal = tracers.variables['vosaline']
ssh = tracers.variables['sossheig']
tem = tracers.variables['votemper']
zlv = tracers.variables['deptht']
zlv = -zlv[:]

tem.shape

!head thalweg.txt
thalweg = np.loadtxt('thalweg.txt', dtype=int, unpack=True)
thalweg[1].shape

def sal_fxn(ax, sal, t, part):
    sal0 = sal[:,0,:,:]
    cmin = sal0.min()
    cmax = sal0.max()
    
    cmap = plt.get_cmap('ocean_r')
    cmap.set_bad('burlywood')
    ax.set_axis_bgcolor('burlywood')
    viz_tools.set_aspect(ax)
    
    sal_t = np.ma.masked_values(sal0[t], 0)
    #mesh = ax.pcolormesh(sal_t, cmap=cmap)
    levels = np.arange(0, 33, 4)
    mesh = ax.contourf(sal_t, levels, cmap=cmap, extend='both')
    ax.set_xlabel('x Index'); ax.set_ylabel('y Index')
    ax.set_title('{label}'.format(label=sal.long_name.title()))
    if part == 'init':
        cbar = fig.colorbar(mesh, ax=ax)
        cbar.set_label('[{units}]'.format(units=sal.units))
    else:
        pass

def tem_fxn(ax, tem, t, part):
    tem0 = tem[:,0,:,:]
    cmin = tem0.min()
    cmax = tem0.max()
    
    cmap = plt.get_cmap('gist_rainbow_r')
    cmap.set_bad('burlywood')
    ax.set_axis_bgcolor('burlywood')
    viz_tools.set_aspect(ax)

    tem_t = np.ma.masked_values(tem0[t], 0)
    #mesh = ax.pcolormesh(tem_t, cmap=cmap)
    levels = np.arange(round(cmin), round(cmax), 2)
    mesh = ax.contourf(tem_t, levels, cmap=cmap, extend='both')
    ax.set_xlabel('x Index'); ax.set_ylabel('y Index')
    ax.set_title('{label}'.format(label=tem.long_name.title()))
    if part == 'init':
        cbar = fig.colorbar(mesh, ax=ax)
        cbar.set_label('[{units}]'.format(units=tem.units))
    else:
        pass

def ssh_fxn(ax, ssh, t, part):
    cmin = ssh[:].min()
    cmax = ssh[:].max()
    
    cmap = plt.get_cmap('terrain_r')
    cmap.set_bad('burlywood')
    ax.set_axis_bgcolor('burlywood')
    viz_tools.set_aspect(ax)
    
    ssh_t = np.ma.masked_values(ssh[t], 0)
    #mesh = ax.pcolormesh(ssh_t, cmap=cmap)
    levels = np.arange(round(cmin, 1), round(cmax, 1)+0.5, 0.25)
    mesh = ax.contourf(ssh_t, levels, cmap=cmap, extend='both')
    ax.set_xlabel('x Index'); ax.set_ylabel('y Index')
    ax.set_title('{label}'.format(label=ssh.long_name.title()))
    if part == 'init':
        cbar = fig.colorbar(mesh, ax=ax)
        cbar.set_label('[{units}]'.format(units=ssh.units))
        cbar.set_ticks(np.arange(-1.4, 1.9, 0.4))
    else:
        pass

def twg_fxn(ax, grid):
    
    cmap = plt.get_cmap('winter_r')
    cmap.set_bad('burlywood')
    viz_tools.set_aspect(ax)
    ax.set_axis_bgcolor('burlywood')
    
    bathy = grid.variables['Bathymetry']
    x_slice = np.arange(bathy.shape[1])
    y_slice = np.arange(0, 898)
    
    ax.pcolormesh(x_slice, y_slice, bathy[y_slice, x_slice], cmap=cmap)
    ax.plot(thalweg[1], thalweg[0], linestyle='-', marker='.', color='red', label='Thalweg Points')
    ax.set_xlabel('x Index'); ax.set_ylabel('y Index')
    ax.set_title('Thalweg Points')

def Twg_fxn(ax, sal, zlv, t, part):
    
    cmap = plt.get_cmap('hsv')
    cmap.set_bad('burlywood')
    ax.set_axis_bgcolor('burlywood')
    
    x, z = np.meshgrid(np.arange(thalweg.shape[1]), zlv)
    thw = np.ma.masked_values(sal[t, :, thalweg[0], thalweg[1]], 0)
    cs = [26,27,28,29,30,30.2,30.4,30.6,30.8,31,32,33,34]
    mesh=ax.contourf(x, z, thw, cs, cmap=cmap, extend='both')
    ax.set_xlim(0, thalweg[0][-1])
    ax.set_xlabel('Thalweg Points'); ax.set_ylabel('Depth')
    ax.set_title('{label} Profile along Thalweg'.format(label=sal.long_name.title()))
    
    if part == 'init':
        cbar = fig.colorbar(mesh, ax=ax)
        cbar.set_ticks([26,27,28,29,30,30.2,30.4,30.6,30.8,31,32,33,34])
        cbar.set_label('[{units}]'.format(units=sal.units))
    else:
        pass

fig = plt.figure(figsize=(18,12))
gs = gridspec.GridSpec(2, 4, height_ratios=[1.5,1])
ax0 = fig.add_subplot(gs[0, 0])
ax1 = fig.add_subplot(gs[0, 1])
ax2 = fig.add_subplot(gs[0, 2])
ax3 = fig.add_subplot(gs[0, 3])
ax4 = fig.add_subplot(gs[1, :])
twg_fxn(ax3, bathy)

def init():
    sal_fxn(ax0, sal, 0, 'init')
    tem_fxn(ax1, tem, 0, 'init')
    ssh_fxn(ax2, ssh, 0, 'init')
    Twg_fxn(ax4, sal, zlv, 0, 'init')
    timestamp = nc_tools.timestamp(tracers,0)
    plt.suptitle(timestamp.format('MMM DD'), fontsize=20, verticalalignment='bottom')
    
def animate(t):
    sal_fxn(ax0, sal, t, 'anim')
    tem_fxn(ax1, tem, t, 'anim')
    ssh_fxn(ax2, ssh, t, 'anim')
    Twg_fxn(ax4, sal, zlv, t, 'anim')
    timestamp = nc_tools.timestamp(tracers,t)
    plt.suptitle(timestamp.format('MMM DD'), fontsize=20, verticalalignment='bottom')

anim = animation.FuncAnimation(fig, animate, init_func=init, frames=365)
mywriter = animation.FFMpegWriter()
anim.save('Seminar_Profiles.mp4',writer=mywriter,fps=1, dpi=200, bitrate=1000000, codec="libx264")