# Round 7 copy of round6/attribution/regional_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).
# Conditional version of round5/attribution/regional_members.py: per 1850
# state, its HIST strand against its own NAT ensemble (the 197 NAT climates
# run from the same 1850 glaciers, round6). HIST and all NAT runs of a
# state share its pre-1850 history, hence the same 1820-1850 baseline.
#
# Per state and variable: figs/regional_{volume,area}_conditional_state###_<NAME>_w<W>.png
#   top: Δ (trailing W-yr mean minus 1820-1850) of HIST and the conditional
#        NAT (5-95%, 25-75%, median); dashed grey: the median of the round5
#        NAT ensemble (all states, each with its own baseline), i.e. the
#        counterfactual of the current method;
#   bottom: conditional SNR (solid) and current SNR (dashed; round5 HIST of
#        the same member vs all round5 NAT).
# Summary: figs/regional_snr_conditional_<NAME>_w<W>.png (all states, both
# methods), figs/regional_snr_conditional_only_<NAME>_w<W>.png (all states,
# conditional SNR only), figs/regional_snr_conditional_vol_area_state###_
# <NAME>_w<W>.png (one state, env ONE_STATE (default 010), volume and area
# SNR on one panel) and figs/snr_conditional_<NAME>_w<W>.csv (SNR at
# selected years + ToE), figs/snr_conditional_timeseries_<NAME>_w<W>.csv
# (all SNR series, conditional and current method).
#
# Run from round6/attribution/ with the oggm_env python. Needs
# ../preprocess_perm.py <NAME>. Env vars: NAME, ENS (round5 dataset), WINDOW.
import os

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

NAME = os.environ.get('NAME', 'perm_a5v10')
ENS = os.environ.get('ENS', 'ens_v8_adapbs_a5v10')
WINDOW = int(os.environ.get('WINDOW', 15))
BASELINE = (1820, 1850)
T_END = 2021
STATES = ['023', '007', '083', '169', '118']   # by 1850 volume
COLORS = plt.cm.Reds(np.linspace(0.35, 0.95, len(STATES)))
# IPCC likelihood terms (one-sided normal probabilities 68/90/95/99%)
SNR_THRESH = [(1, 'Likely (SNR>1)', '#f5b041'), (1.65, 'Very likely (SNR>1.65)', '#e67e22'),
              (2, 'Extremely likely (SNR>2)', '#c0392b'),
              (2.58, 'Virtually certain (SNR>2.58)', '#7b241c')]
NAT_C, NAT_DARK = '#2e86c1', '#1b4f72'
FILT6 = f'../perm/filtered/{NAME}'
FILT5 = f'../filtered_ens/{ENS}'
ONE_STATE = os.environ.get('ONE_STATE', '083')
YEARS_TAB = [1875, 1900, 1925, 1950, 1975, 2000, 2021]
os.makedirs('figs', exist_ok=True)


def read(f):
    df = pd.read_csv(f, index_col=0)
    df.index = df.index.astype(int)
    return df


def delta(df):
    """Trailing WINDOW-year mean minus the baseline mean (NaN for partial windows)."""
    return (df.rolling(WINDOW, min_periods=WINDOW).mean()
            - df.loc[BASELINE[0]:BASELINE[1]].mean())


def toe(snr, thr=2):
    """First year after which SNR stays > thr until the end (NaN if never)."""
    above = (snr > thr)[::-1].cummin()[::-1]
    return int(above.idxmax()) if above.any() else np.nan


rows = []
snr_one = {}
snr_all = {}
fig_c, axs_c = plt.subplots(2, 1, figsize=(10, 8), sharex=True,
                            gridspec_kw=dict(hspace=0.12))
fig_s, axs_s = plt.subplots(2, 1, figsize=(10, 8), sharex=True,
                            gridspec_kw=dict(hspace=0.12))
