# ---
# 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
# ---

# %% [markdown]
# # round3's talk_7.png, for a round4 ensemble dataset
#
# Same figure as round3/plot_talk_hist_ens.py (talk_7.png): the NAT ensemble
# over the full period, with the HIST realisations on top. Differences:
# - ens123 is left out (it was only the previous random member);
# - the HIST members are ranked by the 1850 volume of their NAT counterpart
#   in *this* dataset (round3's ranking came from the old v4 NAT runs);
# - Hugonnet is summed per glacier over the plotted glaciers, instead of the
#   regional value for all of region 11;
# - "Volume OGGM 2003" is the inversion volume of the v5 gdirs used here.
#
# Run from round4/ensemble/ (after `python ../preprocess_mira.py <exp>`).

# %%
import os

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from oggm import cfg, utils

cfg.initialize(logging_level='CRITICAL')

EXP = os.environ.get('EXP', 'lam8_jjas_dyn')         # dataset in filtered/ (and output/)
GDIR_EXP = os.environ.get('GDIR_EXP', 'lam8_jjas')   # its v5 gdirs, for the inversion volume
HIST_MEMBERS = ['ens053', 'ens124', 'ens153', 'ens178', 'ens183']
FILTERED_DIR = 'filtered'
V5_STATS = ('/home/www/fmaussion/runs/lmr_2025/lmr_gdirs_v5_202610/exps/'
            f'{GDIR_EXP}/RGI62/b_160/L3/summary/glacier_statistics_11.csv')

# %%
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)[HIST_MEMBERS]
# The glaciers in the CSVs (the common running subset)
ids = pd.read_csv(f'{exp_dir}/glacier_statistics_reg11.csv', index_col=0).index
st = pd.read_csv(V5_STATS, index_col=0, low_memory=False)
inv_volume_2003 = st.loc[ids, '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]

# Hugonnet 2000-2020, per glacier, summed over the plotted glaciers
hug = utils.get_geodetic_mb_dataframe()
hug = hug.loc[hug.period == '2000-01-01_2020-01-01'].reindex(ids)
# m w.e. yr-1 x m2 x 20 yr -> m3 w.e. -> km3 ice (900 kg m-3)
dvol = float((hug.dmdtda * hug.area).sum()) * 20 * 1000 / 900 * 1e-9
dvol_err = float(np.sqrt(((hug.err_dmdtda * hug.area) ** 2).sum())) * 20 * 1000 / 900 * 1e-9

# %%
# Rank the HIST members by the 1850 volume of their NAT counterpart (the
# same member before 1850), as a percentile of the whole NAT ensemble
nat_1850 = vol_nat_ens.loc[1850]
pct = {m: (nat_1850 < nat_1850[m]).mean() * 100 for m in HIST_MEMBERS}
order = sorted(HIST_MEMBERS, key=lambda m: pct[m])
print('NAT 1850 volume percentile of the HIST members:',
      {m: round(pct[m]) for m in order})

# %%
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='NAT 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(order)))
    for m, c in zip(order, colors):
        vol_hist_ens[m].plot(label=f'HIST {m} (NAT 1850: P{pct[m]:.0f})',
                             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.title(f'{EXP}: {len(ids)} glaciers, {vol_nat_ens.shape[1]} NAT members')
    plt.legend(loc='lower left', fontsize=9)
    os.makedirs('talk', exist_ok=True)
    plt.savefig(fname=f'talk/talk_7_{EXP}.png', bbox_inches='tight', dpi=150)