# Figure 8 of paper "A locally smoothed terrain-following coordinate"
# submitted to JAMES
# Author: Stephanie Westerhuis
# Date: February 1, 2021

import matplotlib
import numpy as np
import xarray as xr
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import datetime as dt
import sys
from scipy import interpolate, stats
import math
from pathlib import Path
import cartopy
import matplotlib.image as mpimg
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
import seaborn as sns
import pickle

# compare FLS fraction bias for december 2016 for 
# SLEVE (1) and LOSVEC (6)

def load_obs_fcst_bias(setup=''):
    obs       = pickle.load(open(f'data/pickle_fls_fractions_{setup}/obs.p',       "rb"))
    fcst      = pickle.load(open(f'data/pickle_fls_fractions_{setup}/fcst.p',      "rb"))
    frac_bias = pickle.load(open(f'data/pickle_fls_fractions_{setup}/frac_bias.p', "rb"))

    return obs, fcst, frac_bias

def count_cases(b1, b6, b_min, d_min):
    '''
    count cases where sleve and losvec differ
    considerably

    input:
    b1: pandas.dataframe bias sleve
    b6: pandas.dataframe bias losvec
    b_min: minimal bias
    d_min: minimal difference
    '''

    dt1 = list(b1[b1<b_min].index)
    dt6 = list(b6[b6<b_min].index)
    dt1.extend(dt6)
    dt_all = set(dt1)

    count_sleve=0
    l_sleve = []
    count_losvec=0
    l_losvec = []
    count_neutral=0
    l_neutral = []

    for dd in dt_all: 
        if (b1[dd]-b6[dd] > d_min):
            count_sleve=count_sleve+1 
            l_sleve.append(dd)
        elif (b1[dd]-b6[dd] < -d_min):
            count_losvec=count_losvec+1 
            l_losvec.append(dd)
        else:
            count_neutral=count_neutral+1
            l_neutral.append(dd)

    print(f'  sleve better: {count_sleve}')
    print(f'  losvec better: {count_losvec}')
    print(f'  neutral cases: {count_neutral}')

    return l_sleve, l_losvec, l_neutral

# load pickled dataset containing forecasted FLS fractions
obs,fcst1,bias1 = load_obs_fcst_bias(100) # SLEVE
obs,fcst6,bias6 = load_obs_fcst_bias(673) # LOSVEC
lts = fcst1.columns  # leadtimes
valid_times = [0,3,6,9,12,15,18,21]

# disregard uninteresting cases
criteria = (obs.hc<0.1)&(obs.ccl>0.2)
oo = obs.loc[criteria]
f1 = fcst1.loc[criteria]
b1 = bias1.loc[criteria]
f6 = fcst6.loc[criteria]
b6 = bias6.loc[criteria]

color_1 = 'gainsboro'
color_6 = 'dimgrey'
color_l = 'darkgrey'    # whiskers

width = 4.8
fig, ax = plt.subplots(figsize=(width, width/1.618))

# prepare counting of dates
init_1 = [] # initialisations SLEVE
count_total_1 = 0
init_6 = [] # initialisations LOSVEC
count_total_6 = 0
b_min = -0.05  # minimum bias
d_min = 0.05  # minimum difference in bias

# loop over leadtimes
for (ind, lt) in enumerate(lts):
    print(f'\n leadtime = {lt}')

    # A) plot median fls fractions

    # f1
    obs_data = obs[obs.index.hour==lt]
    ax.bar(ind-0.2,f1[lt].median(),
           yerr=[[f1[lt].median()-f1[lt].quantile(0.25)],
                 [f1[lt].quantile(0.75)-f1[lt].median()]], 
           error_kw = {'linewidth':0.5},
           ecolor = color_l,
           color=color_1,
           capsize=0.6,
           width=0.4)
    # f6
    ax.bar(ind+0.2,f6[lt].median(),
           yerr=[[f6[lt].median()-f6[lt].quantile(0.25)],
                 [f6[lt].quantile(0.75)-f6[lt].median()]], 
           error_kw = {'linewidth':0.5},
           ecolor = color_l,
           color=color_6,
           capsize=0.6,
           width=0.4)

    # B) calculate related t-test statistics

    if len(f1[lt].dropna()) == len(f6[lt].dropna()):
        st, pvalue = stats.ttest_rel(f1[lt].dropna(), f6[lt].dropna())
        print(f'  sample size = {len(f1[lt].dropna())}')
        print(f'  p-value = {pvalue:.4f}%')
    
    # C) count cases where one configuration is clearly better

    l_1, l_6, l_neutral = count_cases(b1[lt], b6[lt], b_min, d_min)
    count_total_1 = count_total_1 + len(l_1)
    count_total_6 = count_total_6 + len(l_6)
    [init_1.append(ts - dt.timedelta(hours=lt)) for ts in l_1]
    [init_6.append(ts - dt.timedelta(hours=lt)) for ts in l_6]

# forecast runs which are clearly better for
#  SLEVE (1) or LOSVEC (6)
s_init_1 = set(init_1)
s_init_6 = set(init_6)

print(f'----------------')
print(f' # of forecast runs better with sleve:  {len(s_init_1)}')
print(f'  total cases: {count_total_1}')
print(f' # of forecast runs better with losvec: {len(s_init_6)}')
print(f'  total cases: {count_total_6}')

# plot median
ax.axhline(oo.ccl.median(), color='orchid')
ax.set_ylim(0,oo.ccl.median()+0.03)

# legend
patch_1 = Patch(color=color_1, label='SLEVE')
patch_6 = Patch(color=color_6, label='LOSVEC')
line_o  = Line2D([], [], color='orchid', label='OBS') 
ax.legend(handles = [patch_1, patch_6, line_o],
          loc = 'upper right',
          bbox_to_anchor = (0.995,0.95))


ax.set_xticks(np.arange(0,len(lts)))
ax.set_xticklabels(lts)
ax.set_xlabel('Leadtime [h]')
ax.set_ylabel('FLS fraction')
plt.tight_layout()

# save
out_name=f'8_fls_fraction'
fig.savefig(f'plots/{out_name}.png')

n_obs = len(oo)
print(f'Number of observations: {n_obs}')
print(f'Median observed FLS fraction: {oo.ccl.median():.2f}')
for vt in [0,3,6,9,12,15,18,21]:
    oo_vt = oo.between_time(f'{vt:02}:00',f'{vt:02}:00').ccl.median()
    print(f' at {vt:02} UTC: {oo_vt:.2f}')