sns.set_color_codes()

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')

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)

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

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

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])

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'])

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

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

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