"""LIA validation of the round 5 HIST ensemble (ens_pf3.5_mf5_lam12) against the
Reinthaler & Paul LIA outlines, regionally and per outline.

Inputs:
- lia_model_hist.nc (extract_lia.py): per member and glacier, the LIA state
  (max area over WINDOW, volume in the year of that max), 1850, 2003, 2015;
- data/inventories/lia_outlines_summary.csv, lia_rgi6_links.csv,
  rgi6_without_lia.csv (the LIA outlines and their RGI6 links).

Choices (see README.md):
- LIA window 1800-1875 (project default; WINDOW=1820_1860 for the check),
  sum of per-glacier maxima, volume at the year of the area max.
- Model glaciers without output in a member (~104 of 3892, 6 km2 RGI area)
  are filled with their median over the members that have output; the 102
  that never have output with the median LIA / RGI ratio of their size
  class (filled_* columns say which).
- Reference run ens010 (the calibration member, P54 of the 1820-1850 NAT
  volumes), plus ens178 (the attribution default state, P99) and the 192
  HIST members for the spread (HIST == NAT before 1850).
- Per outline: groups = the LIA outline holding most of each RGI glacier
  (is_main_lia). 'shared' outlines (no RGI glacier of their own) are merged
  into the neighbour that holds the RGI ice covering them. Glaciers whose
  main outline is a 'remnant' are left out of the per-outline comparison and
  removed from the regional total when the not-modelled outlines are added
  back (otherwise they are counted twice). Alternative: each glacier spread
  over its outlines by frac_of_rgi (FRAC columns).

    python validate.py        # ~30 s
-> lia_regional_<WINDOW>.csv, lia_per_outline_<WINDOW>.csv,
   lia_variance_<WINDOW>.csv, lia_worst20_<WINDOW>.csv, figs/lia_*.png
"""
import os

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(f'{HERE}/../..')
INV = f'{ROOT}/data/inventories'
STATS = (f'{ROOT}/lmr_gdirs_v7_202610/exps/sw_pf3.5_mf5_lam12/RGI62/b_240/L3/'
         'summary/glacier_statistics_11.csv')
RGI_ATTRS = '/home/users/fmaussion/run_data/rgi/RGIV60/00_rgi60_attribs/11_rgi60_CentralEurope.csv'
WIN = os.environ.get('WINDOW', '1800_1875')
REF, ALT = 'ens010', 'ens178'
NOT_MODELLED = ['no_rgi', 'remnant', 'sliver_only']
AREA_ERR, VOL_ERR = 0.051, 0.153            # Reinthaler & Paul, 1 sigma, relative
RT_DV, RT_DV_ERR, RT_A2015 = 180., 39., 1806.  # LIA -> 2015 volume change, 2015 area
FIGS = f'{HERE}/figs'
os.makedirs(FIGS, exist_ok=True)

# Colours (dataviz reference palette)
INK, INK2, GRID = '#0b0b0b', '#52514e', '#e1e0d9'
C_REF, C_ALT, C_ENS, C_OBS = '#2a78d6', '#eb6834', '#86b6ef', '#0b0b0b'
DIV = LinearSegmentedColormap.from_list('div', ['#a12f2f', '#e34948', '#f0efec', '#3987e5', '#184f95'])
plt.rcParams.update({'axes.edgecolor': INK2, 'axes.labelcolor': INK, 'xtick.color': INK2,
                     'ytick.color': INK2, 'axes.grid': True, 'grid.color': GRID,
                     'grid.linewidth': 0.6, 'axes.axisbelow': True, 'font.size': 9,
                     'axes.spines.top': False, 'axes.spines.right': False})

SIZE_BINS = [0, 0.1, 0.5, 1, 5, 20, np.inf]
SIZE_LBL = ['<0.1', '0.1-0.5', '0.5-1', '1-5', '5-20', '>20']
CIT_SHORT = {'Maisch et al. 2000': 'Maisch (CH)', 'Fischer et al. 2015': 'Fischer (AT)',
             'Gardent 2014': 'Gardent (FR)', 'Knoll et al. 2009': 'Knoll (S. Tyrol)',
             'GlaRiskAlp': 'GlaRiskAlp (IT/FR)', 'Zanoner et al. 2017': 'Zanoner (Trentino)',
             'Reinthaler & Paul 2024': 'Reinthaler (new)', 'Lucchesi et al. 2014': 'Lucchesi (Aosta)',
             'Colucci & Žebre 2016': 'Colucci (Julian)', 'Scotti & Brardinoni 2018': 'Scotti (Lomb.)'}


# --------------------------------------------------------------------------
# Inputs
# --------------------------------------------------------------------------
def load():
    s = pd.read_csv(f'{INV}/lia_outlines_summary.csv', index_col=0)
    links = pd.read_csv(f'{INV}/lia_rgi6_links.csv')
    without = pd.read_csv(f'{INV}/rgi6_without_lia.csv')
    st = pd.read_csv(STATS, index_col=0, low_memory=False)
    st = st.loc[st.rgi_subregion == '11-01']
    ds = xr.open_dataset(f'{HERE}/lia_model_hist.nc').sel(rgi_id=st.index.values).load()
    rgi_a = st.rgi_area_km2
    vars_ = [f'area_max_{WIN}', f'vol_at_amax_{WIN}', 'area_1850', 'vol_1850',
             'area_2003', 'vol_2003', 'area_2015', 'vol_2015']
    M, filled = {}, {}
    for v in vars_:
        x = ds[v].to_pandas()                              # member x glacier
        miss = x.isnull()
        x = x.fillna(x.median())                           # member median of the glacier
        never = x.columns[x.isnull().all()]
        if len(never):
            # median ratio value / RGI area of glaciers of the same size class
            ratio = x.median() / rgi_a
            cls = pd.cut(rgi_a, [0, 0.05, 0.1, 0.5, 1, np.inf])
            r_cls = ratio.drop(never).groupby(cls.drop(never), observed=True).median()
            fill = cls[never].map(r_cls).astype(float) * rgi_a[never]
            x[never] = np.broadcast_to(fill.values, (len(x), len(never)))
        M[v] = x
        filled[v] = miss
    M['area'] = M.pop(f'area_max_{WIN}')
    M['vol'] = M.pop(f'vol_at_amax_{WIN}')
    filled_any = filled[f'area_max_{WIN}']
    return s, links, without, st, M, filled_any, ds[f'year_amax_{WIN}'].to_pandas()


