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

# Attempting a seasonal comparison at different depths between model and ONC data.

# In[1]:


import matplotlib.pyplot as plt
import seaborn as sns
import datetime
import glob
import netCDF4 as nc
import pandas as pd
import numpy as np
import os

from salishsea_tools import viz_tools

import ONC_patrols as onc

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


# In[2]:


sns.set_color_codes()


# In[3]:


badQC = [0,3,4,9]
files = glob.glob('/ocean/nsoontie/MEOPAR/ONC/Patrols/*.csv')
cols = sns.color_palette("hls", len(files))
grid_B = nc.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')
mesh_mask = nc.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/mesh_mask_SalishSea2.nc')


# In[4]:


def seasonal_cycle(data, depth_range, field):
    """Group data into months and apply mean/min/max to field in depth range given by drange. """
    
    data_depth = data[(data['Depth Corrected (m)'] >= depth_range[0]) & 
                      (data['Depth Corrected (m)'] <= depth_range[1])]
    
    grouped = data_depth.groupby([data['day'].dt.year, data['day'].dt.month])
    
    months=[]
    means=[]
    mins=[]
    maxes=[]
    for month, group in grouped:
        months.append(datetime.datetime(month[0],month[1],15))
        means.append(group[field].mean())
        mins.append(group[field].min())
        maxes.append(group[field].max())
    return months, np.array(means), np.array(maxes), np.array(mins)


# # Observed Seasonal cycles (like Rich's plots)

# In[5]:


fig,axm = plt.subplots(1,1)
fig,axs = plt.subplots(2,2,figsize=(15,6))
c=0
fields = ['Practical Salinity Corrected (psu)',
         'Temperature Corrected (C)']
dranges = [ (0,1), (119,120)]
for file in files:
    data = onc.load_patrol_csv(file)
    data = onc.exclude_bad(data,['Practical Salinity Corrected QC Flag  '], badQC)
    for field, ax_c in zip(fields,axs):
        for drange, ax in zip(dranges, ax_c):
            months, means,maxes,mins = seasonal_cycle(data,drange,field)
            yerr = [means-mins, maxes-means]
            ax.plot(months, means,'o-',color=cols[c])
            ax.errorbar(months, means, yerr=yerr, color=cols[c], fmt='o')
            ax.set_ylabel(field)
            ax.set_title('{}: Depth range {}-{} m'.format(field,drange[0],drange[1]))
            ax.set_xlim([datetime.datetime(2015,1,1),datetime.datetime(2016,3,31)])
    
    data.plot(x='Longitude Corrected (deg)',y='Latitude Corrected (deg)',kind='scatter',ax=axm,c=cols[c])
    c=c+1

viz_tools.plot_coastline(axm,grid_B,coords='map')
axm.set_ylim([48,51])


# # Model Comparison
# Idea: for each cast, look up model cast but average over area, depth range and month.

# In[6]:


def model_obs_data_frame(data, model_field, obs_field):
    """Build model data frame of daily mean, max, min of equivalent casts in observed data frame"""
    
    model_obs_data = pd.DataFrame({'day':[], 'Model Mean':[],
                               'Model Max':[],'Model Min':[],
                               'Depth Corrected (m)':[],
                               'obs': []})
    data_days, days = onc.list_days(data)
    for d in days:
        daily = data_days.get_group(d).dropna()
        daily_casts = daily.groupby('Cast')
        for c in daily_casts.groups:
            cast = daily_casts.get_group(c)
            lat, lon, date = onc.cast_position_and_time(cast)
            obs_depth = np.array(cast['Depth Corrected (m)'][:])
            try:
                model_d_interp, model_max, model_min = onc.retrieve_nowcast_data(
                    lon, lat, date, obs_depth, model_field, grid_B, mesh_mask)
                day_list = [date for i in range(model_d_interp.shape[0])]
                sub_data = pd.DataFrame({'day': day_list,
                                        'Model Mean': model_d_interp,
                                        'Model Max': model_max,
                                        'Model Min': model_min,
                                        'Depth Corrected (m)': obs_depth,
                                        'obs': np.array(cast[obs_field]),
                                        'Longitude Corrected (deg)':  np.array(cast['Longitude Corrected (deg)']),
                                        'Latitude Corrected (deg)':  np.array(cast['Latitude Corrected (deg)'])})
                model_obs_data = model_obs_data.append(sub_data,ignore_index=True)
            except IndexError:
                print(
                    'No Model Point for {} {}'.format(
                        cast['Longitude Corrected (deg)'].mean(),
                        cast['Latitude Corrected (deg)'].mean()))
    model_obs_data = model_obs_data.dropna()
    return model_obs_data


# In[7]:


def process_all_model_obs_dict(files, model_field, obs_field):
    """Create model-obs dictionary of data for all files listed in files."""
    
    storage_dict = {}
    
    for file in files:
        data = onc.load_patrol_csv(file)
        data = onc.exclude_bad(data,['Practical Salinity Corrected QC Flag  '], badQC)
        data_cast = onc.divide_into_casts(data)
        key = os.path.basename(file).split('_')[0]

        model_obs_data = model_obs_data_frame(data_cast, model_field, obs_field)

        storage_dict[key] = model_obs_data
    return storage_dict


# In[8]:


