import os
import glob

import netCDF4 as nc
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import datetime

import comparisons

from salishsea_tools import tidetools, viz_tools, teos_tools
from nowcast import analyze

import seaborn as sns

import ACTDR
import pandas as pd
%matplotlib inline
import math

sns.set_color_codes()
sns.set_context('talk')

nowcast_dir = '/results/SalishSea/nowcast/'
grid_B = nc.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')    
mesh = nc.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/mesh_mask_SalishSea2.nc')
FIRST_NOWCAST = datetime.datetime(2014, 10,27)

def get_model_cast(year, month, day, lon, lat, field):
    """Get the model daily mean, max and min corresponding to given year, month, day and lat/lon"""
    
    if field =='Salinity':
        var = 'vosaline'
    elif field == 'Temperature':
        var = 'votemper'
        
    date = datetime.datetime(year, month, day)
    early_days = False
    if date < FIRST_NOWCAST:
        early_days = True
    #Look up grid point
    bathy, X, Y= tidetools.get_bathy_data(grid_B)
    j, i = tidetools.find_closest_model_point(lon, lat, X, Y, bathy)
    # Grab data
    depths = mesh.variables['gdept'][0,:,j,i]
    if early_days:
        var_model = comparisons.early_model_data(date, j, i, '1h', 'grid_T', var, nowcast_dir)
    else:
        results_dir = os.path.join(nowcast_dir, date.strftime('%d%b%y').lower())
        grid_T_h = comparisons.results_dataset('1h', 'grid_T', results_dir)
        var_model = grid_T_h.variables[var][:, :, j, i]
    var_model, tmask = mask_model(j, i, var_model)
    # Daily mean/min/max
    mean_daily = np.mean(var_model, axis=0)
    max_daily = np.max(var_model, axis=0)
    min_daily = np.min(var_model, axis=0)
 
    return mean_daily, max_daily, min_daily, depths, tmask

def mask_model(j, i, var_model):
    """Mask a model variable at a grid point"""
    
    tmask = np.abs(1-mesh.variables['tmask'][:, :, j, i])
    tmask = tmask + np.zeros(var_model.shape)
    var_masked = np.ma.array(var_model, mask=tmask)
    return var_masked, tmask[0,:]

def prepare_cast_comparison(dep_obs, var_obs, year, month, day, lon, lat, field):
    """Gather model and observed cast data for comparison"""
    
    # model data 
    mean_mod, max_mod, min_mod, dep_mod, tmask = get_model_cast(year, month, day, lon, lat, field)
    
    # interp observations
    interp_obs = comparisons.interpolate_depth(var_obs, dep_obs, dep_mod)
    interp_obs = np.ma.array(interp_obs, mask=tmask)
    
    return interp_obs, mean_mod, max_mod, min_mod, dep_mod

data = pd.DataFrame(ACTDR.CTD_DAT)
data_fall = data[(data.Year>=2014) & (data.Month >=10)  ]

sns.set_context('talk')

def scatter_compare(data, field, cmin, cmax, ax, units):
    sal_flags=[]
    for dep_obs, var_obs, lon, lat, year, month, day in zip(data['Depth'],
                                                        data[field],
                                                        data['Longitude'],
                                                        data['Latitude'],
                                                        data['Year'],
                                                        data['Month'],
                                                        data['Day']):
        interp_obs, mean_mod, max_mod, min_mod, dep_mod = prepare_cast_comparison(dep_obs, var_obs,
                                                                                  year, month, day,
                                                                                  lon, lat, field)
        thres=1; date=datetime.datetime(year, month, day)
        if field == 'Salinity':
            flag, k=classify_surface_salinity(interp_obs, mean_mod, max_mod, min_mod, thres)
            if flag == 0:
                print('Salinity within {} at depth {}, {}, {}, {}'.format(thres, dep_mod[k], lon, lat, date))
            if flag == 1:
                print('Too salty at depth {}, {}, {}, {}'.format(dep_mod[k], lon, lat, date))
            if flag == -1:
                print('Too fresh at depth {}, {}, {}, {}'.format(dep_mod[k], lon, lat, date))
            sal_flags.append(flag)
        ax.errorbar(interp_obs, mean_mod, yerr=[mean_mod-min_mod,max_mod-mean_mod ],fmt='k:', marker='',
                     zorder=0, ecolor='gray',lw=1,)
        sc = ax.scatter(interp_obs, mean_mod, edgecolor='',c= dep_mod, 
                         norm=mcolors.LogNorm(), cmap='Spectral', vmin=0.5, vmax=300, s=50)
        ax.set_xlim([cmin, cmax])
        ax.set_ylim([cmin, cmax])
        ax.plot([cmin, cmax],[cmin, cmax], 'r-',zorder=0)
    cb = plt.colorbar(sc, ax=ax)
    ticks = [1, 10, 25, 50, 100, 200, 400]
    cb.set_ticks(ticks)
    cb.set_ticklabels(ticks)
    cb.set_label('Depth [m]')
    ax.set_xlabel('Observed {} [{}]'.format(field, units))
    ax.set_ylabel('Modelled {} [{}]'.format(field, units))
    return sal_flags

def compare_region(data, lon_min, lon_max, lat_min, lat_max, title, smin=18, smax=32, tmin=8, tmax=16):
    cmap = sns.diverging_palette(220, 20,as_cmap=True)
    data_region  = comparisons.isolate_region(data,lon_min, lon_max, lat_min, lat_max)
    #sal
    fig2,ax2 = plt.subplots(1,1)
    field = 'Salinity'; 
    sal_flags=scatter_compare(data_region, field, smin,smax,ax2,'g/kg')
    ax2.set_title('{} - {}'.format(title, field))
    #temp
    fig3,ax3 = plt.subplots(1,1)
    field = 'Temperature';
    scatter_compare(data_region, field, tmin,tmax,ax3,'deg C')
    ax3.set_title('{} - {}'.format(title, field))
    #map
    fig1, ax1 = plt.subplots(1,1)
    ax1.scatter(data_region['Longitude'], data_region['Latitude'], s=50,c=sal_flags, cmap=cmap,vmin=-1,vmax=1)
    viz_tools.plot_coastline(ax1,grid_B, coords='map')
    ax1.set_xlim([lon_min, lon_max])
    ax1.set_ylim([lat_min, lat_max])
    ax1.set_title(title)
   

    return fig1, ax1

def classify_surface_salinity(interp_obs, mean_mod, max_mod, min_mod, thres):

    interp_obs = np.ma.masked_invalid(interp_obs)
    nonmasked = np.ma.flatnotmasked_edges(interp_obs)
    k=nonmasked[0]
    if (min_mod[k] -interp_obs[k] >thres) :
        flag = 1 #salty surface
    elif (max_mod[k] - interp_obs[k] )<-thres:
        flag = -1 #fresh surface
    else:
        flag =0  #ok
    
    return flag, k