def groups(s, links):
    """RGI glacier -> group LIA_ID; shared outline -> neighbour group."""
    main = links[links.is_main_lia].set_index('RGIId').LIA_ID
    main_status = main.map(s.status)
    g = main[main_status == 'modelled']
    sh = s.index[s.status == 'shared']
    ll = links[links.LIA_ID.isin(sh)].copy()
    ll['nb'] = ll.RGIId.map(main)
    ll = ll[ll.nb.map(s.status) == 'modelled']
    nb = ll.sort_values('overlap_km2').groupby('LIA_ID').nb.last()
    return g, main, main_status, nb


def glacier_names():
    """RGI6 names, completed with the photo-pair and Leclercq inventories."""
    n = pd.read_csv(RGI_ATTRS, index_col='RGIId', encoding='latin1').Name.str.strip()
    n = n.replace('', np.nan)
    ph = pd.read_csv(f'{INV}/glacier_before_after_pairs_rgi6.csv')
    ph = ph[ph.rgi60_score >= 0.7].drop_duplicates('rgi60_id').set_index('rgi60_id').glacier
    lec = pd.read_csv(f'{ROOT}/climate_analysis/leclercq_alps/leclercq_alps_pre1850.csv')
    lec = lec.drop_duplicates('rgi_id').set_index('rgi_id')['name']
    lec = lec.rename(index={'RGI60-11.03648': 'RGI60-11.03646'})   # Bossons (03648 is Bionnassay)
    return n.fillna(ph.reindex(n.index)).fillna(lec.reindex(n.index))


# --------------------------------------------------------------------------
# Step 1: regional
# --------------------------------------------------------------------------
def regional(s, without, st, M, main_status):
    nm = s.status.isin(NOT_MODELLED)
    obs_a, obs_v = s.lia_area_km2.sum(), s.lia_vol_km3.sum()
    nm_a, nm_v = s.loc[nm, 'lia_area_km2'].sum(), s.loc[nm, 'lia_vol_km3'].sum()
    fa, fv = nm_a / obs_a, nm_v / obs_v
    rem = main_status.index[main_status == 'remnant'].intersection(st.index)
    wo = pd.Index(without.RGIId).intersection(st.index)

    totals = {}
    for k in ['area', 'vol', 'area_1850', 'vol_1850', 'area_2015', 'vol_2015', 'area_2003']:
        totals[k] = M[k].sum(axis=1)
        totals[k + '_rem'] = M[k][rem].sum(axis=1)
        totals[k + '_wo'] = M[k][wo].sum(axis=1)
    T = pd.DataFrame(totals)
    cols = {REF: T.loc[REF], ALT: T.loc[ALT], 'ens P5': T.quantile(0.05),
            'ens P50': T.quantile(0.5), 'ens P95': T.quantile(0.95)}
    rows = {}
    for c, t in cols.items():
        r = {}
        for lab, a, v in [('LIA', 'area', 'vol'), ('1850', 'area_1850', 'vol_1850')]:
            r[f'{lab} area raw'] = t[a]
            r[f'{lab} vol raw'] = t[v]
            r[f'{lab} area on remnant glaciers'] = t[a + '_rem']
            r[f'{lab} vol on remnant glaciers'] = t[v + '_rem']
            r[f'{lab} area corr. additive'] = t[a] - t[a + '_rem'] + nm_a
            r[f'{lab} vol corr. additive'] = t[v] - t[v + '_rem'] + nm_v
            r[f'{lab} area corr. multiplicative'] = (t[a] - t[a + '_rem']) / (1 - fa)
            r[f'{lab} vol corr. multiplicative'] = (t[v] - t[v + '_rem']) / (1 - fv)
        r['LIA area on glaciers without LIA outline'] = t['area_wo']
        r['LIA vol on glaciers without LIA outline'] = t['vol_wo']
        r['area 2003'] = t['area_2003']
        r['area 2015'] = t['area_2015']
        r['vol 2015'] = t['vol_2015']
        # LIA -> 2015: the not-modelled outlines are assumed ice-free by 2015
        r['dV LIA-2015 corr.'] = r['LIA vol corr. additive'] - (t['vol_2015'] - t['vol_2015_rem'])
        r['area ratio LIA/2015 corr.'] = r['LIA area corr. additive'] / t['area_2015']
        rows[c] = r
    R = pd.DataFrame(rows)
    R['obs'] = np.nan
    R.loc['LIA area raw', 'obs'] = R.loc['1850 area raw', 'obs'] = obs_a
    R.loc['LIA vol raw', 'obs'] = R.loc['1850 vol raw', 'obs'] = obs_v
    for lab in ['LIA', '1850']:
        for k in ['additive', 'multiplicative']:
            R.loc[f'{lab} area corr. {k}', 'obs'] = obs_a
            R.loc[f'{lab} vol corr. {k}', 'obs'] = obs_v
    R.loc['dV LIA-2015 corr.', 'obs'] = RT_DV
    R.loc['area ratio LIA/2015 corr.', 'obs'] = obs_a / RT_A2015
    R.loc['area 2015', 'obs'] = RT_A2015
    R['obs 1sigma'] = np.nan
    R.loc[R.index.str.contains('area corr|area raw'), 'obs 1sigma'] = obs_a * AREA_ERR
    R.loc[R.index.str.contains('vol corr|vol raw'), 'obs 1sigma'] = obs_v * VOL_ERR
    R.loc['dV LIA-2015 corr.', 'obs 1sigma'] = RT_DV_ERR
    info = dict(obs_a=obs_a, obs_v=obs_v, nm_a=nm_a, nm_v=nm_v, fa=fa, fv=fv,
                n_nm=int(nm.sum()), n_rem=len(rem), n_wo=len(wo),
                wo_rgi=st.rgi_area_km2[wo].sum())
    return R, T, info


