# Explainer for the attribution of the loss (companion of
# explainer_conditional.py, which explains the detection), one 1850 state
# only (default ens010, "the only Alps we have"):
#   (a) the experiment: the state's history (grey), its 1850 glaciers under
#       the 197 natural climates (blue) and the historical climate (red);
#   (b) the decomposition in REF_YEAR: total loss since 1850 = natural part
#       (V1850 - natural median) + anthropogenic part (natural median - HIST);
#   (c) one share per natural climate: histogram of A_j / L, median and
#       16-84%;
#   (d) the share through time (median, 16-84%), from FRAC_START.
# Same definitions as attr_fraction_conditional.py (annual values, no window):
#   L(t) = V(1850) - V_hist(t),  A_j(t) = V_nat,j(t) - V_hist(t),  f_j = A_j / L.
# REF_YEAR after 2021: each NAT run that has ended is held at the mean of its
# last NAT_LAST years (as attr_fraction_conditional.py; the held part is not
# drawn in (a), the median is dashed there).
# Output: figs/explainer_attribution_<NAME>_ens<STATE>_ref<REF_YEAR>.png
# Env: NAME, STATE (default 010), REF_YEAR (default 2025).
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')
STATE = os.environ.get('STATE', '010')
REF_YEAR = int(os.environ.get('REF_YEAR', 2025))
NAT_C, NAT_DARK, HIST_C, ANTH_C = '#5dade2', '#1b4f72', '#c0392b', '#c0392b'
QLO, QHI = 0.16, 0.84
NAT_END, NAT_LAST = 2021, 5
FRAC_START = 1950

v = pd.read_csv(f'../filtered/{NAME}/volume_km3_perm_reg11.csv', index_col=0)
pre = pd.read_csv(f'../filtered/{NAME}/pre1850_reg11.csv', index_col=0)
v.index, pre.index = v.index.astype(int), pre.index.astype(int)

nat_raw = v[[c for c in v.columns if c.startswith(f'state{STATE}_nat_')]]
hist = v[f'state{STATE}_hist']
T_END = max(NAT_END, REF_YEAR)
nat = nat_raw.reindex(range(nat_raw.index[0], T_END + 1))
last = nat_raw.apply(lambda x: x.dropna().iloc[-NAT_LAST:].mean())
for c in nat.columns:                      # hold ended runs (only used after their end)
    end = nat_raw[c].last_valid_index()
    if REF_YEAR > end:
        nat.loc[end + 1:, c] = last[c]
hist = hist.loc[:T_END]
v0 = hist.loc[1850]
L = v0 - hist
A = nat.sub(hist, axis=0)
f = A.div(L, axis=0)
q = f.quantile([QLO, 0.5, QHI], axis=1).T
nat_med = nat.median(axis=1)

Lr, Ar, Nr = L.loc[REF_YEAR], nat_med.loc[REF_YEAR] - hist.loc[REF_YEAR], v0 - nat_med.loc[REF_YEAR]
fr = f.loc[REF_YEAR]
print(f'ens{STATE}, {REF_YEAR}: V1850 {v0:.1f}, HIST {hist.loc[REF_YEAR]:.1f}, NAT median '
      f'{nat_med.loc[REF_YEAR]:.1f} km3; total loss {Lr:.1f} = natural {Nr:.1f} + '
      f'anthropogenic {Ar:.1f}; share median {100 * fr.median():.0f}% '
      f'(16-84% {100 * fr.quantile(QLO):.0f}-{100 * fr.quantile(QHI):.0f}%), '
      f'{int(nat_raw.loc[min(REF_YEAR, NAT_END)].notnull().sum())} runs reach {min(REF_YEAR, NAT_END)}, '
      f'{(fr > 1).sum()} of {fr.notnull().sum()} > 100%')

fig = plt.figure(figsize=(14, 9.5))
gs = fig.add_gridspec(2, 6, height_ratios=[1.15, 1], hspace=0.32, wspace=0.7)
ax_a = fig.add_subplot(gs[0, :4])
ax_b = fig.add_subplot(gs[0, 4:])
ax_c = fig.add_subplot(gs[1, :3])
ax_d = fig.add_subplot(gs[1, 3:])

# (a) the experiment
p = pre[f'volume_km3_state{STATE}'].loc[:1849]
ax_a.plot(nat_raw.index, nat_raw.values, color=NAT_C, alpha=0.25, lw=0.6)
ax_a.plot(nat_med.loc[:NAT_END].index, nat_med.loc[:NAT_END], color=NAT_DARK, lw=2,
          label='natural climates: median')
if T_END > NAT_END:
    ax_a.plot(nat_med.loc[NAT_END:].index, nat_med.loc[NAT_END:], color=NAT_DARK, lw=2, ls='--')
ax_a.plot([], [], color=NAT_C, lw=1, label=f'natural climates ({nat.shape[1]} runs)')
ax_a.plot(list(p.index) + [1850], list(p.values) + [v0], color='0.3', lw=2,
          label='history of this LIA state')
ax_a.plot(hist.loc[1850:].index, hist.loc[1850:], color=HIST_C, lw=2.4, label='historical climate')
ax_a.axhline(v0, color='0.3', lw=1, ls='--')
ax_a.text(1802, v0 + 3, f'1850 volume: {v0:.0f} km$^3$', fontsize=9, color='0.3')
ax_a.plot(1850, v0, 'o', color='k', ms=7, zorder=5)
ax_a.axvline(REF_YEAR, color='0.5', lw=1, ls=':')
ax_a.set_xlim(1800, T_END + 2)
ax_a.set_ylabel('Regional glacier volume (km$^3$)')
ax_a.set_title(f'(a) One Little Ice Age state (ens{STATE}): what was lost, and what would '
               f'have been lost anyway?', loc='left', fontsize=11)