for iv, (var, short, unit) in enumerate([('volume', 'km3', 'km$^3$'),
                                         ('area', 'km2', 'km$^2$')]):
    perm = read(f'{FILT6}/{var}_{short}_perm_reg11.csv')
    nat5 = read(f'{FILT5}/{var}_{short}_nat_ens_reg11.csv')
    hist5 = read(f'{FILT5}/{var}_{short}_hist_ens_reg11.csv')
    v1850 = read(f'{FILT5}/volume_km3_nat_ens_reg11.csv').loc[1850]
    yrs = np.arange(BASELINE[1] + 1, T_END + 1)

    # Current method: round5 HIST member vs all round5 NAT (own baselines)
    dn5 = delta(nat5).loc[yrs]
    med5, sd5 = dn5.median(axis=1), dn5.std(axis=1)

    for s, c in zip(STATES, COLORS):
        m = f'ens{s}'
        pre = nat5[m].loc[:BASELINE[1] - 1]     # the state's history, all runs share it
        cols = [x for x in perm.columns if x.startswith(f'state{s}_nat_')]
        nat = pd.concat([pd.DataFrame({x: pre for x in cols}), perm[cols].loc[BASELINE[1]:]])
        hist = pd.concat([pre, perm[f'state{s}_hist'].loc[BASELINE[1]:]])
        dn, dh = delta(nat).loc[yrs], delta(hist).loc[yrs]
        med, sd = dn.median(axis=1), dn.std(axis=1)
        snr = -(dh - med) / sd
        snr5 = -(delta(hist5[m]).loc[yrs] - med5) / sd5
        pct = (v1850 < v1850[m]).mean() * 100

        fig, (ax, ax2) = plt.subplots(2, 1, figsize=(10, 7.5), sharex=True,
                                      gridspec_kw=dict(height_ratios=[3, 2], hspace=0.08))
        q = dn.quantile([0.05, 0.25, 0.75, 0.95], axis=1).T
        ax.fill_between(yrs, q[0.05], q[0.95], color=NAT_C, alpha=0.18, lw=0,
                        label=f'NAT from this 1850 state ({len(cols)} climates), 5–95%')
        ax.fill_between(yrs, q[0.25], q[0.75], color=NAT_C, alpha=0.30, lw=0,
                        label='NAT from this 1850 state, 25–75%')
        ax.plot(yrs, med, color=NAT_DARK, lw=2, label='NAT from this 1850 state, median')
        ax.plot(yrs, med5, color='0.45', lw=1.5, ls='--',
                label='NAT median, all 1850 states (current method)')
        ax.plot(yrs, dh, color=c, lw=2.2, label=f'HIST {m} (1850 volume: P{pct:.0f})')
        ax.axhline(0, color='0.4', lw=0.8)
        ax.set_ylabel(f'Δ{var} vs {BASELINE[0]}–{BASELINE[1]} ({unit})')
        ax.set_title(f'Regional {var}, Alps: HIST {m} vs NAT from the same 1850 state\n'
                     f'{WINDOW}-yr mean ending at t, minus {BASELINE[0]}–{BASELINE[1]} '
                     f'({NAME})', fontsize=11)
        ax.legend(loc='lower left', fontsize=8.5, frameon=False)
        ax.grid(alpha=0.3)

        ax2.plot(yrs, snr, color=c, lw=2.2, label='SNR, conditional (this state)')
        ax2.plot(yrs, snr5, color=c, lw=1.5, ls='--', label='SNR, current method')
        for th, lab, tc in SNR_THRESH:
            ax2.axhline(th, color=tc, ls=':', lw=1.3, label=lab)
        ax2.axhline(0, color='0.4', lw=0.8)
        ax2.set_ylabel('SNR (HIST − NAT median) / NAT std')
        ax2.set_xlabel('Year t')
        ax2.legend(loc='upper left', fontsize=8.5, frameon=False, ncol=2)
        ax2.grid(alpha=0.3)
        ax2.set_xlim(yrs[0], yrs[-1])
        fig.savefig(f'figs/regional_{var}_conditional_state{s}_{NAME}_w{WINDOW}.png',
                    dpi=150, bbox_inches='tight')
        plt.close(fig)

        snr_all[f'{var}_cond_{m}'] = snr
        snr_all[f'{var}_current_{m}'] = snr5
        if s == ONE_STATE:
            snr_one[var] = snr
        axs_c[iv].plot(yrs, snr, color=c, lw=1.6, label=f'{m} (P{pct:.0f})')
        a = axs_s[iv]
        a.plot(yrs, snr, color=c, lw=2, label=f'{m} (P{pct:.0f})')
        a.plot(yrs, snr5, color=c, lw=1, ls='--')

        r = dict(var=var, state=m, pct_1850=round(pct), toe_cond=toe(snr), toe_current=toe(snr5))
        for y in YEARS_TAB:
            r[f'snr_cond_{y}'] = snr.loc[y]
            r[f'snr_current_{y}'] = snr5.loc[y]
        rows.append(r)

    a = axs_c[iv]
    for th, lab, tc in SNR_THRESH:
        a.axhline(th, color=tc, ls=':', lw=1.2)
    a.axhline(0, color='0.4', lw=0.8)
    a.set_ylabel(f'SNR, regional {var}')
    a.grid(alpha=0.3)
    a.set_xlim(yrs[0], yrs[-1])
    if iv == 0:
        a.set_title(f'Conditional SNR per 1850 state (HIST vs NAT from the same '
                    f'1850 glaciers), {WINDOW}-yr window', fontsize=11)
        a.legend(loc='upper left', fontsize=8.5, frameon=False, ncol=2,
                 title='HIST strand (1850 volume pct.)', title_fontsize=8.5)

    a = axs_s[iv]
    for th, lab, tc in SNR_THRESH:
        a.axhline(th, color=tc, ls=':', lw=1.2)
    a.axhline(0, color='0.4', lw=0.8)
    a.set_ylabel(f'SNR, regional {var}')
    a.grid(alpha=0.3)
    a.set_xlim(yrs[0], yrs[-1])
    if iv == 0:
        a.set_title(f'SNR per 1850 state: conditional NAT (solid) vs current method '
                    f'(dashed), {WINDOW}-yr window', fontsize=11)
        a.legend(loc='upper left', fontsize=8.5, frameon=False, ncol=2,
                 title='HIST strand (1850 volume pct.)', title_fontsize=8.5)
