#!/usr/bin/env python
# coding: utf-8

# # Notebook to Analyze April 2016 Drift Tracks

# In[1]:


import datetime
import glob
import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np
import scipy.io as sio
from salishsea_tools import tidetools
from salishsea_tools import viz_tools

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


# In[2]:


# Convert from Matlab time to Python Time

def convert_time(matlab_datenum):
    # from http://sociograph.blogspot.ca/2011/04/how-to-avoid-gotcha-when-converting.html
    python_datetime = (datetime.datetime.fromordinal(int(matlab_datenum)) 
                   + datetime.timedelta(days=matlab_datenum%1) 
                   - datetime.timedelta(days = 366))
     # there is a bizarre one hour error that I don't know the origin of
    python_datetime = python_datetime - datetime.timedelta(hours=1)
    return python_datetime


# In[3]:


# Grid for Background, using the same for both nowcast and nowcast green

grid = nc.Dataset('../../nemo-forcing/grid/bathy_meter_SalishSea2.nc')
bathy, nav_lon, nav_lat = tidetools.get_bathy_data(grid)


# In[4]:


# make the plots

def make_drifter_plot(runtype):
    fig, ax = plt.subplots(3, 3, figsize=(15, 20))
    cmap = plt.get_cmap('winter_r')
    cmap.set_bad('burlywood')
    
    directory = '/ocean/sallen/allen/research/MEOPAR/Ariane/AprilDrifters/'


    drifters = sio.loadmat(directory+'driftersPositions.mat')
    for drifterid in [1, 2, 3, 4, 5, 6, 311, 312, 313]:
        mattime = np.array([t[0] for t in drifters['drifters'][0][drifterid-1][4]])
        pytime = []
        for i, time in enumerate(mattime):
            time = convert_time(time)
            pytime.append(time)
    
        drifter ={'name': drifters['drifters'][0][drifterid-1][0][0],
          'id': drifters['drifters'][0][drifterid-1][1][0],
          'lat': np.array([t[0] for t in drifters['drifters'][0][drifterid-1][2]]),
          'lon': np.array([t[0] for t in drifters['drifters'][0][drifterid-1][3]]),
          'time': pytime,
          }

        if drifterid < 4:
            row = 0
            col = drifterid-1
        elif drifterid < 10:
            row = 1
            col = drifterid-1-3
        else:
            row = 2
            col = drifterid-311
        aspect = viz_tools.set_aspect(ax[row, col], coords='map', lats=nav_lat)
        mesh = ax[row, col].pcolormesh(nav_lon[:], nav_lat[:], bathy[:], cmap=cmap)
        cbar = fig.colorbar(mesh, ax=ax[row, col] )
        ax[row, col].axis((-123.6, -122.9, 48.4, 49.5))

        filepattern = 'traj_'+runtype+'_d_'+str(drifterid)+'_'
        for file in glob.glob(directory+filepattern+'*.nc'):
            traj_nowcast = nc.Dataset(file)

            for nt in range(54):
                ax[row, col].plot(traj_nowcast.variables['traj_lon'][:,nt], traj_nowcast.variables['traj_lat'][:,nt], 'ys')
            for nt in range(54,81):
                ax[row, col].plot(traj_nowcast.variables['traj_lon'][:,nt], traj_nowcast.variables['traj_lat'][:,nt], 'y^')
            ax[row, col].plot(traj_nowcast.variables['traj_lon'][:,0], traj_nowcast.variables['traj_lat'][:,0], 'wo')
    
        ax[row, col].plot(drifter['lon'][:], drifter['lat'][:], 'ro', label = 'Observations')
       
        ax[row, col].set_title('Drifter '+str(drifterid)+': Six Hour Segments')
        ax[row, col].legend()
        ax[row, col].set_ylabel('Latitude')
        ax[row, col].set_xlabel('Longitude')


# In[5]:


make_drifter_plot('nowcast')


# In[6]:


make_drifter_plot('nowcast-green')


# In[15]:


# look at a particular segment, for a particular drifter

def detail_drifter_plot(runtype, drifterid=1, segment_start=189, segment_end=161, lonmin=-123.6, lonmax=-122.9, latmin=48.4, latmax=49.5):
    fig, ax = plt.subplots(1, 1, figsize=(15, 20))
    cmap = plt.get_cmap('winter_r')
    cmap.set_bad('burlywood')
    
    directory = '/ocean/sallen/allen/research/MEOPAR/Ariane/AprilDrifters/'


    drifters = sio.loadmat(directory+'driftersPositions.mat')
    mattime = np.array([t[0] for t in drifters['drifters'][0][drifterid-1][4]])
    pytime = []
    for i, time in enumerate(mattime):
        time = convert_time(time)
        pytime.append(time)
    
    drifter ={'name': drifters['drifters'][0][drifterid-1][0][0],
          'id': drifters['drifters'][0][drifterid-1][1][0],
          'lat': np.array([t[0] for t in drifters['drifters'][0][drifterid-1][2]]),
          'lon': np.array([t[0] for t in drifters['drifters'][0][drifterid-1][3]]),
          'time': pytime,
          }

        
    aspect = viz_tools.set_aspect(ax, coords='map', lats=nav_lat)
    mesh = ax.pcolormesh(nav_lon[:], nav_lat[:], bathy[:], cmap=cmap)
    ax.axis((lonmin, lonmax, latmin, latmax))

    file = 'traj_'+runtype+'_d_'+str(drifterid)+'_i_'+str(segment_start)+'.nc'
    print (file)
        
    traj_nowcast = nc.Dataset(directory+file)

    for nt in range(54):
            ax.plot(traj_nowcast.variables['traj_lon'][:,nt], traj_nowcast.variables['traj_lat'][:,nt], 'ys')
    for nt in range(54,81):
            ax.plot(traj_nowcast.variables['traj_lon'][:,nt], traj_nowcast.variables['traj_lat'][:,nt], 'y^')
    ax.plot(traj_nowcast.variables['traj_lon'][:,0], traj_nowcast.variables['traj_lat'][:,0], 'wo')
    
    ax.plot(drifter['lon'][segment_end:segment_start], drifter['lat'][segment_end:segment_start], 'ro', label = 'Observations')
       
    ax.set_title('Drifter '+str(drifterid)+': Six Hour Segment')
    ax.legend()
    ax.set_ylabel('Latitude')
    ax.set_xlabel('Longitude')


# In[22]:


detail_drifter_plot('nowcast-green', lonmin=-123.6, lonmax=-123.1, latmin= 49, latmax=49.2, 
                   segment_end=189, segment_start=225)


# In[20]:


detail_drifter_plot('nowcast-green', lonmin=-123.6, lonmax=-123.1, latmin= 48.9, latmax=49.1)


# In[23]:


detail_drifter_plot('nowcast', lonmin=-123.6, lonmax=-123.1, latmin= 49, latmax=49.2, 
                   segment_end=189, segment_start=225)


# In[24]:


detail_drifter_plot('nowcast', lonmin=-123.6, lonmax=-123.1, latmin= 48.9, latmax=49.1)


# In[ ]: