"""Round 7 version of round6/talk/plot_lia_max.py: the LIA comparison of the
round 7 ensemble (default ens_v8_adapbs_a5v10: AdaPBS round 2 gdirs, LAPrec
precipitation, runs from 1500), with 5 HIST members: the min / P25 / median /
P75 / max of the 1850 regional volume across the 197 NAT members (HIST == NAT
before 1850), labelled by their 1850 volume percentile (INCLUDE_010=1 adds
ens010, the calibration member). All 192 HIST and 197 NAT members are run.
Shown: all HIST members (faint, grey) and their LIA box; the NAT runs are not shown
(user, 2026-11-02: after 1850 they are a climate no comparison here uses).

Left: regional totals on the common subset (filtered_ens/<NAME>/, raw), with
the present-day observations on the same glaciers (RGI 6 area 2003, Paul et
al. 2020 / RGI 7 2015 area as their -14 ± 2.9% change applied to it, the
inversion volume 2003 and Hugonnet's dV 2000-2020 through it).
Right: the sum of the per-glacier 1800-1875 maxima (area_min_h; volume in the
year of the max area), corrected for the outlines the model does not cover as
in round6/lia_validation/validate.py ("additive": + the not-modelled outlines'
observed area / volume, - the modelled ice on 'remnant' glaciers, + validate's
filled round5 ens010 values for the glaciers outside the subset), against the
outline totals (Reinthaler & Paul, regional 1 sigma). Same correction for
every member.

    [NAME=ens_v8_adapbs] python plot_lia_max.py
    -> figs/lia_max_{volume,area}_<NAME>.png (+ figs/lia_max_corrected_<NAME>.csv)
"""
import os
import sys
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, f'{ROOT}/round6/lia_validation')
import validate as V  # noqa: E402

NAME = os.environ.get('NAME', 'ens_v8_adapbs_a5v10')
FDIR = f'{HERE}/filtered_ens/{NAME}'
DDIR = f'{HERE}/output_ens/{NAME}'
OUT = f'{HERE}/figs'
os.makedirs(OUT, exist_ok=True)
WINDOW = (1800, 1875)
X0 = int(os.environ.get('X0', 1650))   # runs start in 1500: hide the spin-up
y0, y1 = WINDOW
T_RGI, T_2015 = 2004, 2016          # OGGM's time: start of 2004 / of 2016
PAUL_DA, PAUL_ERR = 100 * (1806 / 2100 - 1), 100 * 60 / 2100


def load(var):
    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


ids = pd.read_csv(f'{FDIR}/glacier_statistics_reg11.csv', index_col=0).index

# --- the 6 HIST members ------------------------------------------------------
vn, vh = load('volume')
v50 = vn.loc[1850]
QS = {0: 'min', 0.25: 'P25', 0.5: 'median', 0.75: 'P75', 1: 'max'}
members = []
for q in QS:
    m = (v50 - v50.quantile(q)).abs().idxmin()
    assert m in vh.columns, m
    members.append(m)
# 5 members (user, 2026-11-02); INCLUDE_010=1 adds the calibration member
if os.environ.get('INCLUDE_010') == '1' and 'ens010' not in members:
    members.append('ens010')
pct = {m: 100 * (v50 < vh.loc[1850, m]).mean() for m in members}
order = sorted(members, key=lambda m: vh.loc[1850, m])
colors = dict(zip(order, plt.cm.Reds(np.linspace(0.4, 0.95, len(order)))))
print('1850 volume percentiles: ' + ', '.join(
    f'{m} {vh.loc[1850, m]:.1f} km3 P{pct[m]:.0f}' for m in order))


def label(m):
    return f'HIST {m} (1850: P{pct[m]:.0f})' + (' calibration member' if m == 'ens010' else '')


# --- corrected LIA sums ----------------------------------------------------
s, links, without, st, M, filled, yamax = V.load()
g, main, main_status, nb = V.groups(s, links)
R, T, info = V.regional(s, without, st, M, main_status)
rem = main_status.index[main_status == 'remnant'].intersection(st.index)
rem_in = rem.intersection(ids)
outside = st.index.difference(ids)
F_out = {v: float(M[v].loc['ens010', outside].sum()) for v in ['area', 'vol']}
R_out = {v: float(M[v].loc['ens010', outside.intersection(rem)].sum()) for v in ['area', 'vol']}
nm = {'area': info['nm_a'], 'vol': info['nm_v']}
print(f'{len(ids)} glaciers in the subset, {len(outside)} outside (filled: {F_out["area"]:.1f} km2 / '
      f'{F_out["vol"]:.2f} km3), {len(rem_in)} remnant glaciers in the subset; not-modelled '
      f'outlines +{nm["area"]:.1f} km2 / +{nm["vol"]:.2f} km3')


