# Round 7 copy of round6/attribution/explainer_conditional.py (2026-11-02): perm_a5v10 states 023/007/083/169/118;
# temperature from climate_analysis/series/nat_jjas_glacw_newnat_r7.csv (nat_spread_alps.py).
# Explainer figure for the conditional detection (leads to the SNR curves of
# conditional_temperature.py). Regional volume, one 1850 state (STATE):
#   (a) the experiment: the state's history up to 1850 (grey), then the
#       same glaciers under the 197 natural climates (thin blue) and the
#       historical climate (red); grey bands = the 15-yr windows of (c);
#   (b) signal and noise (km3) over time, for the same state;
#   (c) three snapshots: the 197 natural outcomes (histogram of the 15-yr
#       mean volume) vs the historical one: signal = NAT median - HIST,
#       noise = NAT std, SNR = signal / noise;
#   (d) the SNR over time (= the volume curve of
#       attr_temperature_vs_glacier_conditional_*_length.png), snapshots
#       marked.
# One state only: the other states are shown by explainer_two_states.py.
# Absolute 15-yr mean volumes: HIST and all natural runs of a state share
# its 1820-1850 baseline, so it cancels and is not needed here.
# Output: figs/explainer_conditional_<NAME>_<STATE>_w<W>.png
# Env: NAME, STATE (default 010), WINDOW, SNAP (default "1930 1975 2021").
import os

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

NAME = os.environ.get('NAME', 'perm_a5v10')
STATE = os.environ.get('STATE', '083')
WINDOW = int(os.environ.get('WINDOW', 15))
SNAP = [int(y) for y in os.environ.get('SNAP', '1930 1975 2021').split()]
NAT_C, NAT_DARK, HIST_C = '#5dade2', '#1b4f72', '#c0392b'
T_END = 2021

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


def series(s):
    """NAT (one column per climate) and HIST, with the state's pre-1850 history."""
    p = pre[f'volume_km3_state{s}'].loc[:1849]
    cols = [c for c in v.columns if c.startswith(f'state{s}_nat_')]
    nat = pd.concat([pd.DataFrame({c: p for c in cols}), v[cols].loc[1850:]])
    hist = pd.concat([p, v[f'state{s}_hist'].loc[1850:]])
    return nat, hist


def roll(x):
    return x.rolling(WINDOW, min_periods=WINDOW).mean()


yrs = np.arange(1851, T_END + 1)
nat, hist = series(STATE)
nat15, hist15 = roll(nat).loc[yrs], roll(hist).loc[yrs]
signal = nat15.median(axis=1) - hist15
noise = nat15.std(axis=1)
snr = signal / noise
print('SNR at the snapshots:', {y: round(snr.loc[y], 2) for y in SNAP})

fig = plt.figure(figsize=(14, 9.5))
gs = fig.add_gridspec(2, 6, height_ratios=[1.15, 1], hspace=0.38, wspace=0.55)
ax_a = fig.add_subplot(gs[0, :4])
ax_d = fig.add_subplot(gs[0, 4:])
ax_b = [fig.add_subplot(gs[1, i]) for i in range(3)]
ax_c = fig.add_subplot(gs[1, 3:])

# (a) the experiment
t = np.arange(1800, T_END + 1)
ax_a.plot(nat.loc[1850:T_END].index, nat.loc[1850:T_END], color=NAT_C, lw=0.4, alpha=0.25)
ax_a.plot(nat.loc[1850:T_END].median(axis=1), color=NAT_DARK, lw=2,
          label='natural climates: median')
ax_a.plot([], [], color=NAT_C, lw=1, label=f'natural climates ({nat.shape[1]} runs)')
ax_a.plot(hist.loc[1800:1850], color='0.3', lw=2, label='history of this LIA state (shared)')
ax_a.plot(hist.loc[1850:T_END], color=HIST_C, lw=2.4, label='historical climate')
ax_a.plot(1850, hist.loc[1850], 'o', color='k', ms=7, zorder=5)
ax_a.annotate('same glaciers in 1850,\ndifferent climates afterwards',
              (1850, hist.loc[1850]), xytext=(1862, hist.loc[1850] + 22), fontsize=9,
              arrowprops=dict(arrowstyle='->', color='0.3'))
for y in SNAP:
    ax_a.axvspan(y - WINDOW + 1, y, color='0.5', alpha=0.15, lw=0)
    ax_a.text(y - WINDOW / 2, 0.97, f'{y}', transform=ax_a.get_xaxis_transform(),
              ha='center', va='top', fontsize=8.5, color='0.3')
ax_a.set_xlim(1800, T_END)
ax_a.set_ylabel('Regional glacier volume (km$^3$)')
ax_a.set_title(f'(a) One Little Ice Age state (ens{STATE}), two kinds of climate after 1850',
               loc='left', fontsize=11)
