#!/usr/bin/env python
# coding: utf-8
A notebook to make an animation of the surface salinity.
# In[1]:


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

from salishsea_tools import (nc_tools,viz_tools)

import os
import glob
from matplotlib import animation

import datetime

get_ipython().run_line_magic('matplotlib', 'inline')


# Define date range

# In[2]:


start = datetime.datetime(2015,1,1)
end = datetime.datetime(2015,8,26)

numdays = (end-start).days
dates = [start + datetime.timedelta(days=num)
             for num in range(0, numdays+1)]


# In[3]:


results_home = '/data/dlatorne/MEOPAR/SalishSea/nowcast/'
bathy = nc.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')


# In[4]:


def results_dataset(period, grid, results_dir):
    """Return the results dataset for period (e.g. 1h or 1d)
    and grid (e.g. grid_T, grid_U) from results_dir.
    """
    filename_pattern = 'SalishSea_{period}_*_{grid}.nc'
    filepaths = glob.glob(os.path.join(results_dir, filename_pattern.format(period=period, grid=grid)))
    return nc.Dataset(filepaths[0])


# Load files over date range

# In[5]:


#subset
ist=0
ien=350
jst=200
jen=650
depth=0
x=np.arange(ist,ien)
y=np.arange(jst,jen)

#initalize arrays
Us=np.zeros((1,y.shape[0],x.shape[0])); 
Vs=np.zeros((1,y.shape[0],x.shape[0]))
Ss=np.zeros((1,y.shape[0],x.shape[0]))
time=np.array([])

#lats and lons
lat = bathy.variables['nav_lat'][jst:jen,ist:ien]
lon = bathy.variables['nav_lon'][jst:jen,ist:ien]

#loop through files
period = '1d'
for date in dates:
    results_dir = os.path.join(results_home, date.strftime('%d%b%y').lower())
    grid_T = results_dataset(period,'grid_T',results_dir)
    
    #load variables
    S = grid_T.variables['vosaline'][:,0,jst:jen,ist:ien]

    Ss = np.append(Ss,S,axis=0)
    
    t = nc_tools.timestamp(grid_T, np.arange(S.shape[0]))
    t = t.datetime
    time = np.append(time, t)
    
Ss=Ss[1:,:,:];


# Testing out the size and colors before animating

# In[6]:


def salinity(t):
    ax.clear()
    #mesh
    mesh=ax.contourf(lon,lat,Ss[t],cs,cmap=cmap,extend='both')
    #land
    viz_tools.plot_land_mask(ax,bathy,coords='map',xslice=x,yslice=y,color='burlywood')
    
    #title
    timestamp = time[t]
    ax.set_title(timestamp.strftime('%d-%b-%Y %H:%M'))
    ax.set_xlabel('Longitude')
    ax.set_ylabel('Latitude')
    
    return mesh


# In[7]:


smin, smax, dels = 0, 34, 1
cs = np.arange(smin,smax)
cmap = plt.get_cmap('spectral')
st=5

fig, ax = plt.subplots(1, 1, figsize=(10, 10))
t=0


mesh = salinity(0)
cbar = plt.colorbar(mesh, ax=ax)
cbar.set_label('Practical Salinity [PSU]')


# In[8]:


#Setting up first frame and static content
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
mesh = salinity(0)
cbar = plt.colorbar(mesh, ax=ax)
cbar.set_label('Practical Salinity [PSU]')


#frmaes
framess=np.arange(1,Ss.shape[0])

#The animation function    
anim = animation.FuncAnimation(fig, salinity,frames=framess,  blit=True, repeat=False)

#A line that makes it all work
mywriter = animation.FFMpegWriter( fps=3, bitrate=10000)

#Save in current folder
anim.save('salinity-2015.mp4',writer=mywriter)


# In[ ]: