# Round 7 copy of round6/attribution/conditional_temperature.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/attribution_temperature.py
# (the "_notoy" figure): the SNR of the glacier-weighted JJAS temperature
# next to the SNR of the OGGM regional volume and area, with the
# counterfactual conditioned on the 1850 state.
#
# For temperature, the conditional NAT of state i is i's own climate before
# 1850 and NAT climate j from 1850, so all 197 counterfactuals share i's
# 1820-1850 baseline and it cancels: the conditional temperature SNR is the
# same for every state by construction (computed for the 7 states as a
# check). For the glacier it collapses onto one curve as a result (round6).
#
# Inputs: round5/attribution/figs/nat_jjas_temp_glacw_<ENS>.csv (197 NAT +
# HIST ens010, glacier-weighted JJAS, cache of nat_spread.py) and
# figs/snr_conditional_timeseries_<NAME>_w<W>.csv (regional_conditional.py).
# Output: figs/attr_temperature_vs_glacier_conditional_<NAME>_<STATE>_w<W>.png
# and .csv (with the current method dashed), and ..._length.png: conditional
# only, with the regional length (sum of the glacier lengths; conditional SNR
# computed here from ../filtered/<NAME>/length_km_perm_reg11.csv and
# pre1850_reg11.csv, as regional_conditional.py does for volume and area).
# ..._length_smb.png: the same plus the regional specific mass balance
# (SMB = (V(t) - V(t-1)) / A(t-1), m w.e. yr-1 with 900 kg m-3, indexed like
# the shifted temperature: year t = climate year t-1), conditional SNR
# computed here from the volume and area CSVs.
# Env: NAME, ENS, WINDOW, STATE (default ens010).
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))
STATE = os.environ.get('STATE', 'ens083')
STATES = ['ens023', 'ens007', 'ens083', 'ens169', 'ens118']
BASELINE = (1820, 1850)
T_END = 2021
# IPCC likelihood terms: likely / very likely / extremely likely / virtually certain
THRESH = [1, 1.65, 2, 2.58]
THRESH_C = ['#f5b041', '#e67e22', '#c0392b', '#7b241c']

temp = pd.read_csv('../../climate_analysis/series/nat_jjas_glacw_newnat_r7.csv',
                   index_col=0)
temp.index = temp.index.astype(int)
hist_t = temp.pop('hist')     # identical for all strands from 1850
nat_t = temp


def delta(df):
    base = df.loc[BASELINE[0]:BASELINE[1]].mean()
    return df.rolling(WINDOW, min_periods=WINDOW).mean() - base


def shift(x):
    """OGGM volume in year t = climate up to t-1: align as round5 did."""
    x = x.copy()
    x.index = x.index + 1
    return x


def toe(s, th=2):
    above = (s > th)[::-1].cummin()[::-1]
    return int(above.idxmax()) if above.any() else np.nan


yrs = np.arange(BASELINE[1] + 1, T_END + 1)
snr_t = {}
for st in STATES:
    # state st's history before 1850, each NAT climate from 1850 (climate years)
    nat = nat_t.copy()
    nat.loc[:1849] = np.repeat(nat_t[st].loc[:1849].values[:, None], nat.shape[1], axis=1)
    hist = hist_t.copy()
    hist.loc[:1849] = nat_t[st].loc[:1849].values
    dn, dh = delta(shift(nat)).loc[yrs], delta(shift(hist)).loc[yrs]
    snr_t[st] = (dh - dn.median(axis=1)) / dn.std(axis=1)
snr_t = pd.DataFrame(snr_t)
spread = (snr_t.max(axis=1) - snr_t.min(axis=1))
print(f'temperature SNR, max spread across the 7 states: {spread.max():.3f} '
      f'(year {spread.idxmax()}); from 1866 on: {spread.loc[1866:].max():.2e}')

# Current method (round5): HIST ens010 vs all NAT, own baselines
dn5 = delta(shift(nat_t)).loc[yrs]
cur_t = (delta(shift(hist_t.where(hist_t.index >= 1850, nat_t[STATE]))).loc[yrs]
         - dn5.median(axis=1)) / dn5.std(axis=1)

g = pd.read_csv(f'figs/snr_conditional_timeseries_{NAME}_w{WINDOW}.csv', index_col=0)
series = {
    'JJAS temperature': (snr_t[STATE], cur_t, '#e67e22', 1.8),
    'OGGM regional volume': (g[f'volume_cond_{STATE}'], g[f'volume_current_{STATE}'], '#c0392b', 2.6),
    'OGGM regional area': (g[f'area_cond_{STATE}'], g[f'area_current_{STATE}'], '#641e16', 2.6),
}