# --------------------------------------------------------------------------
# Step 2: per outline
# --------------------------------------------------------------------------
def per_outline(s, links, st, M, filled, g, nb):
    """One row per group (head = the modelled outline)."""
    s = s.copy()
    s['group'] = np.nan
    s.loc[s.status == 'modelled', 'group'] = s.index[s.status == 'modelled']
    s.loc[nb.index, 'group'] = nb.values
    gs = s.dropna(subset=['group'])
    gs = gs.assign(group=gs.group.astype(int))
    obs = gs.groupby('group')[['lia_area_km2', 'lia_vol_km3']].sum()
    obs.columns = ['obs_area', 'obs_vol']
    head = s.loc[obs.index]
    po = obs.join(head[['Citation', 'reg_ID', 'MIN', 'MAX', 'MEDIAN', 'lat', 'lon', 'Mean_thk']])
    po['obs_area_own'] = head.lia_area_km2
    po['obs_vol_own'] = head.lia_vol_km3
    shd = gs[gs.status == 'shared']
    po['shared_ids'] = pd.Series(shd.index.astype(str), index=shd.index).groupby(
        shd.group).agg(';'.join).reindex(po.index)
    po['n_shared'] = po.shared_ids.fillna('').str.count(';') + po.shared_ids.notnull()
    gg = g[g.index.isin(st.index)]
    po['n_glaciers'] = gg.groupby(gg).size().reindex(po.index).fillna(0).astype(int)
    po['rgi_area'] = st.rgi_area_km2.reindex(gg.index).groupby(gg).sum().reindex(po.index)
    po['rgi_ids'] = pd.Series(gg.index, index=gg.index).groupby(gg).apply(';'.join).reindex(po.index)
    big = st.rgi_area_km2.reindex(gg.index).groupby(gg).idxmax().reindex(po.index)
    po['main_rgi'] = big
    names = glacier_names()
    po['main_name'] = names.reindex(big.values).values
    po['n_filled_glaciers'] = filled[gg.index].loc[REF].groupby(gg).sum().reindex(po.index).astype(int)
    po['shrink'] = po.rgi_area / po.obs_area

    lr = {}
    for k in ['area', 'vol']:
        mg = M[k][gg.index].T.groupby(gg).sum().T.reindex(columns=po.index)   # member x group
        lmg = np.log(mg.clip(lower=1e-9))
        for c, lab in [(REF, 'ref'), (ALT, 'alt')]:
            po[f'mod_{k}_{lab}'] = mg.loc[c]
            po[f'lr_{k}_{lab}'] = np.log(mg.loc[c] / po[f'obs_{k}'])
            po[f'd_{k}_{lab}'] = mg.loc[c] - po[f'obs_{k}']
        for q in [0.05, 0.5, 0.95]:
            po[f'mod_{k}_p{int(q * 100):02d}'] = mg.quantile(q)
        po[f'mod_{k}_min'], po[f'mod_{k}_max'] = mg.min(), mg.max()
        po[f'lr_{k}_p50'] = np.log(po[f'mod_{k}_p50'] / po[f'obs_{k}'])
        po[f'sd_ln_{k}'] = lmg.std()
        po[f'obs_in_p05_p95_{k}'] = (po[f'obs_{k}'] >= po[f'mod_{k}_p05']) & (po[f'obs_{k}'] <= po[f'mod_{k}_p95'])
        po[f'obs_in_minmax_{k}'] = (po[f'obs_{k}'] >= po[f'mod_{k}_min']) & (po[f'obs_{k}'] <= po[f'mod_{k}_max'])
        lr[k] = mg
        # drop option: shared outlines dropped (obs = head outline only)
        po[f'lr_{k}_ref_drop'] = np.log(po[f'mod_{k}_ref'] / po[f'obs_{k}_own'])

    # present-day state at the RGI date (2003) against RGI area / inversion volume
    for k, ref in [('area', st.rgi_area_km2), ('vol', st.inv_volume_km3)]:
        m03 = M[f'{k}_2003'].loc[REF][gg.index].groupby(gg).sum().reindex(po.index)
        r03 = ref.reindex(gg.index).groupby(gg).sum().reindex(po.index)
        po[f'mod_{k}_2003_ref'] = m03
        po[f'lr_{k}_2003_ref'] = np.log(m03.clip(lower=1e-3) / r03)
    po['inv_vol'] = st.inv_volume_km3.reindex(gg.index).groupby(gg).sum().reindex(po.index)

    # FRAC option: every glacier spread over all its outlines by frac_of_rgi
    lk = links[links.RGIId.isin(st.index)]
    fr = {}
    for k in ['area', 'vol']:
        val = lk.RGIId.map(M[k].loc[REF]) * lk.frac_of_rgi
        fr[k] = val.groupby(lk.LIA_ID).sum()
    return po, lr, fr


