# ---
# 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]
# # What we actually compare to the LIA reference
#
# The Reinthaler & Paul (2025) LIA outlines are each glacier's *last*
# maximum, reached at different times. So the model value to compare is the
# sum over glaciers of each glacier's own maximum over a window, not the
# regional total in one year. Two figures (volume, area), each with:
# - left: the regional time series (as talk_7), the window shaded;
# - right: the sum of per-glacier maxima over the window, for the 197 NAT
#   members (box + points), the HIST members, and the reference.
#
# Run from round4/ensemble/ (after `python ../preprocess_mira.py <exp>`).
# The per-glacier maxima need all the run files: they are cached in
# filtered/<exp>/lia_max_members.csv (delete it to recompute).

# %%
import glob
import os
import re

import numpy as np
import pandas as pd
import xarray as xr
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
WINDOW = (1800, 1875)
WINDOWS = [(1800, 1875), (1800, 1900), (1800, 1975)]   # all cached; WINDOW is plotted
HIST_MEMBERS = ['ens053', 'ens124', 'ens153', 'ens178', 'ens183']
FILTERED_DIR = 'filtered'
OUTPUT_DIR = 'output'
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}'
ids = pd.read_csv(f'{exp_dir}/glacier_statistics_reg11.csv', index_col=0).index

# %% [markdown]
# ## Per-glacier maxima, per member (cached)

# %%
cache = f'{exp_dir}/lia_max_members.csv'
if os.path.exists(cache):
    lm = pd.read_csv(cache, index_col=0)
else:
    rows = {}
    files = sorted(glob.glob(f'{OUTPUT_DIR}/{EXP}/lmr_mira_*_ens*_reg11.nc'))
    for i, f in enumerate(files):
        kind, num = re.search(r'lmr_mira_(\w+?)_(ens\d+)_reg11', f).groups()
        with xr.open_dataset(f) as ds:
            ds = ds[['volume', 'area_min_h']].sel(rgi_id=ids).load()
        r = {'kind': kind, 'member': num}
        for y0, y1 in WINDOWS:
            w = ds.sel(time=slice(y0, y1))
            r[f'area_max_{y0}_{y1}'] = float(w.area_min_h.max('time').sum()) * 1e-6
            r[f'vol_max_{y0}_{y1}'] = float(w.volume.max('time').sum()) * 1e-9
            iy = w.area_min_h.argmax('time')
            r[f'vol_at_amax_{y0}_{y1}'] = float(w.volume.isel(time=iy).sum()) * 1e-9
        rows[f'{kind}_{num}'] = r
        if (i + 1) % 25 == 0:
            print(f'{i + 1}/{len(files)}')
    lm = pd.DataFrame(rows).T
    lm.to_csv(cache)
for c in lm.columns:
    if c not in ('kind', 'member'):
        lm[c] = lm[c].astype(float)

# %%
y0, y1 = WINDOW
nat = lm.loc[lm.kind == 'nat']
hist = lm.loc[lm.kind == 'hist'].set_index('member').loc[HIST_MEMBERS]
print(f'Sum of per-glacier maxima over {y0}-{y1} ({len(ids)} glaciers):')
for var, ref in [('area_max', 4211.12), ('vol_max', 280.95)]:
    col = f'{var}_{y0}_{y1}'
    print(f'  {var}: NAT median {nat[col].median():.0f} '
          f'[{nat[col].quantile(0.05):.0f}-{nat[col].quantile(0.95):.0f}], '
          f'HIST {hist[col].min():.0f}-{hist[col].max():.0f}, reference {ref:.0f}')

# %% [markdown]
# ## Time series and references (as talk_7)

# %%
vol_nat = pd.read_csv(f'{exp_dir}/volume_km3_nat_ens_reg11.csv', index_col=0)
vol_hist = pd.read_csv(f'{exp_dir}/volume_km3_hist_ens_reg11.csv', index_col=0)[HIST_MEMBERS]
area_nat = pd.read_csv(f'{exp_dir}/area_km2_nat_ens_reg11.csv', index_col=0)
area_hist = pd.read_csv(f'{exp_dir}/area_km2_hist_ens_reg11.csv', index_col=0)[HIST_MEMBERS]

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

# Reinthaler & Paul (2025) uncertainties (relative, 1 sigma): area 5.1%
# (digitisation experiments; abstract: 4244 +- 214 km2), volume 15.3% (area,
# surface and bed uncertainties combined; abstract: 280 +- 43 km3)
REF_AREA_ERR = 0.051
REF_VOL_ERR = 0.153

# Members ranked (and coloured) by the 1850 volume of their NAT counterpart
nat_1850 = vol_nat.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])
colors = dict(zip(order, plt.cm.Reds(np.linspace(0.45, 0.95, len(order)))))


