import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from salishsea_tools import stormtools
from matplotlib import pylab

import netCDF4 as nc
from salishsea_tools import (
    nc_tools,
    viz_tools,)

%matplotlib inline

tracers = nc.Dataset('/ocean/imachuca/MEOPAR/data/results/plume/jun2003/SalishSea_1h_20030531_20030602_grid_T.nc')

tracersU = nc.Dataset('/ocean/imachuca/MEOPAR/data/results/plume/jun2003/SalishSea_1h_20030531_20030602_grid_U.nc')

tracersV = nc.Dataset('/ocean/imachuca/MEOPAR/data/results/plume/jun2003/SalishSea_1h_20030531_20030602_grid_V.nc')

nc_tools.show_variables(tracersU)

nc_tools.show_variables(tracersV)

nc_tools.show_dimensions(tracersU)

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

timesteps = tracers.variables['time_counter']

ugrid = tracersU.variables['vozocrtx']
vgrid = tracersV.variables['vomecrty']
zlevels = tracersU.variables['depthu']
timesteps = tracersU.variables['time_counter']
lats = tracers.variables['nav_lat']
lons = tracers.variables['nav_lon']

ugridshape = ugrid.shape
ugridshape

salvar = tracers.variables['vosaline']
salvar.shape

sal = salvar[:,0,:,:]

cmin = sal.min()
cmax = sal.max()

def sal_fxn(timeindex):
    
    sal_t = np.ma.masked_values(sal[timeindex], zlevel)

    fig, ax = plt.subplots(1, 1, figsize=(10, 8))
    viz_tools.set_aspect(ax)
    
    cmap = plt.get_cmap('Blues')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(sal_t, cmap=cmap)
    cbar = fig.colorbar(mesh)
    plt.xlim(xstart,xend)
    plt.ylim(ystart,yend)
    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    ax.set_title('t = {:.1f}h'.format((timesteps[timeindex] / 3600)-72))

def vel_fxn(timeindex):
    
    y_slice = np.arange(ystart, yend)
    x_slice = np.arange(xstart, xend)

    ugrid_tzyx = np.ma.masked_values(ugrid[timeindex, zlevel, y_slice, x_slice], 0)
    vgrid_tzyx = np.ma.masked_values(vgrid[timeindex, zlevel, y_slice, x_slice], 0)
    u_tzyx, v_tzyx = viz_tools.unstagger(ugrid_tzyx, vgrid_tzyx)
    speeds = np.sqrt(np.square(u_tzyx) + np.square(v_tzyx))
    max_speed = viz_tools.calc_abs_max(speeds)

    fig, ax = plt.subplots(1, 1, figsize=(10, 8))
    viz_tools.set_aspect(ax)

    ax.streamplot(x_slice[1:], y_slice[1:], u_tzyx, v_tzyx, 
                  density=1.2, linewidth=7*speeds/max_speed,
                  color='black')
    viz_tools.plot_land_mask(ax, bathy, xslice=x_slice, yslice=y_slice)

    ax.set_xlim(x_slice[0], x_slice[-1])
    ax.set_ylim(y_slice[0], y_slice[-1])
    ax.grid()

    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    ax.set_title('t = {:.1f}h'.format((timesteps[timeindex] / 3600)-72))

xstart=0
xend=ugridshape[3]
ystart=200
yend=600
timeindex=0
zlevel=0

sal_fxn(timeindex)

vel_fxn(timeindex)

def animateQ(i):
    
    plt.cla()
    
    sal_t = np.ma.masked_values(sal[i], zlevel)
    
    y_slice = np.arange(ystart, yend)
    x_slice = np.arange(xstart, xend)
    
    arrow_step = 25
    y_slice_a = y_slice[::arrow_step]
    x_slice_a = x_slice[::arrow_step]

    ugrid_tzyx = np.ma.masked_values(ugrid[i, zlevel, y_slice_a, x_slice_a], 0)
    vgrid_tzyx = np.ma.masked_values(vgrid[i, zlevel, y_slice_a, x_slice_a], 0)
    u_tzyx, v_tzyx = viz_tools.unstagger(ugrid_tzyx, vgrid_tzyx)
    speeds = np.sqrt(np.square(u_tzyx) + np.square(v_tzyx))
    
    fig1 = ax.pcolormesh(sal_t, cmap=cmap)
    
    if i==0:
        cbar = fig.colorbar(fig1)
    
    fig2 = ax.quiver(x_slice_a[1:], y_slice_a[1:], u_tzyx, v_tzyx, speeds,
                     pivot='mid', cmap='Reds', width=0.005)

    viz_tools.plot_land_mask(ax, bathy, xslice=x_slice, yslice=y_slice, color='burlywood')
        
    ax.set_xlim(x_slice[0], x_slice[-1])
    ax.set_ylim(y_slice[0], y_slice[-1])
    ax.grid()

    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    timestamp = nc_tools.timestamp(tracers,i)
    ax.set_title(timestamp)
    
    return [fig1,fig2]

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('Blues')
cmap.set_bad('burlywood')
animateQ(47)