def outline_csv(s, po, fr):
    """One row per modelled / shared outline, joinable on LIA_ID."""
    rows = s[s.status.isin(['modelled', 'shared'])][['status', 'Citation', 'reg_ID',
                                                       'lia_area_km2', 'lia_vol_km3']].copy()
    rows.columns = ['status', 'Citation', 'reg_ID', 'obs_area_km2', 'obs_vol_km3']
    grp = pd.Series(po.index, index=po.index)
    for sid, gid in po.shared_ids.dropna().items():
        for i in str(sid).split(';'):
            grp[int(i)] = gid
    rows['group_LIA_ID'] = grp.reindex(rows.index)
    keep = ['obs_area', 'obs_vol', 'rgi_ids', 'main_name', 'n_glaciers', 'rgi_area', 'shrink',
            'n_filled_glaciers',
            'mod_area_ref', 'mod_vol_ref', 'lr_area_ref', 'lr_vol_ref', 'd_area_ref', 'd_vol_ref',
            'mod_area_alt', 'mod_vol_alt', 'lr_area_alt', 'lr_vol_alt',
            'mod_area_p05', 'mod_area_p50', 'mod_area_p95', 'mod_vol_p05', 'mod_vol_p50',
            'mod_vol_p95', 'sd_ln_area', 'sd_ln_vol', 'obs_in_p05_p95_area', 'obs_in_p05_p95_vol', 'mod_area_2003_ref', 'lr_area_2003_ref', 'inv_vol', 'mod_vol_2003_ref', 'lr_vol_2003_ref']
    gcols = po[keep].rename(columns={'obs_area': 'group_obs_area_km2', 'obs_vol': 'group_obs_vol_km3'})
    gcols = gcols.rename(columns=lambda c: c.replace('_ref', f'_{REF}').replace('_alt', f'_{ALT}'))
    out = rows.join(gcols, on='group_LIA_ID')
    out[f'frac_mod_area_{REF}'] = fr['area'].reindex(out.index)
    out[f'frac_mod_vol_{REF}'] = fr['vol'].reindex(out.index)
    out[f'frac_lr_area_{REF}'] = np.log(out[f'frac_mod_area_{REF}'] / out.obs_area_km2)
    out[f'frac_lr_vol_{REF}'] = np.log(out[f'frac_mod_vol_{REF}'] / out.obs_vol_km3)
    out.index.name = 'LIA_ID'
    return out


# --------------------------------------------------------------------------
# Step 3a: systematic vs scatter
# --------------------------------------------------------------------------
def design(po, which):
    X = []
    if 'size' in which:
        la = np.log10(po.obs_area)
        X += [la, la ** 2]
    if 'shrink' in which:
        X += [np.log(po.shrink)]
    if 'elev' in which:
        X += [po.MEDIAN / 1000, po.MIN / 1000]
    if 'p03' in which:
        X += [po.lr_area_2003_ref]
    X = [np.ones(len(po))] + [np.asarray(x, float) for x in X]
    for cat in ['reg', 'cit']:
        if cat in which:
            c = po.reg_ID if cat == 'reg' else po.Citation
            d = pd.get_dummies(c, drop_first=True).astype(float).values
            X += list(d.T)
    return np.column_stack(X)


def r2(X, y, w=None, cv=0, seed=0):
    w = np.ones(len(y)) if w is None else w
    def fit(Xa, ya, wa):
        sw = np.sqrt(wa)
        return np.linalg.lstsq(Xa * sw[:, None], ya * sw, rcond=None)[0]
    if cv:
        rng = np.random.default_rng(seed)
        fold = rng.integers(0, cv, len(y))
        pred = np.empty(len(y))
        for k in range(cv):
            tr = fold != k
            pred[~tr] = X[~tr] @ fit(X[tr], y[tr], w[tr])
    else:
        pred = X @ fit(X, y, w)
    ybar = np.average(y, weights=w)
    return 1 - np.sum(w * (y - pred) ** 2) / np.sum(w * (y - ybar) ** 2)


def variance(po):
    sets = {'size class (log area, quadratic)': ['size'], 'RGI/LIA area (shrinkage)': ['shrink'],
            'elevation (median, min)': ['elev'], 'subregion (reg_ID)': ['reg'],
            'source inventory (Citation)': ['cit'], 'subregion + inventory': ['reg', 'cit'],
            'size + elevation + subregion + inventory': ['size', 'elev', 'reg', 'cit'],
            'all incl. shrinkage': ['size', 'shrink', 'elev', 'reg', 'cit'],
            'model 2003 area / RGI area (diagnostic)': ['p03'],
            'all incl. shrinkage + 2003 mismatch': ['size', 'shrink', 'elev', 'reg', 'cit', 'p03']}
    rows = []
    d = po.dropna(subset=['MEDIAN', 'MIN'])
    for k in ['area', 'vol']:
        y = d[f'lr_{k}_ref'].values
        w = d.obs_area.values
        for lab, which in sets.items():
            X = design(d, which)
            rows.append({'target': f'ln(model/obs) {k}', 'predictors': lab, 'n_params': X.shape[1],
                         'R2': r2(X, y), 'R2_cv5': np.mean([r2(X, y, cv=5, seed=i) for i in range(5)]),
                         'R2_area_weighted': r2(X, y, w)})
    V = pd.DataFrame(rows)
    # noise levels to compare with the residual
    noise = {'n_outlines': len(d),
             'sd_lr_area': d.lr_area_ref.std(), 'sd_lr_vol': d.lr_vol_ref.std(),
             'rms_ens_sd_area': np.sqrt((d.sd_ln_area ** 2).mean()),
             'rms_ens_sd_vol': np.sqrt((d.sd_ln_vol ** 2).mean())}
    return V, noise


# --------------------------------------------------------------------------
# Figures
# --------------------------------------------------------------------------
def size_class(x):
    return pd.cut(x, SIZE_BINS, labels=SIZE_LBL)


