# %% [markdown]
# # One HIST member against the NAT ensemble, for the talk
#
# As talk/lia_max_{volume,area}_<exp>.png (left panel), but with one HIST
# member only, the NAT ensemble drawn up to 1850 (where HIST and NAT start
# to differ), and only the modern references (no LIA reference). Also a
# NAT-only version, and one with the NAT runs to their end (2006-2021; the
# NAT mean is over the members available each year, which shifts it by
# < 1 km3 when some end).
#
# Run from round4/ensemble/ with the oggm_env python. Env vars: EXP,
# GDIR_EXP, MEMBER.

# %%
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')
GDIR_EXP = os.environ.get('GDIR_EXP', 'lam8_jjas')
MEMBER = os.environ.get('MEMBER', 'ens183')
HIST_MEMBERS = ['ens053', 'ens124', 'ens153', 'ens178', 'ens183']   # ens123 left out
NAT_END = 1850
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/{EXP}'
ids = pd.read_csv(f'{exp_dir}/glacier_statistics_reg11.csv', index_col=0).index
dfo = pd.read_csv("../../round1/filtered_update/Area_volume_change_LIA_Alps.csv",
                  encoding="ISO-8859-1", index_col=0, skiprows=[1]).loc[17]
st = pd.read_csv(V5_STATS, index_col=0, low_memory=False)
inv_volume_2003 = st.loc[ids, 'inv_volume_km3'].sum()
rgi_area = st.loc[ids, 'rgi_area_km2'].sum()
hug = utils.get_geodetic_mb_dataframe()
hug = hug.loc[hug.period == '2000-01-01_2020-01-01'].reindex(ids)
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


def ref_hugonnet(ax, s):
    v_2000 = s.loc[2000]
    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')


def figure(var, short, ylabel, refs):
    nat = pd.read_csv(f'{exp_dir}/{var}_{short}_nat_ens_reg11.csv', index_col=0)
    hist_all = pd.read_csv(f'{exp_dir}/{var}_{short}_hist_ens_reg11.csv',
                           index_col=0)[HIST_MEMBERS]
    hist = hist_all[MEMBER]
    # The 5 HIST members ranked and coloured by the 1850 percentile of their
    # NAT counterpart (as plot_lia_max.py). Always the volume, so that the
    # area figure has the same ranking and colours
    nat_1850 = pd.read_csv(f'{exp_dir}/volume_km3_nat_ens_reg11.csv', index_col=0).loc[NAT_END]
    pct = {m: (nat_1850 < nat_1850[m]).mean() * 100 for m in HIST_MEMBERS}
    order = sorted(HIST_MEMBERS, key=lambda m: pct[m])
    colors = dict(zip(order, plt.cm.Reds(np.linspace(0.45, 0.95, len(order)))))
    nat_all = nat
    ylim = None
    # Four versions on the same axes: HIST + obs with NAT to 1850, NAT only
    # (to 1850), HIST + obs with NAT to the end of each member, and the same
    # with the 5 HIST members (obs: Hugonnet starts from MEMBER's 2000 value)
    for nat_only, nat_full, all_hist in ((False, False, False), (True, False, False),
                                         (False, True, False), (False, True, True)):
        nat = nat_all if nat_full else nat_all.loc[:NAT_END]
        with sns.axes_style("whitegrid"):
            fig, ax = plt.subplots(figsize=(10, 5.5))
            for i, c in enumerate(nat.columns):
                ax.plot(nat.index, nat[c].values, color='C0', alpha=0.15, linewidth=0.8,
                        label='NAT ensemble' if i == 0 else '_nolegend_')
            ax.plot(nat.index, nat.mean(axis=1), color='navy', linewidth=3,
                    label='NAT ensemble mean')
            if all_hist:
                for m in order:
                    ax.plot(hist_all.index, hist_all[m].values, color=colors[m],
                            linewidth=1.8, label=f'HIST {m} (NAT 1850: P{pct[m]:.0f})')
            elif not nat_only:
                ax.plot(hist.index, hist.values, color='C3', linewidth=2,
                        label=f'HIST {MEMBER}')
            if not nat_only:
                for r in refs:
                    r(ax, hist)
            ax.set_xlim([1300, hist.index.max()])
            if ylim is None:
                ylim = ax.get_ylim()
            ax.set_ylim(ylim)
            ax.set_xlabel('Year')
            ax.set_ylabel(ylabel)
            ax.legend(loc='lower left', fontsize=9 if all_hist else 10,
                      ncol=2 if all_hist else 1)
            name = (f'{var}_{EXP}_nat_only' if nat_only else
                    f'{var}_{EXP}_hist5_natfull' if all_hist else
                    f'{var}_{EXP}_{MEMBER}_natfull' if nat_full else f'{var}_{EXP}_{MEMBER}')
            fig.savefig(f'talk/{name}.png', bbox_inches='tight', dpi=150)
            plt.close(fig)


def ref_point(x, y, **kw):
    return lambda ax, s: ax.plot(x, y, **kw)


# %%
os.makedirs('talk', exist_ok=True)
figure('volume', 'km3', 'Total volume (km$^3$)',
       [ref_point(2003, inv_volume_2003, marker='x', c='k', ls='', ms=8,
                  label='Volume OGGM 2003'),
        ref_hugonnet,
        ref_point(2015, dfo['Volume 2015'], marker='o', c='k', ls='', ms=8,
                  label='Volume Ref. 2015')])
figure('area', 'km2', 'Total area (km$^2$)',
       [ref_point(2003, rgi_area, marker='x', c='k', ls='', ms=8,
                  label='Area RGI (~2003)'),
        ref_point(2003, dfo['Area 2003'], marker='+', c='k', ls='', ms=10,
                  label='Area Reinthaler 2003'),
        ref_point(2015, dfo['Area 2015'], marker='o', c='k', ls='', ms=8,
                  label='Area Reinthaler 2015')])