def animate(i):
    
    plt.cla()
    
    sal_t = np.ma.masked_values(sal[i], zlevel)
    
    y_slice = np.arange(ystart, yend)
    x_slice = np.arange(xstart, xend)

    ugrid_tzyx = np.ma.masked_values(ugrid[i, zlevel, y_slice, x_slice], 0)
    vgrid_tzyx = np.ma.masked_values(vgrid[i, zlevel, y_slice, x_slice], 0)
    u_tzyx, v_tzyx = viz_tools.unstagger(ugrid_tzyx, vgrid_tzyx)
    speeds = np.sqrt(np.square(u_tzyx) + np.square(v_tzyx))
    max_speed = viz_tools.calc_abs_max(speeds)
    
    fig1 = ax.pcolormesh(sal_t, cmap=cmap)
    
    if i==0:
        cbar = fig.colorbar(fig1)
        
    fig2 = ax.streamplot(x_slice[1:], y_slice[1:], u_tzyx, v_tzyx, 
                  density=1.2, linewidth=7*speeds/max_speed,
                  color='black')
    viz_tools.plot_land_mask(ax, bathy, xslice=x_slice, yslice=y_slice, color='burlywood')
    
    ax.set_xlim(x_slice[0], x_slice[-1])
    ax.set_ylim(y_slice[0], y_slice[-1])
    ax.grid()

    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    timestamp = nc_tools.timestamp(tracers,i)
    ax.set_title(timestamp)
    
    return [fig1,fig2]

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('Blues')
cmap.set_bad('burlywood')
animate(47)

fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('Blues')
cmap.set_bad('burlywood')
plt.xlim(xstart,xend)
plt.ylim(ystart,yend)
ax.set_xlabel('x Index')
ax.set_ylabel('y Index')

def init():
    return [ax]

#The animation function    
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=48, interval=300, blit=True, repeat=False)

#A line that makes it all work
mywriter = animation.FFMpegWriter()

#Save in current folder
anim.save('SalVelJun_Animation.mp4',writer=mywriter,fps=1)

#Show as a pop-up window
#plt.show() 


fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('Blues')
cmap.set_bad('burlywood')
plt.xlim(xstart,xend)
plt.ylim(ystart,yend)
ax.set_xlabel('x Index')
ax.set_ylabel('y Index')

def init():
    return [ax]

#The animation function    
anim = animation.FuncAnimation(fig, animateQ, init_func=init, frames=48, interval=300, blit=True, repeat=False)

#A line that makes it all work
mywriter = animation.FFMpegWriter()

#Save in current folder
anim.save('SalVelJunQ_Animation.mp4',writer=mywriter,fps=1)

#Show as a pop-up window
#plt.show() 