ax_a.legend(loc='lower left', fontsize=9, frameon=False)
ax_a.grid(alpha=0.3)

# (b) the decomposition in REF_YEAR
x = 0
ax_b.bar(x, Nr, bottom=hist.loc[REF_YEAR] + Ar, color=NAT_C, width=0.5, edgecolor='k', lw=0.6)
ax_b.bar(x, Ar, bottom=hist.loc[REF_YEAR], color=ANTH_C, width=0.5, edgecolor='k', lw=0.6)
ax_b.text(x + 0.3, hist.loc[REF_YEAR] + Ar + Nr / 2,
          f'natural: {Nr:.0f} km$^3$\n(V$_{{1850}}$ − natural median)', va='center', fontsize=9)
ax_b.text(x + 0.3, hist.loc[REF_YEAR] + Ar / 2,
          f'anthropogenic: {Ar:.0f} km$^3$\n(natural median − historical)', va='center',
          fontsize=9, color=ANTH_C)
for y, lab, c in [(v0, f'1850: {v0:.0f}', '0.3'), (nat_med.loc[REF_YEAR],
                  f'natural {REF_YEAR}: {nat_med.loc[REF_YEAR]:.0f}', NAT_DARK),
                  (hist.loc[REF_YEAR], f'historical {REF_YEAR}: {hist.loc[REF_YEAR]:.0f}', HIST_C)]:
    ax_b.plot([x - 0.4, x + 0.25], [y, y], color=c, lw=2)
    ax_b.text(x - 0.42, y, lab, ha='right', va='center', fontsize=9, color=c)
ax_b.set_xlim(-1.6, 1.6)
ax_b.set_ylim(hist.loc[REF_YEAR] - 15, v0 + 15)
ax_b.set_xticks([])
ax_b.set_yticks([])       # the levels are labelled
ax_b.set_title(f'(b) Loss 1850–{REF_YEAR}: {Lr:.0f} km$^3$\n'
               f'anthropogenic share = {Ar:.0f} / {Lr:.0f} = {100 * Ar / Lr:.0f}%',
               loc='left', fontsize=11)
for s in ('top', 'right', 'left', 'bottom'):
    ax_b.spines[s].set_visible(False)

# (c) one share per natural climate
vals = 100 * fr.dropna()
ax_c.hist(vals, bins=np.arange(np.floor(vals.min() / 5) * 5, vals.max() + 5, 5),
          color=NAT_C, edgecolor='white')
ax_c.axvspan(100 * fr.quantile(QLO), 100 * fr.quantile(QHI), color='0.6', alpha=0.2, lw=0,
             label='16–84% of the natural climates')
ax_c.axvline(100 * fr.median(), color=ANTH_C, lw=2.4, label=f'median: {100 * fr.median():.0f}%')
ax_c.axvline(100, color='0.3', lw=1, ls='--')
ax_c.text(101, ax_c.get_ylim()[1] * 0.75, '> 100%: the glaciers would\nhave grown without us',
          fontsize=8.5, color='0.3', va='top')
ax_c.set_xlabel(f'Anthropogenic share of the 1850–{REF_YEAR} loss (%)')
ax_c.set_ylabel('number of natural climates')
ax_c.set_title(f'(c) Each natural climate gives one share\nmedian {100 * fr.median():.0f}%, '
               f'shaded 16–84%: {100 * fr.quantile(QLO):.0f}–{100 * fr.quantile(QHI):.0f}%',
               loc='left', fontsize=11)
ax_c.grid(alpha=0.3)

# (d) the share through time
yrs = np.arange(FRAC_START, T_END + 1)
qq = 100 * q.loc[yrs]
ax_d.fill_between(yrs, qq[QLO], qq[QHI], color=ANTH_C, alpha=0.18, lw=0, label='16–84%')
ax_d.plot(yrs, qq[0.5], color=ANTH_C, lw=2.4, label='median')
if T_END > NAT_END:
    ax_d.axvspan(NAT_END, T_END, color='0.85', alpha=0.5, lw=0)
    ax_d.text(NAT_END + 0.5, 5, 'natural runs\nheld', fontsize=7.5, color='0.4')
ax_d.plot(REF_YEAR, 100 * fr.median(), 'o', color='k', ms=6, zorder=5)
ax_d.axhline(100, color='0.3', lw=1, ls='--')
ax_d.set_xlim(FRAC_START, T_END)
ax_d.set_ylim(0, max(130, qq[QHI].max() + 5))
ax_d.set_ylabel('Anthropogenic share of the loss since 1850 (%)')
ax_d.set_xlabel('Year')
ax_d.set_title('(d) The same, every year\n(share of the loss since 1850)', loc='left', fontsize=11)
ax_d.legend(loc='lower left', fontsize=9, frameon=False)
ax_d.grid(alpha=0.3)

os.makedirs('figs', exist_ok=True)
out = f'figs/explainer_attribution_{NAME}_ens{STATE}_ref{REF_YEAR}.png'
fig.savefig(out, dpi=150, bbox_inches='tight')
print('wrote', out)
print((100 * q.loc[[1950, 1975, 2000, 2015, REF_YEAR]]).round(0).to_string())