ax_a.legend(loc='lower left', fontsize=8.5, frameon=False)
ax_a.grid(alpha=0.3)

# (b) snapshots
lo = min(nat15.loc[SNAP].min().min(), hist15.loc[SNAP].min()) - 12
hi = nat15.loc[SNAP].max().max() + 5
bins = np.arange(np.floor(lo), np.ceil(hi) + 3, 3)
for ax, y in zip(ax_b, SNAP):
    d, h = nat15.loc[y].dropna(), hist15.loc[y]
    med, sd = d.median(), d.std()
    ax.hist(d, bins=bins, orientation='horizontal', color=NAT_C, alpha=0.7)
    xm = ax.get_xlim()[1]
    ax.axhline(med, color=NAT_DARK, lw=2)
    ax.axhspan(med - sd, med + sd, color=NAT_DARK, alpha=0.12, lw=0)
    ax.axhline(h, color=HIST_C, lw=2.4)
    ax.annotate('', xy=(0.82 * xm, h), xytext=(0.82 * xm, med),
                arrowprops=dict(arrowstyle='<->', color='k', lw=1.2))
    ax.set_ylim(lo, hi)
    ax.set_title(f'{y - WINDOW + 1}–{y}\nSNR = {med - h:.0f} / {sd:.0f} = {(med - h) / sd:.1f}',
                 fontsize=9.5)
    ax.set_xlabel('number of\nnatural runs', fontsize=8.5)
    ax.tick_params(labelsize=8)
    if ax is ax_b[0]:
        ax.set_ylabel(f'{WINDOW}-yr mean volume (km$^3$)')
    else:
        ax.set_yticklabels([])
fig.text(0.125, 0.475, '(c) Snapshots: where does the historical glacier (red) fall among the '
         'natural ones (histogram)?', fontsize=11)
fig.text(0.125, 0.455, 'signal = natural median − historical (arrow, km$^3$),  noise = std of the '
         'natural runs (shaded: ± 1 std, km$^3$),  SNR = signal / noise', fontsize=8.5, color='0.25')

# (c) SNR
ax_c.plot(yrs, snr, color=HIST_C, lw=2.4)
# IPCC likelihood thresholds (one-sided normal): likely / very likely /
# extremely likely / virtually certain
for th, c, lab in zip([1, 1.65, 2, 2.58], ['#f5b041', '#e67e22', '#c0392b', '#7b241c'],
                      ['likely', 'very likely', 'extremely likely', 'virtually certain']):
    ax_c.axhline(th, color=c, lw=1, ls=':')
    ax_c.text(1.01, th, f'{lab} ({th:g})', transform=ax_c.get_yaxis_transform(),
              fontsize=8, color=c, va='center', ha='left')
    above = (snr > th)[::-1].cummin()[::-1]     # crosses and stays above until T_END
    print(f'SNR > {th:g} ({lab}) for good from: '
          f'{int(above.idxmax()) if above.any() else "never"}')
for y in SNAP:
    ax_c.plot(y, snr.loc[y], 'o', color='k', ms=6, zorder=5)
    ax_c.annotate(f'{y}', (y, snr.loc[y]), xytext=(-6, 8 if y < T_END else -14),
                  textcoords='offset points', ha='right', fontsize=8.5)
ax_c.axhline(0, color='0.4', lw=0.8)
ax_c.set_xlim(yrs[0], T_END)
ax_c.set_ylabel('SNR = signal / noise')
ax_c.set_xlabel(f'Year (end of the {WINDOW}-yr window)')
ax_c.set_title('(d) The same, every year', loc='left', fontsize=11)
ax_c.grid(alpha=0.3)

# (d) signal and noise of this state
ax_d.plot(yrs, signal, color=HIST_C, lw=2.2, label='signal: natural median − historical')
ax_d.plot(yrs, noise, color=NAT_DARK, lw=2, ls='--', label='noise: std of the natural runs')
for y in SNAP:
    ax_d.plot(y, signal.loc[y], 'o', color='k', ms=5, zorder=5)
ax_d.axhline(0, color='0.4', lw=0.8)
ax_d.set_xlim(yrs[0], T_END)
ax_d.set_ylabel('km$^3$')
ax_d.set_title(f'(b) Signal and noise ({WINDOW}-yr means)', loc='left', fontsize=11)
ax_d.legend(loc='upper left', fontsize=8, frameon=False)
ax_d.grid(alpha=0.3)

os.makedirs('figs', exist_ok=True)
out = f'figs/explainer_conditional_{NAME}_ens{STATE}_w{WINDOW}.png'
fig.savefig(out, dpi=150, bbox_inches='tight')
print('wrote', out)