fig, ax = plt.subplots(figsize=(11, 5))
rows = []
for name, (cond, cur, c, lw) in series.items():
    ax.plot(yrs, cond.loc[yrs], color=c, lw=lw, label=f'{name}, conditional')
    ax.plot(yrs, cur.loc[yrs], color=c, lw=1, ls='--', alpha=0.6,
            label=f'{name}, current method')
    rows.append(dict(series=name, toe_cond=toe(cond.loc[yrs]), toe_current=toe(cur.loc[yrs]),
                     **{f'snr_cond_{y}': cond.loc[y] for y in [1900, 1950, 1975, 2000, T_END]},
                     **{f'snr_current_{y}': cur.loc[y] for y in [1900, 1950, 1975, 2000, T_END]}))
for th, c in zip(THRESH, THRESH_C):
    ax.axhline(th, color=c, lw=1, ls=':')
ax.axhline(0, color='0.4', lw=0.8)
ax.axvline(1900, color='0.4', lw=1, ls=':')
ax.set_xlim(yrs[0], yrs[-1])
ax.set_ylabel('SNR  (HIST − NAT median) / NAT std')
ax.set_xlabel('Year t (end of the 15-yr window)')
ax.set_title(f'Forcing vs glacier, NAT conditioned on the 1850 state (HIST {STATE}, '
             f'{WINDOW}-yr mean minus {BASELINE[0]}–{BASELINE[1]})', fontsize=10.5, loc='left')
ax.legend(loc='upper left', fontsize=8.5, frameon=False, ncol=2)
ax.grid(alpha=0.3)
os.makedirs('figs', exist_ok=True)
out = f'figs/attr_temperature_vs_glacier_conditional_{NAME}_{STATE}_w{WINDOW}'
fig.savefig(out + '.png', dpi=150, bbox_inches='tight')
tab = pd.DataFrame(rows).set_index('series')
tab.round(2).to_csv(out + '.csv')
with pd.option_context('display.width', 250):
    print(tab.round(2).to_string())


# --- Conditional only, with the regional length
def cond_glacier(var):
    """Conditional SNR of a regional glacier variable for STATE (as regional_conditional.py)."""
    perm = pd.read_csv(f'../perm/filtered/{NAME}/{var}_perm_reg11.csv', index_col=0)
    pre = pd.read_csv(f'../perm/filtered/{NAME}/pre1850_reg11.csv', index_col=0)[f'{var}_state{STATE[3:]}']
    perm.index, pre.index = perm.index.astype(int), pre.index.astype(int)
    cols = [c for c in perm.columns if c.startswith(f'state{STATE[3:]}_nat_')]
    nat = pd.concat([pd.DataFrame({c: pre for c in cols}), perm[cols].loc[BASELINE[1]:]])
    hist = pd.concat([pre, perm[f'state{STATE[3:]}_hist'].loc[BASELINE[1]:]])
    dn, dh = delta(nat).loc[yrs], delta(hist).loc[yrs]
    return -(dh - dn.median(axis=1)) / dn.std(axis=1)


snr_len = cond_glacier('length_km')
chk = (cond_glacier('volume_km3') - g[f'volume_cond_{STATE}'].loc[yrs]).abs().max()
print(f'check: volume SNR from pre1850_reg11.csv vs regional_conditional.py: max |diff| {chk:.1e}')
series2 = {
    'JJAS temperature': (snr_t[STATE], '#e67e22', 1.8),
    'OGGM regional volume': (g[f'volume_cond_{STATE}'], '#c0392b', 2.6),
    'OGGM regional area': (g[f'area_cond_{STATE}'], '#641e16', 2.6),
    'OGGM regional length': (snr_len, '#1f618d', 2.6),
}
fig, ax = plt.subplots(figsize=(11, 5))
rows = []
for name, (cond, c, lw) in series2.items():
    ax.plot(yrs, cond.loc[yrs], color=c, lw=lw, label=name)
    rows.append(dict(series=name, toe_cond=toe(cond.loc[yrs]),
                     **{f'snr_{y}': cond.loc[y] for y in [1900, 1950, 1975, 2000, T_END]}))
for th, c in zip(THRESH, THRESH_C):
    ax.axhline(th, color=c, lw=1, ls=':')