def plot_map(files,data_dict, ax, cols, grid_B):
    """Plot casts color coded on a map"""
    c=0
    for file in files:
        key = os.path.basename(file).split('_')[0]
        data_region = data_dict[key]
        data_region.plot(x='Longitude Corrected (deg)',y='Latitude Corrected (deg)',kind='scatter',ax=ax,c=cols[c])
        c=c+1
    
    viz_tools.plot_coastline(ax,grid_B,coords='map')
    ax.set_ylim([48,51])


# In[9]:


def fake_legend(ax):
    """Add a legend for model --s and obs -o"""
    lo,=ax.plot([0,0],[0,0],'o-',color='gray')
    lm,=ax.plot([0,0],[0,0],'s--',color='gray')
    ax.legend([lo,lm],['obs','model'])


# ### Process salinity and temperature

# In[10]:


data_dict =  process_all_model_obs_dict(files, 'vosaline', 'Practical Salinity Corrected (psu)')
data_dict_temp =  process_all_model_obs_dict(files, 'votemper', 'Temperature Corrected (C)')


# ### Overview

# In[11]:


def plot_overview(obs_field, data_dict, dranges, lims, cols, files):
    """Plot an overview of seasonal trends in data_dict at depth ranges listed in dranges.
    Compare model and observations."""
    
    fig,axs = plt.subplots(1,len(dranges),figsize=(20,3))

    c=0 
    for file in files:
        key = os.path.basename(file).split('_')[0]
        data_region = data_dict[key]
        for drange, ax,ylim in zip(dranges, axs,lims):
            obs_months, obs_means, _,_ = seasonal_cycle(data_region,drange,'obs')
            months, means,_,_ = seasonal_cycle(data_region,drange,'Model Mean')
            months, mins,_,_ = seasonal_cycle(data_region,drange,'Model Min')
            months, maxes,_,_ = seasonal_cycle(data_region,drange,'Model Max')
            yerr= [means-mins, maxes-means]
            ax.plot(obs_months, obs_means,'o-',color=cols[c])
            ax.plot(months, means,'s--',color=cols[c])
            ax.errorbar(months, means, yerr=yerr, color=cols[c], fmt='o--')
            ax.set_ylabel(obs_field)
            ax.set_title('{}: Depth range {}-{} m'.format(obs_field,drange[0],drange[1]))
            ax.set_xlim([datetime.datetime(2015,1,1),datetime.datetime(2016,3,31)])
            plt.setp( ax.xaxis.get_majorticklabels(), rotation=70 )
            fake_legend(ax)
            ax.set_ylim(ylim)

        c=c+1


# In[12]:


fig,ax=plt.subplots(1,1,figsize=(5,3))
plot_map(files, data_dict, ax, cols, grid_B)
#Salinity
obs_field = 'Practical Salinity Corrected (psu)'
dranges = [ (1,2), (25,50), (80,100)]
lims=[(10,32), (27,32.5), (29,34)]
plot_overview(obs_field, data_dict, dranges, lims, cols,files)
#Temperature
obs_field = 'Temperature Corrected (C)'
lims=[(5,20), (5,15), (7.5,12)]
plot_overview(obs_field, data_dict_temp, dranges, lims, cols,files)


# ### Biases

# In[13]:


def calculate_bias(data, drange):
    """Calculat difference between model means and observed means in a depth range binned by month."""
    obs_months, obs_means, _,_ = seasonal_cycle(data,drange,'obs')
    months, mod_means,_,_ = seasonal_cycle(data,drange,'Model Mean')
    months, mod_mins,_,_ = seasonal_cycle(data,drange,'Model Min')
    months, mod_maxes,_,_ = seasonal_cycle(data,drange,'Model Max')
    bias = mod_means - obs_means
    upper_lim = mod_maxes - obs_means
    lower_lim = mod_mins - obs_means
    return bias, upper_lim, lower_lim, months


# In[14]:


def plot_regional_bais(obs_field, data_dict, dranges, lims, cols,files):
    """Compare model and observed seasonal trends in depth ranges region by region."""
    c=0 
    fig,axs = plt.subplots(1,len(dranges),figsize=(20,3))
    for file in files:
        key = os.path.basename(file).split('_')[0]
        data_region = data_dict[key]
        for drange, ax, ylim in zip(dranges, axs,lims):
            model_bias, upper_lim, lower_lim, months = calculate_bias(data_region, drange)
            ax.plot(months, model_bias,'o-',color=cols[c])
            yerr=[model_bias - lower_lim, upper_lim-model_bias ]
            ax.errorbar(months,model_bias, yerr=yerr, color=cols[c], fmt='o--')
            ax.set_ylabel(obs_field)
            ax.set_title('Model bias: Depth range {}-{} m'.format(drange[0],drange[1]))
            ax.set_xlim([datetime.datetime(2015,1,1),datetime.datetime(2016,3,31)])
            ax.set_ylim(ylim)
            plt.setp( ax.xaxis.get_majorticklabels(), rotation=70 )

        c=c+1


# In[15]:


fig,ax=plt.subplots(1,1,figsize=(5,3))
plot_map(files, data_dict, ax, cols, grid_B)
#Salinity
obs_field = 'Practical Salinty Corrected (psu)'
dranges = [ (1,2), (25,50), (80,100)]
lims=[(-8,8), (-2,2), (-1.5,1.5)]
plot_regional_bais(obs_field, data_dict, dranges, lims, cols,files)
#Temperature
obs_field = 'Temperature Corrected (C)'
lims=[(-3,3), (-3,3), (-1.5,1.5)]
plot_regional_bais(obs_field, data_dict_temp, dranges, lims, cols,files)


# In[ ]: