# Round 7 copy of round6/attribution/explainer_two_states.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).
# Companion of explainer_conditional.py, without the SNR: the two extreme
# 1850 states (smallest and largest regional volume), regional volume,
# annual values.
#   (a), (b) each state: its history (grey), then the same glaciers under the
#       197 natural climates (thin blue) and the historical climate (red);
#   (c) both states together: historical runs and the natural median with
#       its 16-84% range -> both pairs converge, whatever the 1850 state;
#   (d) the gap natural median - historical (the anthropogenic loss, km3)
#       with the natural 16-84% range, for both states.
# Output: figs/explainer_two_states_<NAME>_ens<A>_ens<B>.png
# Env: NAME, STATES (default "123 178").
import os

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

NAME = os.environ.get('NAME', 'perm_a5v10')
PAIR = os.environ.get('STATES', '023 118').split()
T_END = 2021
QLO, QHI = 0.16, 0.84
STYLE = {PAIR[0]: dict(nat='#85c1e9', nat_dark='#2e86c1', hist='#f1948a', ls='--', z=4),
         PAIR[1]: dict(nat='#1f618d', nat_dark='#154360', hist='#922b21', ls='-', z=3)}

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):
    p = pre[f'volume_km3_state{s}'].loc[:1849]
    cols = [c for c in v.columns if c.startswith(f'state{s}_nat_')]
    return v[cols].loc[1850:T_END], pd.concat([p, v[f'state{s}_hist'].loc[1850:T_END]])


data = {s: series(s) for s in PAIR}
fig, axes = plt.subplots(2, 2, figsize=(14, 9), gridspec_kw=dict(hspace=0.3, wspace=0.18))
ylim = (60, 300)

# (a), (b)
for ax, s, lab in zip(axes[0], PAIR, ['(a)', '(b)']):
    nat, hist = data[s]
    ax.plot(nat.index, nat, color='#5dade2', lw=0.4, alpha=0.25)
    ax.plot(nat.median(axis=1), color='#1b4f72', lw=2, label='natural climates: median')
    ax.plot([], [], color='#5dade2', lw=1, label=f'natural climates ({nat.shape[1]} runs)')
    ax.plot(hist.loc[1800:1850], color='0.3', lw=2, label='history of this LIA state')
    ax.plot(hist.loc[1850:], color='#c0392b', lw=2.4, label='historical climate')
    ax.plot(1850, hist.loc[1850], 'o', color='k', ms=7, zorder=5)
    if s == PAIR[0]:
        ax.text(1852, hist.loc[1850] - 16, f'{hist.loc[1850]:.0f} km$^3$ in 1850', fontsize=9)
    else:
        ax.text(1848, hist.loc[1850] + 10, f'{hist.loc[1850]:.0f} km$^3$ in 1850', fontsize=9,
                ha='right')
    ax.set_xlim(1800, T_END)
    ax.set_ylim(*ylim)
    ax.grid(alpha=0.3)
    which = 'smallest' if s == PAIR[0] else 'largest'
    ax.set_title(f'{lab} ens{s}: the {which} Little Ice Age state of the 197', loc='left',
                 fontsize=11)
axes[0, 0].set_ylabel('Regional glacier volume (km$^3$)')
axes[0, 0].legend(loc='lower left', fontsize=8.5, frameon=False)
axes[0, 1].set_yticklabels([])

# (c) both together
ax = axes[1, 0]
for s in PAIR:
    nat, hist = data[s]
    st = STYLE[s]
    q = nat.quantile([QLO, 0.5, QHI], axis=1).T
    ax.fill_between(q.index, q[QLO], q[QHI], color=st['nat'], alpha=0.25, lw=0)
    ax.plot(q.index, q[0.5], color=st['nat_dark'], lw=2, label=f'ens{s}: natural median, 16–84%')
    ax.plot(hist.loc[1820:1850], color=st['hist'], lw=1.2, ls=':')
    ax.plot(hist.loc[1850:], color=st['hist'], lw=2.2, label=f'ens{s}: historical')
ax.set_xlim(1820, T_END)
ax.set_ylim(*ylim)
ax.set_ylabel('Regional glacier volume (km$^3$)')
ax.set_title('(c) Both states: same end points, whatever the start\n(dotted: before 1850)', loc='left', fontsize=11)
ax.legend(loc='lower left', fontsize=8.5, frameon=False)
ax.grid(alpha=0.3)

# (d) the gap
ax = axes[1, 1]
for s in PAIR:
    nat, hist = data[s]
    st = STYLE[s]
    gap = nat.sub(hist.loc[1850:], axis=0)
    q = gap.quantile([QLO, 0.5, QHI], axis=1).T
    ax.fill_between(q.index, q[QLO], q[QHI], color=st['hist'], alpha=0.2, lw=0)
    ax.plot(q.index, q[0.5], color=st['hist'], lw=2.2, ls=st['ls'], zorder=st['z'],
            label=f'ens{s}: median, 16–84%')
    print(f'ens{s}: gap 1900 / 1950 / 2000 / 2015: ' +
          ' / '.join(f'{q.loc[y, 0.5]:.1f}' for y in [1900, 1950, 2000, 2015]))
ax.axhline(0, color='0.4', lw=0.8)
ax.set_xlim(1820, T_END)
ax.set_ylabel('natural − historical volume (km$^3$)')
ax.set_title('(d) The anthropogenic loss: the gap between natural and\nhistorical (the two states overlap)', loc='left', fontsize=11)
ax.legend(loc='upper left', fontsize=8.5, frameon=False)
ax.grid(alpha=0.3)

fig.suptitle('Why do the smallest and the largest Little Ice Age states have the same time '
             'of emergence?', y=0.97, fontsize=13)

os.makedirs('figs', exist_ok=True)
out = f'figs/explainer_two_states_{NAME}_ens{PAIR[0]}_ens{PAIR[1]}.png'
fig.savefig(out, dpi=150, bbox_inches='tight')
print('wrote', out)