# Anthropogenic share of the regional volume loss since 1850, with the NAT
# counterfactual conditioned on the 1850 state (round6 permutation runs).
# Conditional version of round5/attribution/attributable_fraction.py.
#
# For state i and NAT climate j (annual values, no window):
#   total loss          L_i(t)  = V_i(1850) - V_hist,i(t)
#   anthropogenic loss  A_ij(t) = V_nat,ij(t) - V_hist,i(t)
#   share               f_ij(t) = A_ij(t) / L_i(t)      (median / 5-95% over j)
# Both HIST and the conditional NAT forget the 1850 state, so A is almost
# the same for every state and f ~ A / (V(1850) - V_hist): a hyperbola in
# the 1850 volume. A share > 100% means the glaciers would have grown
# without anthropogenic forcing.
#
# Reinthaler & Paul (2025): the observed LIA -> 2015 volume change, -180 ± 39
# km3, put on the model as V(1850) = model HIST 2015 + 180 ± 39 (as the
# anchored LIA reference of round5's talk figures; their map is the ~1850
# extent). In the model the per-glacier 1800-1875 maximum differs from the
# 1850 regional volume by 3-55 km3 for these states.
#
# Panels: (a) regional volume, HIST and conditional NAT median per state,
# Reinthaler at 1850; (b) the share since 1850, 1950-2021; (c) the share in
# REF_YEAR vs the 1850 volume, the hyperbola, the 197 NAT 1850 volumes and
# Reinthaler.
# Output: figs/attr_fraction_conditional_<NAME>.png / .csv (REF_YEAR 2015),
# ..._ref<REF_YEAR>.png / .csv otherwise.
# REF_YEAR after 2021 (the last year of the NAT runs): each NAT run is held
# at the mean of its last NAT_LAST years after it ends (they are flat since
# ~1900, as in attr_fraction_committed.py); dashed / shaded in (a) and (b).
# Reinthaler's V(1850) is always anchored on the model HIST of 2015 (their
# change is LIA -> 2015), whatever REF_YEAR: share = A / (V_rt - V_hist(REF_YEAR)).
# Run from round6/attribution/ with the oggm_env python (needs
# ../preprocess_perm.py <NAME>). Env: NAME, ENS, REF_YEAR (2015), VAR (volume).
# VAR=area: the same for the regional area (area_min_h), output suffix
# _area. Reinthaler & Paul (2025) area change LIA -> 2015: 4211 -> 1806 km2,
# i.e. 2405 ± 215 km2 (their ±5.1% on the LIA area; the 2015 error is
# smaller), anchored on the model HIST 2015 area as for the volume.
import os

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

NAME = os.environ.get('NAME', 'perm_pf3.5_mf5_lam12')
ENS = os.environ.get('ENS', 'ens_pf3.5_mf5_lam12')
REF_YEAR = int(os.environ.get('REF_YEAR', 2015))
VAR = os.environ.get('VAR', 'volume')
CFG = {'volume': dict(col='volume_km3', unit='km$^3$', u='km3', dv=180., err=39.),
       'area': dict(col='area_km2', unit='km$^2$', u='km2', dv=4211.12 - 1806.24,
                    err=0.051 * 4211.12)}[VAR]
COL, UNIT, U = CFG['col'], CFG['unit'], CFG['u']
STATES = ['123', '023', '117', '054', '010', '041', '178']
COLORS = dict(zip(STATES, plt.cm.Reds(np.linspace(0.35, 0.95, len(STATES)))))
BLUES = dict(zip(STATES, plt.cm.Blues(np.linspace(0.35, 0.95, len(STATES)))))
REINTHALER_DV, REINTHALER_ERR = CFG['dv'], CFG['err']
# Ranges: 16-84% (~ ±1 sigma) everywhere, to match Reinthaler's ±1 sigma. At
# the Reinthaler volume both uncertainties are combined by sampling: A from
# the NAT runs (pooled over the states), V(1850) ~ N(v_rt, REINTHALER_ERR).
QLO, QHI = 0.16, 0.84
QTAG = '16–84%'
N_MC = 200_000
ANCHOR_YEAR = 2015          # Reinthaler's change is LIA -> 2015
NAT_END = 2021              # last year of the NAT runs
NAT_LAST = 5
FRAC_START = 1950
T_END = max(NAT_END, REF_YEAR)
HOLD = REF_YEAR > NAT_END
SUFFIX = ('' if REF_YEAR == 2015 else f'_ref{REF_YEAR}') + ('' if VAR == 'volume' else f'_{VAR}')