def fig_regional(R, info, po):
    fig, axs = plt.subplots(2, 2, figsize=(11, 7.5))
    for j, (k, unit, err) in enumerate([('area', 'km²', AREA_ERR), ('vol', 'km³', VOL_ERR)]):
        ax = axs[0, j]
        obs = info['obs_a'] if k == 'area' else info['obs_v']
        ax.axhspan(obs * (1 - err), obs * (1 + err), color=GRID, lw=0, label='Reinthaler & Paul ±1σ')
        ax.axhline(obs, color=C_OBS, lw=1.5, label=f'Reinthaler & Paul: {obs:.0f} {unit}')
        cats = ['raw', 'corr. additive', 'corr. multiplicative']
        x = np.arange(len(cats))
        for dx, (c, col, lab) in zip([-0.15, 0.15], [(REF, C_REF, f'{REF} (calibration member)'),
                                                     (ALT, C_ALT, f'{ALT} (attribution state)')]):
            y = [R.loc[f'LIA {k} {cc}', c] for cc in cats]
            ax.plot(x + dx, y, 'o', color=col, ms=8, mec='white', mew=1.5, label=lab, zorder=3)
            for xi, yi in zip(x + dx, y):
                ax.annotate(f'{yi:.0f}', (xi, yi), xytext=(7, 0), textcoords='offset points',
                            va='center', fontsize=8, color=INK2)
        lo = [R.loc[f'LIA {k} {cc}', 'ens P5'] for cc in cats]
        hi = [R.loc[f'LIA {k} {cc}', 'ens P95'] for cc in cats]
        md = [R.loc[f'LIA {k} {cc}', 'ens P50'] for cc in cats]
        ax.errorbar(x, md, yerr=[np.subtract(md, lo), np.subtract(hi, md)], fmt='s', color=C_ENS,
                    ms=6, lw=2, capsize=0, label='192 HIST members: median, P5-P95', zorder=2)
        ax.set_xticks(x, ['model', f'+ not-modelled\noutlines (+{info["nm_a"] if k == "area" else info["nm_v"]:.1f})',
                          f'÷ (1 − {info["fa"] if k == "area" else info["fv"]:.3f})'])
        ax.set_xlim(-0.5, 2.6)
        ax.set_ylabel(f'LIA {"area" if k == "area" else "volume"} ({unit})')
        ax.set_title(f'({"ab"[j]}) Regional LIA {"area" if k == "area" else "volume"}, '
                     f'{WIN.replace("_", "-")}', loc='left', fontsize=10)
        if j == 0:
            ax.legend(fontsize=7.5, frameon=False, loc='lower right')
        # (c, d): where the difference sits, by LIA size class
        ax = axs[1, j]
        sc = size_class(po.obs_area)
        x = np.arange(len(SIZE_LBL) + 1)
        for dx, (lab, col) in zip([-0.2, 0.2], [('ref', C_REF), ('alt', C_ALT)]):
            d = po[f'd_{k}_{lab}'].groupby(sc, observed=False).sum()
            vals = list(d.values) + [-(info['nm_a'] if k == 'area' else info['nm_v'])]
            ax.bar(x + dx, vals, width=0.38, color=col, label=REF if lab == 'ref' else ALT)
        ax.axhline(0, color=INK2, lw=0.8)
        ax.set_xticks(x, SIZE_LBL + ['not in\nRGI6'])
        ax.set_xlabel('LIA outline size class (km²)')
        ax.set_ylabel(f'model − obs ({unit})')
        ax.set_title(f'({"cd"[j]}) Where the difference sits (per-outline sums)', loc='left', fontsize=10)
        if j == 0:
            ax.legend(fontsize=8, frameon=False)
    fig.tight_layout()
    fig.savefig(f'{FIGS}/lia_regional_{WIN}.png', dpi=150)
    plt.close(fig)


def ecdf(ax, x, w=None, **kw):
    x = np.asarray(x)
    o = np.argsort(x)
    w = np.ones(len(x)) if w is None else np.asarray(w)[o]
    ax.step(x[o], np.cumsum(w) / w.sum(), where='post', **kw)


def fig_distributions(po, out):
    fig, axs = plt.subplots(2, 2, figsize=(11, 7))
    fr = out[out.status == 'modelled']
    for j, (k, lab) in enumerate([('area', 'area'), ('vol', 'volume')]):
        ax = axs[0, j]
        bins = np.linspace(-3, 3, 61)
        ax.hist(po[f'lr_{k}_ref'].clip(-3, 3), bins=bins, color=C_REF, alpha=0.85, label=REF)
        ax.hist(po[f'lr_{k}_alt'].clip(-3, 3), bins=bins, histtype='step', color=C_ALT, lw=1.5, label=ALT)
        ax.axvline(0, color=INK, lw=0.8)
        med = po[f'lr_{k}_ref'].median()
        ax.set_title(f'({"ab"[j]}) ln(model / obs), LIA {lab}, {len(po)} outlines', loc='left', fontsize=10)
        ax.text(0.02, 0.95, f'{REF}: median {med:+.2f} (×{np.exp(med):.2f})\n'
                f'{ALT}: median {po[f"lr_{k}_alt"].median():+.2f} (×{np.exp(po[f"lr_{k}_alt"].median()):.2f})',
                transform=ax.transAxes, va='top', fontsize=8, color=INK2)
        ax.set_xlabel('ln(model / obs)  (clipped at ±3)')
        ax.set_ylabel('outlines')
        ax.legend(frameon=False, fontsize=8)
        ax = axs[1, j]
        ecdf(ax, po[f'lr_{k}_ref'], color=C_REF, lw=2, label=f'{REF}')
        ecdf(ax, po[f'lr_{k}_alt'], color=C_ALT, lw=2, label=f'{ALT}')
        ecdf(ax, po[f'lr_{k}_p50'], color=C_ENS, lw=2, label='member median')
        ecdf(ax, po[f'lr_{k}_ref'], w=po[f'obs_{k}'], color=C_REF, lw=1.2, ls='--',
             label=f'{REF}, weighted by obs {lab}')
        ecdf(ax, fr[f'frac_lr_{k}_{REF}'].replace([np.inf, -np.inf], np.nan).dropna(),
             color=INK2, lw=1.2, ls=':', label=f'{REF}, frac_of_rgi weighting')
        ax.axvline(0, color=INK, lw=0.8)
        ax.set_xlim(-3, 3)
        ax.set_xlabel('ln(model / obs)')
        ax.set_ylabel('cumulative fraction')
        ax.set_title(f'({"cd"[j]}) ECDF, LIA {lab}', loc='left', fontsize=10)
        ax.legend(frameon=False, fontsize=7.5, loc='upper left')
    fig.tight_layout()
    fig.savefig(f'{FIGS}/lia_distributions_{WIN}.png', dpi=150)
    plt.close(fig)