def remnant_max(f):
    """Sum of the per-glacier 1800-1875 max area (and volume then) on the remnant glaciers."""
    with xr.open_dataset(f) as ds:
        d = ds[['volume', 'area_min_h']].sel(rgi_id=list(rem_in), time=slice(y0, y1)).load()
    j = d.area_min_h.fillna(-1).argmax('time')
    return (os.path.basename(f), float(d.area_min_h.max('time').sum()) * 1e-6,
            float(d.volume.isel(time=j).sum()) * 1e-9)


cache = f'{OUT}/lia_max_remnants_{NAME}.csv'
if not os.path.exists(cache):
    fs = sorted(f'{DDIR}/{x}' for x in os.listdir(DDIR) if x.endswith('_reg11.nc')
                and '_ens' in x)
    with ProcessPoolExecutor(16) as ex:
        res = list(ex.map(remnant_max, fs))
    pd.DataFrame(res, columns=['file', 'rem_area', 'rem_vol']).to_csv(cache, index=False)
rm = pd.read_csv(cache)
rm['kind'] = rm.file.str.extract(r'lmr_mira_(hist|nat)_')[0]
rm['member'] = rm.file.str.extract(r'_(ens\d+)_')[0]
rm = rm.set_index(rm.kind + '_' + rm.member)

lm = pd.read_csv(f'{FDIR}/lia_max_members.csv', index_col=0)
lm = lm.join(rm[['rem_area', 'rem_vol']])
assert lm.rem_area.notna().all()
lm['area_corr'] = (lm[f'area_max_{y0}_{y1}'] + F_out['area'] - lm.rem_area - R_out['area']
                   + nm['area'])
lm['vol_corr'] = (lm[f'vol_at_amax_{y0}_{y1}'] + F_out['vol'] - lm.rem_vol - R_out['vol']
                  + nm['vol'])
lm.to_csv(f'{OUT}/lia_max_corrected_{NAME}.csv')
nat_lm = lm.loc[lm.kind == 'nat']
hist_lm = lm.loc[lm.kind == 'hist'].set_index('member')

# --- present-day observations on the subset (as round5/talk_common.py) -----
from oggm import cfg, utils  # noqa: E402
cfg.initialize(logging_level='CRITICAL')
st8 = pd.read_csv(f'{ROOT}/lmr_gdirs_v8_202610/exps/pf3.5_mf5_lam12/RGI62/b_240/L3/summary/'
                  'glacier_statistics_11.csv', index_col=0, low_memory=False)
inv_vol = float(st8.loc[ids, 'inv_volume_km3'].sum())
rgi_area = float(st8.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
_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]
dvol_err = abs(dvol) * float(_reg.err_dmdtda / abs(_reg.dmdtda))


def refs_area(ax, hist):
    ax.plot(T_RGI, rgi_area, 'o', c='k', ms=6, label='RGI 6 area (2003)')
    ax.errorbar([T_2015], [rgi_area * (1 + PAUL_DA / 100)], yerr=[rgi_area * PAUL_ERR / 100],
                fmt='s', mfc='white', mec='k', ecolor='k', ms=6, capsize=3,
                label='2015 (Paul et al. 2020 / RGI 7: −14 ± 2.9%)')


def refs_volume(ax, hist):
    """Inversion volume at the RGI date; Hugonnet's dV 2000-2020 from the mean 2000 volume of
    the highlighted HIST members (they differ by < 5 km3 in 2000). Uncertainty: the regional
    one of Hugonnet et al. (2021, OGGM's hugonnet_2021_regional_avg.csv, RGI 11, 2000-2020) as
    a relative error on the summed dV, growing from 0 in 2000."""
    ax.plot(T_RGI, inv_vol, 'o', c='k', ms=6, label='inversion volume (2003)')
    v0 = float(hist.loc[2000].mean())
    x = np.array([2000, 2020])
    y = v0 + dvol * (x - 2000) / 20
    e = dvol_err * (x - 2000) / 20
    ax.fill_between(x, y - e, y + e, color='k', alpha=0.25, lw=0, zorder=5)
    ax.plot(x, y, ls=(0, (4, 2)), c='k', lw=1.8, zorder=6,
            label='Hugonnet dV 2000–2020 (± regional unc.), from HIST 2000')