def sal_fxn_broch(timeindex):
    
    sal_t = np.ma.masked_values(sal[timeindex], zlevel)

    fig, ax = plt.subplots(1, 1, figsize=(20, 15))
    viz_tools.set_aspect(ax)
    
    cmap = plt.get_cmap('gist_earth_r')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(sal_t, cmap=cmap)
    cbar = fig.colorbar(mesh)
    plt.xlim(xstart,xend)
    plt.ylim(ystart,yend)
    ax.set_xlabel('x Index', **{'size': '16'})
    ax.set_ylabel('y Index', **{'size': '16'})
    timestamp = nc_tools.timestamp(tracers,timeindex)
    ax.set_title(timestamp.format('DD-MMM-YYYY, HH:mm'), **{'size': '16'})
    cbar.set_label('{salvar.long_name} [{salvar.units}]'.format(salvar=salvar), **{'size': '16'})
    
    y_slice = np.arange(ystart, yend)
    x_slice = np.arange(xstart, xend)
    arrow_step = 15
    y_slice_a = y_slice[::arrow_step]
    x_slice_a = x_slice[::arrow_step]
    ugrid_tzyx = np.ma.masked_values(ugrid[timeindex, zlevel, y_slice_a, x_slice_a], 0)
    vgrid_tzyx = np.ma.masked_values(vgrid[timeindex, zlevel, y_slice_a, x_slice_a], 0)
    u_tzyx, v_tzyx = viz_tools.unstagger(ugrid_tzyx, vgrid_tzyx)
    speeds = np.sqrt(np.square(u_tzyx) + np.square(v_tzyx))
    ax.quiver(x_slice_a[1:], y_slice_a[1:], u_tzyx, v_tzyx, speeds,pivot='mid', cmap='autumn_r', width=0.0065)
    viz_tools.plot_land_mask(ax, bathy, xslice=x_slice, yslice=y_slice, color='burlywood')
    ax.set_xlim(x_slice[0], x_slice[-1])
    ax.set_ylim(y_slice[0], y_slice[-1])
    
    fig.patch.set_facecolor('#BEE3E1')

xstart=0
xend=ugridshape[3]
ystart=0
yend=ugridshape[2]
timeindex=42
zlevel=0
sal_fxn_broch(timeindex)

def unstagger(data_array,dim):
    varin=data_array
    size = data_array.shape;
    sizeU=size[1];
    sizeV=size[0];
    if dim == 'X':
        varout = 0.5*(data_array[:,0:sizeU-1] + data_array[:,1:sizeU]);
        varout = varout[1:sizeV,:]
    else:
        varout = 0.5*(varin[0:sizeV-1,:] + varin[1:sizeV,:]);
        varout = varout[:,1:sizeU];
    return varout

def sal_fxn_broch2(timeindex, zlevel):
    
    sal_t = np.ma.masked_values(sal[timeindex], zlevel)

    fig, ax = plt.subplots(1, 1, figsize=(20, 10))
    viz_tools.set_aspect(ax, coords = 'map', lats=lats)
    
    cmap = plt.get_cmap('gist_earth_r')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(lons[:], lats[:], sal_t[:,:], cmap=cmap)
    cbar = fig.colorbar(mesh)
    plt.axis((xstart, xend, ystart, yend))
    ax.set_xlabel('Longitude', **{'size': '16'})
    ax.set_ylabel('Latitude', **{'size': '16'})
    timestamp = nc_tools.timestamp(tracers,timeindex)
    ax.set_title('Surface Salinity and Currents Velocity Field\n'+ timestamp.format('DD-MMM-YYYY, HH:mm'), **{'size': '16'})
    cbar.set_label('{salvar.long_name} [{salvar.units}]'.format(salvar=salvar), **{'size': '16'})
    
    U = ugrid[timeindex,zlevel,:,:]
    V = vgrid[timeindex,zlevel,:,:]
    U_unstag = unstagger(U,'X')
    V_unstag = unstagger(V,'Y')

    Tsize = lats.shape
    end_lat=Tsize[0];
    end_lon=Tsize[1];
    T_lat_plot = lats[1:end_lat,1:end_lon];
    T_lon_plot = lons[1:end_lat,1:end_lon];
    theta=29 * np.pi / 180 #rotation angle in radians. 
    U_true = U_unstag * np.cos(theta) - V_unstag * np.sin(theta) #E velocity
    V_true = U_unstag * np.sin(theta) + V_unstag * np.cos(theta) #N velocity
    st=25
    speeds = np.sqrt(np.square(U_true) + np.square(V_true))
    ax.quiver(T_lon_plot[::st, ::st],T_lat_plot[::st, ::st],U_true[::st, ::st],V_true[::st, ::st],scale=15,color='OrangeRed',pivot='mid', width=0.0065)
    
    fig.patch.set_facecolor('#BEE3E1')

xstart=-124.8
xend=-122.4
ystart=48.2
yend=49.6
timeindex=42
zlevel=0
sal_fxn_broch2(timeindex, zlevel)