def breakdown_cats(po):
    elev_med = pd.cut(po.MEDIAN, [0, 2600, 2800, 3000, 3200, 5000],
                      labels=['<2600', '2600-2800', '2800-3000', '3000-3200', '>3200'])
    elev_min = pd.cut(po.MIN, [0, 2000, 2300, 2600, 2900, 5000],
                      labels=['<2000', '2000-2300', '2300-2600', '2600-2900', '>2900'])
    shrink = pd.cut(po.shrink, [0, 0.1, 0.25, 0.5, 0.75, 10],
                    labels=['<0.1', '0.1-0.25', '0.25-0.5', '0.5-0.75', '>0.75'])
    cit = po.Citation.map(CIT_SHORT)
    cit = pd.Categorical(cit, categories=cit.value_counts().index)
    return {'LIA size class (km²)': size_class(po.obs_area),
            'Source inventory': pd.Series(cit, index=po.index),
            'reg_ID (subregion)': po.reg_ID,
            'Median elevation (m)': elev_med, 'Minimum elevation (m)': elev_min,
            'RGI6 area / LIA area': shrink}


def fig_breakdown(po, k):
    cats = breakdown_cats(po)
    fig, axs = plt.subplots(2, 3, figsize=(14, 8))
    y = po[f'lr_{k}_ref']
    for ax, (name, c) in zip(axs.flat, cats.items()):
        c = pd.Series(c, index=po.index)
        lev = [l for l in (c.cat.categories if hasattr(c, 'cat') else sorted(c.dropna().unique()))
               if (c == l).any()]
        data = [y[c == l].values for l in lev]
        bp = ax.boxplot(data, whis=(5, 95), showfliers=False, widths=0.55, patch_artist=True,
                        medianprops=dict(color=INK, lw=1.5))
        for b in bp['boxes']:
            b.set(facecolor='#cde2fb', edgecolor=C_REF)
        agg = [np.log(po.loc[c == l, f'mod_{k}_ref'].sum() / po.loc[c == l, f'obs_{k}'].sum()) for l in lev]
        agg_alt = [np.log(po.loc[c == l, f'mod_{k}_alt'].sum() / po.loc[c == l, f'obs_{k}'].sum()) for l in lev]
        xs = np.arange(1, len(lev) + 1)
        ax.plot(xs, agg, 'D', color=C_REF, ms=7, mec='white', mew=1.2, label=f'Σ model / Σ obs, {REF}')
        ax.plot(xs, agg_alt, 'D', color=C_ALT, ms=6, mec='white', mew=1.2, label=f'Σ model / Σ obs, {ALT}')
        ax.axhline(0, color=INK, lw=0.8)
        for xi, l in zip(xs, lev):
            n = int((c == l).sum())
            ax.text(xi, -2.85, f'{n}', ha='center', fontsize=7, color=INK2)
        ax.set_xticks(xs, [str(l) for l in lev], rotation=35 if len(lev) > 6 else 0, ha='right' if len(lev) > 6 else 'center')
        ax.set_ylim(-3, 2)
        ax.set_title(name, loc='left', fontsize=10)
        ax.set_ylabel(f'ln(model / obs), {"area" if k == "area" else "volume"}')
    axs[0, 0].legend(frameon=False, fontsize=7.5, loc='upper right')
    fig.suptitle(f'Per-outline LIA {"area" if k == "area" else "volume"} mismatch, {REF}; '
                 f'boxes = P25-P75, whiskers P5-P95, numbers = outlines', fontsize=10, x=0.01, ha='left')
    fig.tight_layout()
    fig.savefig(f'{FIGS}/lia_breakdown_{k}_{WIN}.png', dpi=150)
    plt.close(fig)