v = pd.read_csv(f'../filtered/{NAME}/{COL}_perm_reg11.csv', index_col=0)
pre = pd.read_csv(f'../filtered/{NAME}/pre1850_reg11.csv', index_col=0)
v5 = pd.read_csv(f'../../round5/filtered_ens/{ENS}/{COL}_nat_ens_reg11.csv', index_col=0)
for df in (v, pre, v5):
    df.index = df.index.astype(int)

res, rows = {}, []
for s in STATES:
    nat = v[[c for c in v.columns if c.startswith(f'state{s}_nat_')]]
    if HOLD:   # each run held at the mean of its last NAT_LAST years
        last = nat.apply(lambda x: x.dropna().iloc[-NAT_LAST:].mean())
        nat = nat.reindex(range(nat.index[0], T_END + 1))
        for c in nat.columns:
            end = nat[c].last_valid_index()
            nat.loc[end + 1:, c] = last[c]
    hist = v[f'state{s}_hist']
    v0 = hist.loc[1850]
    A = nat.sub(hist, axis=0)                       # anthropogenic loss per NAT climate
    L = v0 - hist                                   # total loss since 1850
    f = A.div(L, axis=0)
    q = f.quantile([QLO, 0.5, QHI], axis=1).T
    res[s] = dict(nat=nat, hist=hist, v0=v0, A=A, L=L, q=q)
    for y in sorted({2000, 2015, REF_YEAR, T_END}):
        rows.append(dict(state=f'ens{s}', year=y, V1850=v0, V_hist=hist.loc[y],
                         V_nat_median=nat.loc[y].median(), total_loss=L.loc[y],
                         anthro_loss_median=A.loc[y].median(),
                         share_plo=q.loc[y, QLO], share_median=q.loc[y, 0.5],
                         share_phi=q.loc[y, QHI], n_nat=int(nat.loc[y].notnull().sum())))
tab = pd.DataFrame(rows).set_index(['state', 'year'])

# The common quantities in REF_YEAR (mean over the states)
r = tab.xs(REF_YEAR, level='year')
A_med, Vh = r.anthro_loss_median.mean(), r.V_hist.mean()
A_q = {p: np.mean([res[s]['A'].loc[REF_YEAR].quantile(p) for s in STATES]) for p in (QLO, QHI)}
Vh_anchor = np.mean([res[s]['hist'].loc[ANCHOR_YEAR] for s in STATES])
v_rt = Vh_anchor + REINTHALER_DV
f_rt = A_med / (v_rt - Vh)
f_rt_lo, f_rt_hi = A_med / (v_rt + REINTHALER_ERR - Vh), A_med / (v_rt - REINTHALER_ERR - Vh)
# combined uncertainty at the Reinthaler volume
A_all = np.concatenate([res[s]['A'].loc[REF_YEAR].dropna().values for s in STATES])
rng = np.random.default_rng(0)
f_mc = rng.choice(A_all, N_MC) / (rng.normal(v_rt, REINTHALER_ERR, N_MC) - Vh)
f_c = {p: np.quantile(f_mc, p) for p in (0.05, QLO, 0.5, QHI, 0.95)}
print(f'Reinthaler, NAT and Reinthaler combined: median {100 * f_c[0.5]:.0f}%, '
      f'{QTAG} {100 * f_c[QLO]:.0f}-{100 * f_c[QHI]:.0f}%, 5-95% {100 * f_c[0.05]:.0f}-{100 * f_c[0.95]:.0f}%')
print(f'{REF_YEAR}: anthropogenic loss A = {A_med:.1f} {U} (states {r.anthro_loss_median.min():.1f}'
      f'-{r.anthro_loss_median.max():.1f}; NAT {QTAG} {A_q[QLO]:.1f}-{A_q[QHI]:.1f}), '
      f'HIST {VAR} {Vh:.1f}')
