"""Shared inputs of the round 5 talk figures (plot_talk*.py, plot_lia_max.py).

Env vars: NAME (dataset in filtered_ens/, default ens_pf3.5_mf5_lam12),
GDIR_EXP (its v7 gdirs, default sw_pf3.5_mf5_lam12), BORDER (240).

- HIST_MEMBERS: the 5 quantile members of the 1820-1850 mean regional
  volume (chosen 2026-10-19 from the round4 NAT runs), ranked and coloured
  by their percentile of the 1820-1850 mean volume in *this* NAT ensemble
  (HIST == NAT before 1850, so it is the member's own baseline).
- Volume references (comparison rules of 2026-10-23, see CLAUDE.md): the
  RGI-date (2003) volume is the inversion volume of these gdirs; the LIA
  reference is anchored on it: V_inv(2003) + 180 km3 (Reinthaler & Paul
  2025, dV LIA -> 2015) - |observed dV 2003 -> 2015| (the Hugonnet rate x
  12 years), with Reinthaler's ±39 km3. The 2015 Reinthaler volume is not
  shown (it depends on their bed).
- Area references: RGI (~2003), Reinthaler 2003 / 2015, and the LIA area
  4211 ± 5.1%.
"""
import os

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

cfg.initialize(logging_level='CRITICAL')

HERE = os.path.dirname(os.path.abspath(__file__))
NAME = os.environ.get('NAME', 'ens_pf3.5_mf5_lam12')
GDIR_EXP = os.environ.get('GDIR_EXP', 'sw_pf3.5_mf5_lam12')
BORDER = int(os.environ.get('BORDER', 240))
HIST_MEMBERS = ['ens066', 'ens008', 'ens010', 'ens055', 'ens041']
BASE = (1820, 1850)
FDIR = f'{HERE}/filtered_ens/{NAME}'
TALK = f'{HERE}/talk'
os.makedirs(TALK, exist_ok=True)

ids = pd.read_csv(f'{FDIR}/glacier_statistics_reg11.csv', index_col=0).index
st = pd.read_csv(f'{HERE}/../lmr_gdirs_v7_202610/exps/{GDIR_EXP}/RGI62/b_{BORDER:03d}/L3/'
                 'summary/glacier_statistics_11.csv', 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()

dfo = pd.read_csv(f'{HERE}/../round1/filtered_update/Area_volume_change_LIA_Alps.csv',
                  encoding='ISO-8859-1', index_col=0, skiprows=[1]).loc[17]
REF_DV_LIA_2015, REF_DV_ERR = 180., 39.     # Reinthaler & Paul 2025, km3
REF_AREA_ERR = 0.051                        # relative, 1 sigma

# 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)
dvol = float((hug.dmdtda * hug.area).sum()) * 20 * 1000 / 900 * 1e-9      # km3 ice, < 0
# Uncertainty: the regional one of Hugonnet et al. (2021) for RGI 11, 2000-2020
# (OGGM's regional table), as a relative error on the summed dV. The paper's
# regional errors account for the spatial correlation of the per-glacier
# errors (neither independent nor fully correlated).
_reg = pd.read_csv(utils.file_downloader('https://cluster.klima.uni-bremen.de/~oggm/'
                                         'geodetic_ref_mb/hugonnet_2021_regional_avg.csv'))
_reg = _reg.loc[(_reg.reg == 11) & (_reg.period == '2000-01-01_2020-01-01')].iloc[0]
HUG_REL_ERR = float(_reg.err_dmdtda / abs(_reg.dmdtda))
dvol_err = abs(dvol) * HUG_REL_ERR
# Anchored LIA volume reference
ref_lia_volume = inv_volume_2003 + REF_DV_LIA_2015 + dvol * 12 / 20


def load(var):
    """var: 'volume' or 'area' -> (nat, hist) DataFrames (years x members)."""
    short = 'km3' if var == 'volume' else 'km2'
    nat = pd.read_csv(f'{FDIR}/{var}_{short}_nat_ens_reg11.csv', index_col=0)
    hist = pd.read_csv(f'{FDIR}/{var}_{short}_hist_ens_reg11.csv', index_col=0)
    nat.index = nat.index.astype(int)
    hist.index = hist.index.astype(int)
    return nat, hist


# Ranking and colours of the 5 HIST members (always from the volume)
_vn, _vh = load('volume')
base_nat = _vn.loc[BASE[0]:BASE[1]].mean()
base_hist = _vh.loc[BASE[0]:BASE[1]].mean()
pct = {m: 100 * (base_nat < base_hist[m]).mean() 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)))))


def label(m):
    return f'HIST {m} ({BASE[0]}–{BASE[1]}: P{pct[m]:.0f})'


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


def ref_hugonnet(ax, s):
    """Observed dV 2000-2019 from the HIST value(s) in 2000 (mean if several)."""
    v_2000 = (s.mean(axis=1) if s.ndim == 2 else 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 ref_lia_vol(ax, s):
    ax.errorbar([1850], [ref_lia_volume], yerr=[REF_DV_ERR], fmt='o', c='grey',
                capsize=4, label='Volume Ref. LIA (anchored on 2003)')


VOLUME_REFS = [ref_point(2003, inv_volume_2003, marker='x', c='k', ls='', ms=8,
                         label='Volume OGGM 2003 (inversion)'),
               ref_hugonnet]
AREA_REFS = [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')]

print(f'{NAME}: {len(ids)} glaciers; inversion volume 2003 {inv_volume_2003:.1f} km3, '
      f'observed dV 2000-2020 {dvol:.1f} ± {dvol_err:.1f} km3 (regional rel. err. '
      f'{100 * HUG_REL_ERR:.1f}%), anchored LIA volume ref '
      f'{ref_lia_volume:.1f} ± {REF_DV_ERR:.0f} km3; 1820-1850 percentiles: '
      + ', '.join(f'{m} P{pct[m]:.0f}' for m in order))