def fig_spread(po):
    fig, axs = plt.subplots(1, 3, figsize=(15, 4.6))
    ax = axs[0]
    o = po.sort_values('obs_area')
    ax.vlines(o.obs_area, o.mod_area_p05, o.mod_area_p95, color=C_ENS, lw=0.8, label='members P5-P95')
    ax.plot(o.obs_area, o.mod_area_ref, '.', color=C_REF, ms=3, label=REF)
    lim = [0.005, 200]
    ax.plot(lim, lim, color=INK, lw=0.8)
    ax.set_xscale('log'); ax.set_yscale('log'); ax.set_xlim(lim); ax.set_ylim(lim)
    ax.set_xlabel('LIA area, obs (km²)'); ax.set_ylabel('LIA area, model (km²)')
    ax.set_title('(a) Per outline: model vs obs', loc='left', fontsize=10)
    ax.legend(frameon=False, fontsize=8)
    ax = axs[1]
    for k, col, lab in [('area', C_REF, 'area'), ('vol', C_ALT, 'volume')]:
        ratio = (po[f'lr_{k}_ref'].abs() / po[f'sd_ln_{k}']).clip(upper=100)
        ecdf(ax, np.log10(ratio), color=col, lw=2, label=f'{lab}: median {ratio.median():.1f}')
    ax.axvline(0, color=INK, lw=0.8)
    ax.set_xlabel('log10( |ln(model/obs)| / member spread σ(ln model) )')
    ax.set_ylabel('cumulative fraction of outlines')
    ax.set_title('(b) Mismatch relative to the ensemble spread', loc='left', fontsize=10)
    ax.legend(frameon=False, fontsize=8)
    ax = axs[2]
    sc = size_class(po.obs_area)
    x = np.arange(len(SIZE_LBL))
    for dx, (k, col, lab) in zip([-0.2, 0.2], [('area', C_REF, 'area'), ('vol', C_ALT, 'volume')]):
        f = po[f'obs_in_p05_p95_{k}'].groupby(sc, observed=False).mean() * 100
        ax.bar(x + dx, f.values, width=0.38, color=col, label=lab)
    ax.set_xticks(x, SIZE_LBL)
    ax.set_xlabel('LIA outline size class (km²)')
    ax.set_ylabel('% of outlines with obs inside P5-P95')
    ax.set_title('(c) Obs inside the 192-member range', loc='left', fontsize=10)
    ax.legend(frameon=False, fontsize=8)
    fig.tight_layout()
    fig.savefig(f'{FIGS}/lia_spread_{WIN}.png', dpi=150)
    plt.close(fig)


def fig_map(po):
    fig, axs = plt.subplots(1, 2, figsize=(15, 5.6))
    ax = axs[0]
    o = po.sort_values('obs_area', ascending=False)
    norm = TwoSlopeNorm(0, -1.5, 1.5)
    sca = ax.scatter(o.lon, o.lat, c=o.lr_area_ref.clip(-1.5, 1.5), s=4 + 6 * np.sqrt(o.obs_area),
                     cmap=DIV, norm=norm, edgecolor=INK2, linewidth=0.2)
    fig.colorbar(sca, ax=ax, label='ln(model / obs), LIA area (red: model too small)', shrink=0.85)
    ax.set_title(f'(a) Per outline, {REF} (marker size ~ √LIA area)', loc='left', fontsize=10)
    ax = axs[1]
    d = 0.5
    cx, cy = np.floor(po.lon / d) * d, np.floor(po.lat / d) * d
    agg = po.groupby([cx, cy]).agg(m=('mod_area_ref', 'sum'), o=('obs_area', 'sum'), n=('obs_area', 'size'))
    for (x0, y0), r in agg.iterrows():
        v = np.log(r.m / r.o)
        ax.add_patch(plt.Rectangle((x0, y0), d, d, color=DIV(norm(np.clip(v, -1.5, 1.5))), ec='white', lw=1))
        ax.text(x0 + d / 2, y0 + d / 2, f'{np.exp(v):.2f}\n{r.o:.0f}', ha='center', va='center', fontsize=6.5,
                color=INK)
    ax.set_xlim(po.lon.min() - 0.3, po.lon.max() + 0.3)
    ax.set_ylim(po.lat.min() - 0.3, po.lat.max() + 0.3)
    ax.set_title('(b) 0.5° cells: Σ model / Σ obs (top) and obs LIA area km² (bottom)', loc='left', fontsize=10)
    for a in axs:
        a.set_aspect(1 / np.cos(np.deg2rad(46.3)))
        a.set_xlabel('lon'); a.set_ylabel('lat')
    fig.tight_layout()
    fig.savefig(f'{FIGS}/lia_map_{WIN}.png', dpi=150)
    plt.close(fig)


def fig_shrink(po):
    """Does the model reproduce how much each glacier shrank since the LIA?"""
    fig, axs = plt.subplots(1, 3, figsize=(16, 4.8))
    ax = axs[2]
    x3, y3 = po.lr_area_2003_ref.clip(-3, 2), po.lr_area_ref.clip(-3, 3)
    ax.scatter(x3, y3, s=3 + 4 * np.sqrt(po.obs_area), color=C_REF, alpha=0.5, lw=0)
    ax.plot([-3, 2], [-3, 2], color=INK, lw=0.8, label='1:1')
    ax.axhline(0, color=INK2, lw=0.6); ax.axvline(0, color=INK2, lw=0.6)
    big = po.obs_area > 5
    rr = np.corrcoef(po.lr_area_2003_ref, po.lr_area_ref)[0, 1]
    rb = np.corrcoef(po.lr_area_2003_ref[big], po.lr_area_ref[big])[0, 1]
    for gid in po[po.obs_area > 10].index.union(po.index[po.main_rgi == 'RGI60-11.02072']):
        if pd.notnull(po.main_name[gid]):
            ax.annotate(str(po.main_name[gid]).split(',')[-1].strip()[:18], (x3[gid], y3[gid]),
                        xytext=(4, -3), textcoords='offset points', fontsize=6.5, color=INK2)
    ax.set_xlabel('ln(model area 2003 / RGI area)  — present-day misfit (clipped)')
    ax.set_ylabel('ln(model / obs), LIA area')
    ax.set_title(f'(c) {REF}: r = {rr:.2f} (outlines > 5 km²: {rb:.2f})', loc='left', fontsize=10)
    ax.legend(frameon=False, fontsize=8, loc='upper left')
    x = np.log(po.obs_area / po.rgi_area)
    for ax, (lab, col, c) in zip(axs[:2], [('ref', C_REF, REF), ('alt', C_ALT, ALT)]):
        y = np.log(po[f'mod_area_{lab}'] / po.rgi_area)
        ax.scatter(x, y, s=3 + 4 * np.sqrt(po.obs_area), color=col, alpha=0.5, lw=0)
        lim = [-1, 6]
        ax.plot(lim, lim, color=INK, lw=0.8, label='1:1')
        w = po.obs_area.values
        b = np.polyfit(x, y, 1, w=np.sqrt(w))
        b0 = np.polyfit(x, y, 1)
        xx = np.array(lim)
        ax.plot(xx, np.polyval(b, xx), color=INK2, lw=1.5, ls='--',
                label=f'fit, area weighted: slope {b[0]:.2f}')
        ax.plot(xx, np.polyval(b0, xx), color=INK2, lw=1, ls=':', label=f'fit, unweighted: slope {b0[0]:.2f}')
        ax.set_xlim(lim); ax.set_ylim(-1, 4)
        ax.set_xlabel('ln(LIA area obs / RGI area)  — observed shrinkage')
        ax.set_ylabel('ln(LIA area model / RGI area)  — modelled growth')
        r = np.corrcoef(x, y)[0, 1]
        ax.set_title(f'({"ab"[lab == "alt"]}) {c}: r = {r:.2f}', loc='left', fontsize=10)
        ax.legend(frameon=False, fontsize=8, loc='upper left')
    fig.tight_layout()
    fig.savefig(f'{FIGS}/lia_shrinkage_{WIN}.png', dpi=150)
    print(f'r(2003 misfit, LIA lr): all {rr:.2f}, > 5 km2 {rb:.2f}; '
          f'2003 area misfit sd {po.lr_area_2003_ref.std():.2f}, |misfit| > ln 1.25: '
          f'{(po.lr_area_2003_ref.abs() > np.log(1.25)).mean():.0%} of outlines, '
          f'{np.average(po.lr_area_2003_ref.abs() > np.log(1.25), weights=po.obs_area):.0%} of LIA area')
    plt.close(fig)