print(f'Reinthaler: V(1850) = {v_rt:.0f} ± {REINTHALER_ERR:.0f} {U} -> share {100 * f_rt:.0f}% '
      f'({100 * f_rt_lo:.0f}-{100 * f_rt_hi:.0f}% from its ±1σ; NAT {QTAG} at the central '
      f'value: {100 * A_q[QLO] / (v_rt - Vh):.0f}-{100 * A_q[QHI] / (v_rt - Vh):.0f}%)')

fig = plt.figure(figsize=(15, 5.2))
gs = fig.add_gridspec(1, 3, width_ratios=[1.25, 1, 1], wspace=0.28)
ax_a, ax_b, ax_c = (fig.add_subplot(gs[0, i]) for i in range(3))

# (a) volumes
for s in STATES:
    d = res[s]
    yrs = np.arange(1820, T_END + 1)
    h = pd.concat([pre[f'{COL}_state{s}'].loc[1820:1849], d['hist'].loc[1850:T_END]])
    n = d['nat'].median(axis=1).loc[1850:T_END]
    ax_a.plot(h.loc[:1850].index, h.loc[:1850], color='0.5', lw=1)
    ax_a.plot(n.loc[:NAT_END].index, n.loc[:NAT_END], color=BLUES[s], lw=1.6)
    if HOLD:
        ax_a.plot(n.loc[NAT_END:].index, n.loc[NAT_END:], color=BLUES[s], lw=1.6, ls='--')
    ax_a.plot(h.loc[1850:].index, h.loc[1850:], color=COLORS[s], lw=1.8)
ax_a.errorbar([1850], [v_rt], yerr=[REINTHALER_ERR], fmt='D', color='k', ms=6, capsize=4,
              zorder=5, label=f'Reinthaler & Paul (2025), anchored\non model {ANCHOR_YEAR}: '
                              f'{v_rt:.0f} ± {REINTHALER_ERR:.0f} {UNIT}')
ax_a.plot([], [], color=COLORS['054'], lw=1.8, label='HIST, 7 states (colours as in b)')
ax_a.plot([], [], color=BLUES['117'], lw=1.6, label='NAT from the same 1850 state (median)')
ax_a.plot([], [], color='0.5', lw=1, label='before 1850 (shared)')
if HOLD:
    ax_a.axvspan(NAT_END, T_END, color='0.5', alpha=0.1, lw=0)
    ax_a.plot([], [], color=BLUES['117'], lw=1.6, ls='--', label='NAT held constant (runs end ≤ 2021)')
ax_a.set_xlim(1820, T_END)
ax_a.set_ylabel(f'Regional {VAR} ({UNIT})')
ax_a.set_title(f'(a) Regional {VAR}: HIST (red) and natural (blue)', loc='left', fontsize=10.5)
ax_a.legend(loc='lower left', fontsize=7.5, frameon=False)
ax_a.grid(alpha=0.3)

# (b) share since 1850 over time
for s in STATES:
    q = res[s]['q'].loc[FRAC_START:T_END]
    if s == '010':
        ax_b.fill_between(q.index, 100 * q[QLO], 100 * q[QHI], color=COLORS[s], alpha=0.12,
                          lw=0, label=f'ens010, NAT {QTAG}')
    ax_b.plot(q.index, 100 * q[0.5], color=COLORS[s], lw=1.8, label=f'ens{s}')
ax_b.axhline(100, color='0.3', lw=0.9)
ax_b.axvline(REF_YEAR, color='0.5', lw=0.8, ls=':')
if HOLD:
    ax_b.axvspan(NAT_END, T_END, color='0.5', alpha=0.1, lw=0)
    ax_b.text(NAT_END + 0.3, 5, 'NAT\nheld', fontsize=7, color='0.4')
ax_b.set_xlim(FRAC_START, T_END)
ax_b.set_ylim(0, 200)
ax_b.set_ylabel(f'Anthropogenic share of the\n{VAR} loss since 1850 (%)')
ax_b.set_title('(b) Share of the loss since 1850, per 1850 state', loc='left', fontsize=10.5)
ax_b.legend(loc='upper right', fontsize=7.5, frameon=False, ncol=2)
ax_b.grid(alpha=0.3)

# (c) share in REF_YEAR vs the 1850 volume
if VAR == 'volume':
    XLIM, HBINS = (Vh_anchor + 60, 340), np.arange(170, 345, 5)