axs_c[1].set_xlabel('Year t')

fig1, ax1 = plt.subplots(figsize=(10, 4.5))
for var, c in [('volume', '#922b21'), ('area', '#2874a6')]:
    ax1.plot(snr_one[var].index, snr_one[var], color=c, lw=2, label=f'regional {var}')
for th, lab, tc in SNR_THRESH:
    ax1.axhline(th, color=tc, ls=':', lw=1.3, label=lab)
ax1.axhline(0, color='0.4', lw=0.8)
ax1.set_xlim(snr_one['volume'].index[0], snr_one['volume'].index[-1])
ax1.set_ylabel('SNR (HIST − NAT median) / NAT std')
ax1.set_xlabel('Year t')
ax1.set_title(f'Conditional SNR, HIST ens{ONE_STATE} vs NAT from the same 1850 glaciers, '
              f'{WINDOW}-yr window', fontsize=11)
ax1.legend(loc='upper left', fontsize=9, frameon=False, ncol=2)
ax1.grid(alpha=0.3)
fig1.savefig(f'figs/regional_snr_conditional_vol_area_state{ONE_STATE}_{NAME}_w{WINDOW}.png',
             dpi=150, bbox_inches='tight')
fig_c.savefig(f'figs/regional_snr_conditional_only_{NAME}_w{WINDOW}.png', dpi=150,
              bbox_inches='tight')
axs_s[1].set_xlabel('Year t')
fig_s.savefig(f'figs/regional_snr_conditional_{NAME}_w{WINDOW}.png', dpi=150,
              bbox_inches='tight')

ts = pd.DataFrame(snr_all)
ts.index.name = 'year'
ts.round(4).to_csv(f'figs/snr_conditional_timeseries_{NAME}_w{WINDOW}.csv')
tab = pd.DataFrame(rows)
tab.round(2).to_csv(f'figs/snr_conditional_{NAME}_w{WINDOW}.csv', index=False)
with pd.option_context('display.width', 200):
    for var in ['volume', 'area']:
        t = tab[tab['var'] == var].set_index('state')
        print(f'\n{var}: SNR conditional / current, and ToE (SNR>2 for good)')
        out = pd.DataFrame({y: t[f'snr_cond_{y}'].round(2).astype(str) + ' / ' +
                            t[f'snr_current_{y}'].round(2).astype(str) for y in YEARS_TAB})
        out['ToE cond / current'] = t.toe_cond.astype('Int64').astype(str) + ' / ' + \
            t.toe_current.astype('Int64').astype(str)
        out.insert(0, 'P1850', t.pct_1850)
        print(out.to_string())