"""Round 7 version of round6/attribution/attr_fraction_conditional.py (REF_YEAR 2025):
the anthropogenic share of the regional volume loss since 1850, on the round 7 ensemble
(default ens_v8_adapbs_a5v10: AdaPBS round 2 gdirs, LAPrec precipitation).

For state i (a HIST member) and NAT member j (annual values, no window):
    total loss          L_i(t)  = V_i(1850) - V_hist,i(t)
    anthropogenic loss  A_ij(t) = V_nat,j(t) - V_hist,i(t)
    share               f_ij(t) = A_ij(t) / L_i(t)      (median / 16-84% over j)

**The NAT side is conditional** (user, 2026-11-02: always from now on): the permutation
runs of round7/perm/ (PERM, default perm_a5v10): each state's 1850 glaciers restarted with
all NAT climates, and the state's HIST strand from the same restart. Without them (PERM
missing) it falls back to the full NAT ensemble, the same 197 runs for every state (output
suffix _uncond); on round 5 / round 6 that gives A within 1.5 km3 of the conditional one
from 2000 on, but up to 6 km3 off around 1900.

States: the 5 HIST members at the min / P25 / median / P75 / max of the 1850 regional volume
across the NAT members (as round7/plot_lia_max.py).

Reinthaler & Paul (2025): LIA -> 2015 volume change −180 ± 39 km3, **minus the ice of the
outlines that the model does not cover** (no_rgi / remnant / sliver_only: 9.9 km3 at the
LIA, round6/lia_validation/validate.py; assumed gone by 2015), i.e. the change on the modelled
glaciers, anchored on the model HIST 2015 volume: V(1850) = V_hist(2015) + (180 − 9.9) ± 39.
NAT runs end in 2006-2021; after that each is held at the mean of its last 5 years.

    [NAME=...] [REF_YEAR=2025] python attr_fraction.py
    -> figs/attr_fraction_<NAME>[_uncond]_ref<REF_YEAR>.png / .csv (+ _reinthaler_combined.csv)
"""
import os
import sys

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

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

NAME = os.environ.get('NAME', 'ens_v8_adapbs_a5v10')
REF_YEAR = int(os.environ.get('REF_YEAR', 2025))
FDIR = f'{R7}/filtered_ens/{NAME}'
OUT = f'{HERE}/figs'
os.makedirs(OUT, exist_ok=True)
QLO, QHI, QTAG = 0.16, 0.84, '16–84%'
N_MC = 200_000
ANCHOR_YEAR = 2015
NAT_LAST = 5
FRAC_START = 1950
RT_DV, RT_ERR = 180., 39.

vn = pd.read_csv(f'{FDIR}/volume_km3_nat_ens_reg11.csv', index_col=0)
vh = pd.read_csv(f'{FDIR}/volume_km3_hist_ens_reg11.csv', index_col=0)
vn.index, vh.index = vn.index.astype(int), vh.index.astype(int)
NAT_END = int(vn.apply(lambda x: x.last_valid_index()).max())
T_END = max(NAT_END, REF_YEAR)
HOLD = REF_YEAR > NAT_END

# the NAT runs, held at the mean of their last NAT_LAST years after they end
last = vn.apply(lambda x: x.dropna().iloc[-NAT_LAST:].mean())
nat = vn.reindex(range(vn.index[0], T_END + 1))
for c in nat.columns:
    end = nat[c].last_valid_index()
    nat.loc[end + 1:, c] = last[c]

# the 5 states: 1850 volume quantiles across the NAT members
v50 = vn.loc[1850]
STATES = []
for q in [0, 0.25, 0.5, 0.75, 1]:
    STATES.append((v50 - v50.quantile(q)).abs().idxmin())
pct = {s: 100 * (v50 < vh.loc[1850, s]).mean() for s in STATES}
COLORS = dict(zip(STATES, plt.cm.Reds(np.linspace(0.4, 0.95, len(STATES)))))
BLUE, BLUE_L = '#184e8f', '#7fa9de'