# %%
def lia_figure(series_nat, series_hist, var, ref_lia, ref_err, ylabel, refs, fname):
    """Left: time series. Right: sum of per-glacier maxima over WINDOW."""
    col = f'{var}_{y0}_{y1}'
    with sns.axes_style("whitegrid"):
        fig, (ax, axr) = plt.subplots(
            1, 2, figsize=(12, 5.5), sharey=True,
            gridspec_kw={'width_ratios': [5, 1.3], 'wspace': 0.04})

        # Left: the time series
        ax.axvspan(y0, y1, color='0.85', alpha=0.5, lw=0,
                   label=f'Window of the per-glacier max ({y0}-{y1})')
        for i, c in enumerate(series_nat.columns):
            ax.plot(series_nat.index, series_nat[c].values, color='C0',
                    alpha=0.15, linewidth=0.8,
                    label='NAT ensemble' if i == 0 else '_nolegend_')
        ax.plot(series_nat.index, series_nat.mean(axis=1), color='navy',
                linewidth=2, label='NAT ensemble mean')
        for m in order:
            ax.plot(series_hist.index, series_hist[m], color=colors[m],
                    linewidth=2, label=f'HIST {m} (NAT 1850: P{pct[m]:.0f})')
        for r in refs:
            r(ax, series_hist)
        ax.set_xlim([1300, series_nat.index.max()])
        # Fixed y limits, to compare datasets: YLIM_VOL / YLIM_AREA="lo,hi"
        ylim = os.environ.get('YLIM_VOL' if var.startswith('vol') else 'YLIM_AREA')
        if ylim:
            ax.set_ylim([float(x) for x in ylim.split(',')])
        ax.set_xlabel('Year')
        ax.set_ylabel(ylabel)
        ax.legend(loc='lower left', fontsize=8.5)

        # Right: what is compared with the reference
        v = nat[col].values
        axr.boxplot(v, positions=[0], widths=0.5, showfliers=False,
                    medianprops={'color': 'navy', 'linewidth': 2},
                    boxprops={'color': 'C0'}, whiskerprops={'color': 'C0'},
                    capprops={'color': 'C0'})
        jit = np.random.default_rng(0).uniform(-0.15, 0.15, len(v))
        axr.scatter(jit, v, s=8, color='C0', alpha=0.4, lw=0)
        for m, x in zip(order, np.linspace(0.75, 1.25, len(order))):
            axr.scatter([x], [hist.loc[m, col]], s=60, color=colors[m],
                        edgecolor='white', linewidth=1, zorder=3)
        axr.errorbar([2], [ref_lia], yerr=[ref_err], fmt='D', ms=8, color='k',
                     capsize=5, elinewidth=1.5, zorder=3)
        axr.set_xticks([0, 1, 2])
        axr.set_xticklabels(['NAT', 'HIST', 'Ref.'])
        axr.set_xlim(-0.6, 2.6)
        axr.set_title(f'Sum of per-glacier\nmax, {y0}-{y1}', fontsize=10)
        fig.suptitle(f'{EXP}: {len(ids)} glaciers, {series_nat.shape[1]} NAT '
                     f'members, {len(order)} HIST members', y=0.98)
        os.makedirs('talk', exist_ok=True)
        fig.savefig(fname, bbox_inches='tight', dpi=150)
        return fig


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


def ref_hugonnet(ax, s):
    v_2000 = s.mean(axis=1).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')


# %%
lia_figure(vol_nat, vol_hist, 'vol_max', dfo['Volume LIA'],
           REF_VOL_ERR * dfo['Volume LIA'],
           'Total volume (km$^3$)',
           [ref_point(2003, inv_volume_2003, marker='x', c='k', ls='',
                      label='Volume OGGM 2003'),
            ref_hugonnet,
            ref_point(2015, dfo['Volume 2015'], marker='o', c='k', ls='',
                      label='Volume Ref. 2015')],
           f'talk/lia_max_volume_{EXP}.png')

# %%
lia_figure(area_nat, area_hist, 'area_max', dfo['Area LIA'],
           REF_AREA_ERR * dfo['Area LIA'],
           'Total area (km$^2$)',
           [ref_point(2003, rgi_area, marker='x', c='k', ls='',
                      label='Area RGI (~2003)'),
            ref_point(2003, dfo['Area 2003'], marker='+', c='k', ls='', ms=9,
                      label='Area Reinthaler 2003'),
            ref_point(2015, dfo['Area 2015'], marker='o', c='k', ls='',
                      label='Area Reinthaler 2015')],
           f'talk/lia_max_area_{EXP}.png')