else:
    XLIM = (0.9 * float(v5.loc[1850].min()), 1.03 * (v_rt + REINTHALER_ERR))
    HBINS = np.linspace(*XLIM, 70)
x = np.linspace(*XLIM, 300)
ax_c.fill_between(x, 100 * A_q[QLO] / (x - Vh), 100 * A_q[QHI] / (x - Vh), color='0.6',
                  alpha=0.2, lw=0, label=f'NAT {QTAG}')
ax_c.plot(x, 100 * A_med / (x - Vh), color='0.3', lw=1.5,
          label=f'{A_med:.0f} / (V$_{{1850}}$ − {Vh:.0f})')
for s in STATES:
    q = res[s]['q'].loc[REF_YEAR]
    ax_c.errorbar([res[s]['v0']], [100 * q[0.5]],
                  yerr=[[100 * (q[0.5] - q[QLO])], [100 * (q[QHI] - q[0.5])]],
                  fmt='o', color=COLORS[s], ms=6, capsize=3)
ax_c.axvspan(v_rt - REINTHALER_ERR, v_rt + REINTHALER_ERR, color='k', alpha=0.08, lw=0)
ax_c.axvline(v_rt, color='k', lw=1, ls='--', label=f'Reinthaler: {v_rt:.0f} ± {REINTHALER_ERR:.0f} {UNIT}')
ax_c.errorbar([v_rt], [100 * f_c[0.5]], yerr=[[100 * (f_c[0.5] - f_c[0.05])], [100 * (f_c[0.95] - f_c[0.5])]],
              fmt='none', ecolor='k', elinewidth=1, capsize=3, zorder=6)
ax_c.errorbar([v_rt], [100 * f_c[0.5]], yerr=[[100 * (f_c[0.5] - f_c[QLO])], [100 * (f_c[QHI] - f_c[0.5])]],
              fmt='D', color='k', ms=6, elinewidth=3, capsize=0, zorder=7,
              label=f'Reinthaler + NAT: {100 * f_c[0.5]:.0f}% ({100 * f_c[QLO]:.0f}–'
                    f'{100 * f_c[QHI]:.0f}%; 5–95%: {100 * f_c[0.05]:.0f}–{100 * f_c[0.95]:.0f}%)')
ax_c.axhline(100, color='0.3', lw=0.9)
# the 1850 volumes of all 197 NAT members (where the model's states lie)
ax_r = ax_c.inset_axes([0, 0, 1, 0.12], sharex=ax_c)
ax_r.hist(v5.loc[1850], bins=HBINS, color=BLUES['117'], alpha=0.6)
ax_r.axis('off')
ax_c.text(0.02, 0.135, f'1850 {VAR} of the 197 model members', transform=ax_c.transAxes,
          fontsize=7.5, color=BLUES['054'])
ax_c.set_xlim(*XLIM)
ax_c.set_ylim(0, 200)
ax_c.set_xlabel(f'Regional {VAR} in 1850 ({UNIT})')
ax_c.set_ylabel(f'Anthropogenic share of the loss\n1850–{REF_YEAR} (%)')
ax_c.set_title(f'(c) Share in {REF_YEAR} vs the 1850 state', loc='left', fontsize=10.5)
ax_c.legend(loc='upper right', fontsize=7.5, frameon=True, framealpha=0.9, edgecolor='none')
ax_c.grid(alpha=0.3)

fig.suptitle(f'How much of the Alpine ice loss since 1850 is anthropogenic? '
             f'NAT conditioned on the 1850 state ({NAME})', y=1.0)
os.makedirs('figs', exist_ok=True)
fig.savefig(f'figs/attr_fraction_conditional_{NAME}{SUFFIX}.png', dpi=150, bbox_inches='tight')
tab.round(3).to_csv(f'figs/attr_fraction_conditional_{NAME}{SUFFIX}.csv')
pd.Series({f'q{int(100 * k):02d}': v for k, v in f_c.items()}, name='share').round(4).to_csv(
    f'figs/attr_fraction_conditional_{NAME}{SUFFIX}_reinthaler_combined.csv')
with pd.option_context('display.width', 200):
    print(tab.xs(REF_YEAR, level='year').round(2).to_string())