ax.axhline(0, color='0.4', lw=0.8)
ax.axvline(1900, color='0.4', lw=1, ls=':')
ax.set_xlim(yrs[0], yrs[-1])
ax.set_ylabel('SNR  (HIST − NAT median) / NAT std')
ax.set_xlabel('Year t (end of the 15-yr window)')
ax.set_title(f'Forcing vs glacier, NAT conditioned on the 1850 state (HIST {STATE}, '
             f'{WINDOW}-yr mean minus {BASELINE[0]}–{BASELINE[1]})', fontsize=10.5, loc='left')
ax.legend(loc='upper left', fontsize=9, frameon=False)
ax.grid(alpha=0.3)
fig.savefig(out + '_length.png', dpi=150, bbox_inches='tight')
tab2 = pd.DataFrame(rows).set_index('series')
tab2.round(2).to_csv(out + '_length.csv')
print(tab2.round(2).to_string())


# --- Variant with the specific mass balance
RHO = 0.9   # ice -> w.e.


def smb_series(vol, area):
    """(V(t) - V(t-1)) / A(t-1) in m w.e. yr-1, at index t (vol km3, area km2)."""
    return (vol.diff() / area.shift(1)) * 1e3 * RHO


def cond_smb():
    s_ = STATE[3:]
    vols, areas = {}, {}
    for var, d in [('volume_km3', vols), ('area_km2', areas)]:
        perm = pd.read_csv(f'../perm/filtered/{NAME}/{var}_perm_reg11.csv', index_col=0)
        pre = pd.read_csv(f'../perm/filtered/{NAME}/pre1850_reg11.csv', index_col=0)[f'{var}_state{s_}']
        perm.index, pre.index = perm.index.astype(int), pre.index.astype(int)
        cols = [c for c in perm.columns if c.startswith(f'state{s_}_nat_')]
        d['nat'] = pd.concat([pd.DataFrame({c: pre for c in cols}), perm[cols].loc[BASELINE[1]:]])
        d['hist'] = pd.concat([pre, perm[f'state{s_}_hist'].loc[BASELINE[1]:]])
    nat = smb_series(vols['nat'], areas['nat'])
    hist = smb_series(vols['hist'], areas['hist'])
    dn, dh = delta(nat).loc[yrs], delta(hist).loc[yrs]
    return -(dh - dn.median(axis=1)) / dn.std(axis=1), hist


snr_smb, smb_hist = cond_smb()
series3 = dict(series2)
series3['OGGM specific mass balance'] = (snr_smb, '#27ae60', 2.2)
fig, ax = plt.subplots(figsize=(11, 5))
rows = []
for name, (cond, c, lw) in series3.items():
    ax.plot(yrs, cond.loc[yrs], color=c, lw=lw, label=name)
    rows.append(dict(series=name, toe_cond=toe(cond.loc[yrs]),
                     **{f'snr_{y}': cond.loc[y] for y in [1900, 1950, 1975, 2000, T_END]}))
for th, c in zip(THRESH, THRESH_C):
    ax.axhline(th, color=c, lw=1, ls=':')
ax.axhline(0, color='0.4', lw=0.8)
ax.axvline(1900, color='0.4', lw=1, ls=':')
ax.set_xlim(yrs[0], yrs[-1])
ax.set_ylabel('SNR  (HIST − NAT median) / NAT std')
ax.set_xlabel('Year t (end of the 15-yr window)')
ax.set_title(f'Forcing vs glacier, NAT conditioned on the 1850 state (HIST {STATE}, '
             f'{WINDOW}-yr mean minus {BASELINE[0]}–{BASELINE[1]})', fontsize=10.5, loc='left')
ax.legend(loc='upper left', fontsize=9, frameon=False)
ax.grid(alpha=0.3)
fig.savefig(out + '_length_smb.png', dpi=150, bbox_inches='tight')
tab3 = pd.DataFrame(rows).set_index('series')
tab3.round(2).to_csv(out + '_length_smb.csv')
print(tab3.round(2).to_string())
ok = yrs[yrs >= 1866]
print('r(SNR temperature, SNR x), 1866-2021:',
      {k: round(np.corrcoef(snr_t[STATE].loc[ok], v[0].loc[ok])[0, 1], 2) for k, v in series3.items()
       if k != 'JJAS temperature'})
print('HIST regional SMB (m w.e./yr), 15-yr means ending', {y: round(smb_hist.rolling(15).mean().loc[y], 2)
                                                         for y in [1900, 1950, 2000, 2021]})