def fmt_table(R):
    cols = [REF, ALT, 'ens P5', 'ens P50', 'ens P95', 'obs', 'obs 1sigma']
    return R[cols].round(2)


if __name__ == '__main__':
    s, links, without, st, M, filled, yamax = load()
    g, main, main_status, nb = groups(s, links)
    R, T, info = regional(s, without, st, M, main_status)
    R.to_csv(f'{HERE}/lia_regional_{WIN}.csv')
    print(info)
    print(fmt_table(R).to_string())
    print(f'filled: {int(filled.loc[REF].sum())} glaciers in {REF}, '
          f'{st.rgi_area_km2[filled.loc[REF]].sum():.1f} km2 RGI, model LIA area '
          f'{M["area"].loc[REF, filled.loc[REF]].sum():.1f} km2')

    po, mg, fr = per_outline(s, links, st, M, filled, g, nb)
    out = outline_csv(s, po, fr)
    out.to_csv(f'{HERE}/lia_per_outline_{WIN}.csv')
    print(f'{len(po)} groups ({int((po.n_shared > 0).sum())} with shared outlines merged, '
          f'{int(po.n_shared.sum())} shared outlines); obs {po.obs_area.sum():.1f} km2 / {po.obs_vol.sum():.2f} km3')
    for k in ['area', 'vol']:
        for lab in ['ref', 'alt', 'p50']:
            x = po[f'lr_{k}_{lab}']
            print(f'{k} {lab}: median lr {x.median():+.3f} (x{np.exp(x.median()):.2f}), '
                  f'P16-P84 {x.quantile(.16):+.2f}..{x.quantile(.84):+.2f}, sd {x.std():.2f}, '
                  f'|lr|<ln1.5: {(x.abs() < np.log(1.5)).mean():.0%}')
        print(f'  {k} drop shared: median {po[f"lr_{k}_ref_drop"].median():+.3f}; '
              f'frac weighting (modelled outlines): median '
              f'{out.loc[out.status == "modelled", f"frac_lr_{k}_{REF}"].median():+.3f}, '
              f'(shared outlines): {out.loc[out.status == "shared", f"frac_lr_{k}_{REF}"].median():+.3f}')
        print(f'  {k} obs inside P5-P95: {po[f"obs_in_p05_p95_{k}"].mean():.0%} '
              f'(area-weighted {np.average(po[f"obs_in_p05_p95_{k}"], weights=po.obs_area):.0%}), '
              f'inside min-max {po[f"obs_in_minmax_{k}"].mean():.0%}; median |lr|/sd '
              f'{(po[f"lr_{k}_ref"].abs() / po[f"sd_ln_{k}"]).median():.1f}')

    V, noise = variance(po)
    V.to_csv(f'{HERE}/lia_variance_{WIN}.csv', index=False)
    print(V.round(3).to_string())
    print(noise)

    worst = po.reindex(po.d_area_ref.abs().sort_values(ascending=False).index[:20])
    w20 = worst[['main_name', 'main_rgi', 'n_glaciers', 'Citation', 'reg_ID', 'obs_area', 'mod_area_ref',
                 'mod_area_p05', 'mod_area_p95', 'd_area_ref', 'lr_area_ref', 'obs_vol', 'mod_vol_ref',
                 'd_vol_ref', 'rgi_area', 'mod_area_2003_ref', 'lr_area_2003_ref', 'shrink', 'MIN',
                 'MEDIAN', 'n_shared']]
    w20.to_csv(f'{HERE}/lia_worst20_{WIN}.csv')
    print(w20.round(2).to_string())

    fig_regional(R, info, po)
    fig_distributions(po, out)
    fig_breakdown(po, 'area')
    fig_breakdown(po, 'vol')
    fig_spread(po)
    fig_map(po)
    fig_shrink(po)
    # year of the max area (ens010), area weighted, per reg_ID
    y = yamax.loc[REF].reindex(g.index).dropna()
    wy = st.rgi_area_km2.reindex(y.index)
    reg = g.reindex(y.index).map(s.reg_ID)
    print('area-weighted mean year of max area per reg_ID (ens010):')
    print(((y * wy).groupby(reg).sum() / wy.groupby(reg).sum()).round(0).to_string())