# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.19.3
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %%
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from oggm import utils

FILTERED_DIR = 'filtered'
EXP = 'lmrseas_oldBE'
DROP_ENS = ['ens124']  # only 5 of the 6 alternative hist ensemble members are wanted here

# %%
exp_dir = f'{FILTERED_DIR}/{EXP}'
vol_nat_ens = pd.read_csv(f'{exp_dir}/volume_km3_nat_ens_reg11.csv', index_col=0)
vol_hist_ens = pd.read_csv(f'{exp_dir}/volume_km3_hist_ens_reg11.csv', index_col=0).drop(columns=DROP_ENS)
dfok = pd.read_csv(f'{exp_dir}/glacier_statistics_reg11.csv', index_col=0)
inv_volume_2003 = dfok.inv_volume_km3.sum()

# %%
dfo = pd.read_csv("../round1/filtered_update/Area_volume_change_LIA_Alps.csv",
                   encoding="ISO-8859-1", index_col=0, skiprows=[1]).loc[17]

dfg = utils.get_geodetic_mb_dataframe(regional=True)
dfg = dfg.loc[dfg.period == '2000-01-01_2020-01-01']
dfg = dfg.loc[dfg.index == 11]
dvol = dfg.iloc[0]['dmdt'] * 20 / 900 * 1000  # Gt -> km3
dvol_err = dfg.iloc[0]['err_dmdt'] * 20 / 900 * 1000

# %%
# Quantile ranking of the 5 alt. hist ensemble members by their 1850 (LIA) volume
QUANTILE_ORDER = [
    ('min', 'ens123'),
    ('0.25', 'ens053'),
    ('median', 'ens183'),
    ('0.75', 'ens153'),
    ('max', 'ens178'),
]

# %% [markdown]
# ## Same style as talk/talk_5.png, but with the 5 alternative hist ensemble
# ## members instead of the single original hist run.

# %%
with sns.axes_style("whitegrid"):
    plt.figure(figsize=(10, 5.5))

    vol_nat_ens_plot = vol_nat_ens.loc[:1850]
    for i, col in enumerate(vol_nat_ens_plot.columns):
        plt.plot(vol_nat_ens_plot.index, vol_nat_ens_plot[col].values,
                 color='C0', alpha=0.15, linewidth=0.8,
                 label='Ensemble' if i == 0 else '_nolegend_')
    vol_nat_ens_plot.mean(axis=1).plot(label='Nat Ensemble Mean', color='navy', linewidth=2)

    colors = plt.cm.Reds(np.linspace(0.45, 0.95, len(QUANTILE_ORDER)))
    for (qname, col), c in zip(QUANTILE_ORDER, colors):
        vol_hist_ens[col].plot(label=f'LMR seas. + BE "old" ({qname})', color=c, linewidth=2)

    ax = plt.gca()
    v_2000 = vol_hist_ens.mean(axis=1).loc[2000]
    ax.plot(1850, dfo['Volume LIA'], 'o', c='grey', label='Volume Ref. LIA')
    ax.plot(2003, inv_volume_2003, 'x', c='k', label='Volume OGGM 2003')
    v_2019 = v_2000 + dvol
    ax.plot([2000, 2019], [v_2000, v_2019], '-', c='k', label='Obs. dV (Hugonnet)')
    ax.fill_between([2000, 2019], [v_2000, v_2019 - dvol_err], [v_2000, v_2019 + dvol_err],
                     alpha=0.2, color='k', label='Obs. dV uncertainty')
    ax.plot(2015, dfo['Volume 2015'], 'o', c='k', label='Volume Ref. 2015')

    plt.xlabel('Year'); plt.ylabel('Total volume (km$^3$)')
    plt.xlim([1300, vol_nat_ens.index.max()])
    plt.legend(loc='lower left', fontsize=9)
    plt.savefig(fname='talk/talk_6.png', bbox_inches='tight', dpi=150)

# %% [markdown]
# ## Same as talk_6.png, but the light-blue ensemble lines run to the end
# ## instead of being truncated at 1850.

# %%
with sns.axes_style("whitegrid"):
    plt.figure(figsize=(10, 5.5))

    for i, col in enumerate(vol_nat_ens.columns):
        plt.plot(vol_nat_ens.index, vol_nat_ens[col].values,
                 color='C0', alpha=0.15, linewidth=0.8,
                 label='Ensemble' if i == 0 else '_nolegend_')
    vol_nat_ens.mean(axis=1).plot(label='Nat Ensemble Mean', color='navy', linewidth=2)

    colors = plt.cm.Reds(np.linspace(0.45, 0.95, len(QUANTILE_ORDER)))
    for (qname, col), c in zip(QUANTILE_ORDER, colors):
        vol_hist_ens[col].plot(label=f'LMR seas. + BE "old" ({qname})', color=c, linewidth=2)

    ax = plt.gca()
    v_2000 = vol_hist_ens.mean(axis=1).loc[2000]
    ax.plot(1850, dfo['Volume LIA'], 'o', c='grey', label='Volume Ref. LIA')
    ax.plot(2003, inv_volume_2003, 'x', c='k', label='Volume OGGM 2003')
    v_2019 = v_2000 + dvol
    ax.plot([2000, 2019], [v_2000, v_2019], '-', c='k', label='Obs. dV (Hugonnet)')
    ax.fill_between([2000, 2019], [v_2000, v_2019 - dvol_err], [v_2000, v_2019 + dvol_err],
                     alpha=0.2, color='k', label='Obs. dV uncertainty')
    ax.plot(2015, dfo['Volume 2015'], 'o', c='k', label='Volume Ref. 2015')

    plt.xlabel('Year'); plt.ylabel('Total volume (km$^3$)')
    plt.xlim([1300, vol_nat_ens.index.max()])
    plt.legend(loc='lower left', fontsize=9)
    plt.savefig(fname='talk/talk_7.png', bbox_inches='tight', dpi=150)

# %%