# not-modelled outlines (validate.py)
s_, links, without, st, M, filled, yamax = V.load()
g, main, main_status, nb = V.groups(s_, links)
_, _, info = V.regional(s_, without, st, M, main_status)
NM_V = info['nm_v']
DV_MOD = RT_DV - NM_V

# Conditional NAT (permutation runs, round7/perm/): one NAT ensemble per state, all from its
# 1850 glaciers, and the HIST strand from the same restart. Falls back to the full NAT
# ensemble if the permutation dataset is missing.
PERM = os.environ.get('PERM', 'perm_a5v10')
PFILE = f'{R7}/perm/filtered/{PERM}/volume_km3_perm_reg11.csv'
COND = os.path.exists(PFILE)
if COND:
    pv = pd.read_csv(PFILE, index_col=0)
    pv.index = pv.index.astype(int)
    pv = pv.reindex(range(1850, T_END + 1))
print(f'NAT side: {"conditional, " + PERM if COND else "full NAT ensemble (no permutation runs)"}')


def held(df):
    last = df.apply(lambda x: x.dropna().iloc[-NAT_LAST:].mean())
    df = df.copy()
    for c in df.columns:
        end = df[c].last_valid_index()
        df.loc[end + 1:, c] = last[c]
    return df


res, rows = {}, []
for s in STATES:
    if COND:
        k = s[3:]
        nat_s = held(pv[[c for c in pv.columns if c.startswith(f'state{k}_nat_')]])
        hist = vh[s].reindex(range(vh.index[0], T_END + 1)).copy()
        hist.loc[1850:] = pv[f'state{k}_hist'].loc[1850:T_END].values
    else:
        nat_s = nat
        hist = vh[s].reindex(range(vh.index[0], T_END + 1))
    v0 = hist.loc[1850]
    A = nat_s.sub(hist, axis=0).dropna(how='all')
    L = v0 - hist
    f = A.div(L, axis=0)
    q = f.quantile([QLO, 0.5, QHI], axis=1).T
    res[s] = dict(hist=hist, v0=v0, A=A, L=L, q=q, nat=nat_s)
    for y in sorted({2000, 2015, REF_YEAR, T_END}):
        rows.append(dict(state=s, year=y, V1850=v0, V_hist=hist.loc[y],
                         V_nat_median=nat_s.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]))
tab = pd.DataFrame(rows).set_index(['state', 'year'])

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 + DV_MOD
v_rt_raw = Vh_anchor + RT_DV
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, RT_ERR, N_MC) - Vh)
f_c = {p: np.quantile(f_mc, p) for p in (0.05, QLO, 0.5, QHI, 0.95)}
f_raw = A_med / (v_rt_raw - Vh)
print('states (1850 volume percentile): ' + ', '.join(f'{s} {res[s]["v0"]:.1f} km3 P{pct[s]:.0f}'
                                                      for s in STATES))
print(f'{REF_YEAR}: A = {A_med:.1f} km3 (states {r.anthro_loss_median.min():.1f}-'
      f'{r.anthro_loss_median.max():.1f}; NAT {QTAG} {A_q[QLO]:.1f}-{A_q[QHI]:.1f}), HIST {Vh:.1f} km3')
print(f'Reinthaler on the modelled glaciers: dV {RT_DV:.0f} - {NM_V:.1f} (not-modelled outlines) = '
      f'{DV_MOD:.1f} ± {RT_ERR:.0f} -> V(1850) = {v_rt:.1f} km3 (uncorrected {v_rt_raw:.1f})')
print(f'Reinthaler + NAT combined: median {100 * f_c[0.5]:.0f}%, {QTAG} {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}% '
      f'(without the outline correction: {100 * f_raw:.0f}% at the central values)')

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
BLUES = dict(zip(STATES, plt.cm.Blues(np.linspace(0.4, 0.95, len(STATES)))))
if COND:     # one NAT median per state, from the state's 1850 glaciers
    for s in STATES:
        n = res[s]['nat'].median(axis=1).loc[1850:T_END]
        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([], [], color=BLUES[STATES[2]], lw=1.6, label='NAT from the same 1850 state (median)')