def lia_figure(var, col, ref, ref_err, ylabel, refs, fname):
    series_nat, series_hist = load(var)
    with sns.axes_style('whitegrid'):
        fig, (ax, axr) = plt.subplots(
            1, 2, figsize=(12, 5.5), sharey=True,
            gridspec_kw={'width_ratios': [5, 1.5], 'wspace': 0.18})
        series_nat, series_hist = series_nat.loc[X0:], series_hist.loc[X0:]   # y range too
        ax.axvspan(y0, y1, color='0.85', alpha=0.5, lw=0,
                   label=f'Window of the per-glacier max ({y0}–{y1})')
        # all HIST members (== their NAT twins before 1850; after 1850 the NAT runs are a
        # climate that no comparison here uses, so they are not shown)
        for i, c in enumerate(series_hist.columns):
            ax.plot(series_hist.index, series_hist[c].values, color='0.55', alpha=0.15, lw=0.8,
                    label=f'HIST, all {series_hist.shape[1]} members' if i == 0 else '_nolegend_')
        ax.plot(series_hist.index, series_hist.mean(axis=1), color='0.2', lw=2,
                label='HIST ensemble mean')
        for m in order:
            ax.plot(series_hist.index, series_hist[m], color=colors[m],
                    lw=2.6 if m == 'ens010' else 1.8, label=label(m))
        refs(ax, series_hist[order])
        ax.set_xlim([X0, 2025])
        ax.set_xlabel('Year')
        ax.set_ylabel(ylabel + ', model glaciers, not corrected')
        ax.legend(loc='upper left', fontsize=7.5, ncol=2)

        v = hist_lm[col].astype(float).values      # all HIST members
        axr.boxplot(v, positions=[0], widths=0.5, showfliers=False,
                    medianprops={'color': '0.2', 'linewidth': 2},
                    boxprops={'color': '0.45'}, whiskerprops={'color': '0.45'},
                    capprops={'color': '0.45'})
        jit = np.random.default_rng(0).uniform(-0.15, 0.15, len(v))
        axr.scatter(jit, v, s=8, color='0.5', alpha=0.4, lw=0)
        for m, x in zip(order, np.linspace(0.7, 1.3, len(order))):
            axr.scatter([x], [float(hist_lm.loc[m, col])], s=60, color=colors[m],
                        edgecolor='white', linewidth=1, zorder=3)
        axr.errorbar([2], [ref], 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([f'all HIST\n({len(v)})', '5 HIST', 'outlines'], fontsize=9)
        axr.set_xlim(-0.6, 2.6)
        axr.set_title(f'Sum of per-glacier max, {y0}–{y1},\n+ missing outlines', fontsize=10)
        fig.suptitle(f'Alps-wide: {len(ids)} glaciers, {series_hist.shape[1]} HIST members, '
                     f'{len(order)} highlighted (round 7, {NAME})', y=0.98)
        fig.savefig(fname, bbox_inches='tight', dpi=150)
        plt.close(fig)
    h = hist_lm.loc[order, col].astype(float)
    print(f'{col}: all-HIST median {np.median(v):.0f} [P5–P95 {np.percentile(v, 5):.0f}–'
          f'{np.percentile(v, 95):.0f}], HIST ' + ', '.join(f'{m} {h[m]:.0f}' for m in order)
          + f'; outlines {ref:.0f} ± {ref_err:.0f}  -> {fname}')


lia_figure('volume', 'vol_corr', info['obs_v'], V.VOL_ERR * info['obs_v'], 'Total volume (km$^3$)',
           refs_volume, f'{OUT}/lia_max_volume_{NAME}.png')
lia_figure('area', 'area_corr', info['obs_a'], V.AREA_ERR * info['obs_a'], 'Total area (km$^2$)',
           refs_area, f'{OUT}/lia_max_area_{NAME}.png')