else:
    nq = nat.quantile([QLO, 0.5, QHI], axis=1).T.loc[1820:T_END]
    ax_a.fill_between(nq.loc[:NAT_END].index, nq.loc[:NAT_END, QLO], nq.loc[:NAT_END, QHI],
                      color=BLUE_L, alpha=0.3, lw=0, label=f'NAT ensemble {QTAG} (197 members)')
    ax_a.plot(nq.loc[:NAT_END].index, nq.loc[:NAT_END, 0.5], color=BLUE, lw=1.8,
              label='NAT ensemble median')
if HOLD:
    ax_a.axvspan(NAT_END, T_END, color='0.5', alpha=0.1, lw=0)
for s in STATES:
    h = res[s]['hist']
    ax_a.plot(h.loc[1820:1850].index, h.loc[1820:1850], color='0.5', lw=1)
    ax_a.plot(h.loc[1850:T_END].index, h.loc[1850:T_END], color=COLORS[s], lw=1.8)
ax_a.errorbar([1850], [v_rt], yerr=[RT_ERR], fmt='D', color='k', ms=6, capsize=4, zorder=5,
              label=f'Reinthaler & Paul (2025) − not-modelled\noutlines, anchored on model '
                    f'{ANCHOR_YEAR}: {v_rt:.0f} ± {RT_ERR:.0f} km$^3$')
ax_a.plot([], [], color=COLORS[STATES[2]], lw=1.8, label='HIST, 5 states (colours as in b)')
ax_a.plot([], [], color='0.5', lw=1, label='before 1850 (each state = its NAT member)')
ax_a.set_xlim(1820, T_END)
ax_a.set_ylabel('Regional volume (km$^3$)')
ax_a.set_title('(a) Regional volume: 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 == STATES[2]:
        ax_b.fill_between(q.index, 100 * q[QLO], 100 * q[QHI], color=COLORS[s], alpha=0.12,
                          lw=0, label=f'{s}, NAT {QTAG}')
    ax_b.plot(q.index, 100 * q[0.5], color=COLORS[s], lw=1.8, label=f'{s} (1850: P{pct[s]:.0f})')
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('Anthropogenic share of the\nvolume 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
XLIM = (0.9 * float(v50.min()), 1.05 * (v_rt + RT_ERR))
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 - RT_ERR, v_rt + RT_ERR, color='k', alpha=0.08, lw=0)
ax_c.axvline(v_rt, color='k', lw=1, ls='--',
             label=f'Reinthaler (modelled glaciers): {v_rt:.0f} ± {RT_ERR:.0f} km$^3$')
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)
ax_r = ax_c.inset_axes([0, 0, 1, 0.12], sharex=ax_c)
ax_r.hist(v50, bins=np.linspace(*XLIM, 70), color=BLUE_L, alpha=0.7)
ax_r.axis('off')
ax_c.text(0.02, 0.135, '1850 volume of the 197 model members', transform=ax_c.transAxes,
          fontsize=7.5, color=BLUE)
ax_c.set_xlim(*XLIM)
ax_c.set_ylim(0, 200)
ax_c.set_xlabel('Regional volume in 1850 (km$^3$)')
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('How much of the Alpine ice loss since 1850 is anthropogenic? Round 7 '
             f'({NAME}; ' + (f'NAT conditioned on the 1850 state, {PERM})' if COND else
                             'NAT = full ensemble, no permutation runs)'), y=1.0)
TAGC = ('' if PERM == 'perm_a5v10' else '_' + PERM.replace('perm_a5v10_', '')) if COND else '_uncond'
fig.savefig(f'{OUT}/attr_fraction_{NAME}{TAGC}_ref{REF_YEAR}.png', dpi=150, bbox_inches='tight')
tab.round(3).to_csv(f'{OUT}/attr_fraction_{NAME}{TAGC}_ref{REF_YEAR}.csv')
pd.Series({f'q{int(100 * k):02d}': v for k, v in f_c.items()}, name='share').round(4).to_csv(
    f'{OUT}/attr_fraction_{NAME}{TAGC}_ref{REF_YEAR}_reinthaler_combined.csv')
with pd.option_context('display.width', 200):
    print(tab.xs(REF_YEAR, level='year').round(2).to_string())