# "The glacier adds memory": three panels.
# (a) Memory in time: detection SNR of the JJAS temperature vs the OGGM
#     regional volume, for two contrasting HIST strands (as
#     temperature_members.py: 15-yr mean ending at t minus 1820-1850).
# (b) Memory of the LIA state: across the NAT members, the 1820-1850 mean
#     regional volume against the 1820-1850 mean JJAS temperature.
# (c) How long the memory is: correlation (across the NAT members) between
#     the 1820-1850 volume and the 1820-1850 mean of exponentially weighted
#     past temperature (first-order response, time scale tau), vs tau.
#
# Temperatures: the caches of nat_spread.py and temperature_members.py (run
# those first). The temperature is shifted by one year (OGGM volume in year t
# results from the climate up to t-1), as attribution_temperature.py.
# Run from round4/attribution/ with the oggm_env python. Env vars: EXP, WINDOW.
# -> figs/memory_<EXP>_w<W>.png, figs/memory_rtau_<EXP>.csv
import os

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

EXP = os.environ.get('EXP', 'lam8_jjas_dyn_newnat')
WINDOW = int(os.environ.get('WINDOW', 15))
FILTERED = f'../ensemble/filtered/{EXP}'
BASELINE = (1820, 1850)
T_END = 2021
SHOW = {'ens124': '#e6a09a', 'ens178': '#4a0f0a'}
MEMBERS = ['ens124', 'ens053', 'ens183', 'ens153', 'ens178']
TAUS = np.unique(np.round(np.logspace(0, np.log10(800), 40)).astype(int))
NAT_C = '#2e86c1'

nat_t = pd.read_csv(f'figs/nat_jjas_temp_glacw_{EXP}.csv', index_col=0).drop(columns='hist_ens183')
hist_t = pd.read_csv(f'figs/hist_jjas_temp_glacw_{EXP}.csv', index_col=0)
for df in (nat_t, hist_t):
    df.index = df.index.astype(int) + 1
vol_n = pd.read_csv(f'{FILTERED}/volume_km3_nat_ens_reg11.csv', index_col=0)
vol_h = pd.read_csv(f'{FILTERED}/volume_km3_hist_ens_reg11.csv', index_col=0)
for df in (vol_n, vol_h):
    df.index = df.index.astype(int)


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


def snr(n, h, sign):
    dn, dh = delta(n).loc[yrs], delta(h).loc[yrs]
    return sign * dh.sub(dn.median(axis=1), axis=0).div(dn.std(axis=1), axis=0)


def linear_response(df, tau):
    """x_f += (T - x_f) / tau from year 1 (the toy glacier of attribution_temperature.py)."""
    x = df.values.astype(float)
    out = np.empty_like(x)
    f = np.nanmean(x[:100], axis=0)
    for i in range(len(x)):
        f = f + (x[i] - f) / tau
        out[i] = f
    return pd.DataFrame(out, index=df.index, columns=df.columns)


yrs = np.arange(BASELINE[1] + 1, T_END + 1)
snr_t = snr(nat_t, hist_t, +1)
snr_v = snr(vol_n, vol_h[MEMBERS], -1)

# (b) and (c): baselines across the NAT members (all run through 1850)
pre = nat_t.loc[:BASELINE[1]]
vb = vol_n.loc[BASELINE[0]:BASELINE[1]].mean()
tb = pre.loc[BASELINE[0]:].mean()
r_tau = pd.Series({tau: np.corrcoef(linear_response(pre, tau).loc[BASELINE[0]:].mean(), vb)[0, 1]
                   for tau in TAUS}, name='r')
r_tau.index.name = 'tau'
r_tau.to_csv(f'figs/memory_rtau_{EXP}.csv')
tau_best = r_tau.idxmin()          # the most negative correlation (warm -> small)
r_plain = np.corrcoef(tb, vb)[0, 1]
print(f'r(baseline T, baseline V) = {r_plain:.2f} (R2 {r_plain**2:.2f}); '
      f'best tau {tau_best} yr: r = {r_tau[tau_best]:.2f} (R2 {r_tau[tau_best]**2:.2f})')
print(r_tau.iloc[::4].round(2).to_string())

# %% Figure
fig = plt.figure(figsize=(14, 9))
gs = fig.add_gridspec(2, 2, height_ratios=[1, 1.05], hspace=0.3, wspace=0.22)
ax = fig.add_subplot(gs[0, :])
for m, c in SHOW.items():
    ax.plot(yrs, snr_t[m], color=c, lw=1.2, ls='--', label=f'{m}: JJAS temperature')
    ax.plot(yrs, snr_v[m], color=c, lw=2.6, label=f'{m}: OGGM volume')
for thr, c in zip([1, 1.6, 2], ['#f5b041', '#e67e22', '#c0392b']):
    ax.axhline(thr, color=c, ls=':', lw=1)
ax.axhline(0, color='0.4', lw=0.8)
ax.axvspan(1965, 1985, color='0.85', zorder=0)
ax.text(1975, ax.get_ylim()[1] * 0.92, 'cool\n1965–85', ha='center', va='top', fontsize=8, color='0.35')
ax.set_xlim(yrs[0], yrs[-1])
ax.set_xlabel('Year t (end of the 15-yr window)' if WINDOW == 15 else f'Year t (end of the {WINDOW}-yr window)')
ax.set_ylabel('SNR (HIST − NAT median) / NAT std')
ax.set_title('(a) Memory in time: the glacier integrates the forcing', loc='left', fontweight='bold')
ax.legend(loc='upper left', fontsize=9, frameon=False, ncol=2)
ax.grid(alpha=0.3)

ax = fig.add_subplot(gs[1, 0])
ax.scatter(tb, vb, s=14, color=NAT_C, alpha=0.6, label='NAT members')
for m in MEMBERS:
    ax.scatter(tb[m], vb[m], s=70, color='#c0392b', edgecolor='k', lw=0.6, zorder=3)
    ax.annotate(m[3:], (tb[m], vb[m]), textcoords='offset points', xytext=(5, 4), fontsize=9)
ax.set_xlabel(f'{BASELINE[0]}–{BASELINE[1]} mean JJAS temperature (°C, glacier-weighted)')
ax.set_ylabel(f'{BASELINE[0]}–{BASELINE[1]} mean regional volume (km³)')
ax.set_title(f'(b) LIA state vs LIA temperature: r = {r_plain:.2f}', loc='left', fontweight='bold')
ax.grid(alpha=0.3)

ax = fig.add_subplot(gs[1, 1])
ax.plot(r_tau.index, -r_tau.values, color='k', lw=2)
ax.axvline(tau_best, color='0.5', ls='--', lw=1)
ax.plot(1, -r_plain, 'o', color='#c0392b', ms=7,
        label=f'{BASELINE[0]}–{BASELINE[1]} mean temperature only')
ax.set_xscale('log')
ax.set_xlabel('Memory time scale τ (yr)')
ax.set_ylabel('−r (baseline volume, past temperature)')
ax.set_title(f'(c) Best memory time scale: τ ≈ {tau_best} yr (r = {r_tau[tau_best]:.2f})',
             loc='left', fontweight='bold')
ax.legend(loc='lower right', fontsize=9, frameon=False)
ax.grid(alpha=0.3, which='both')
fig.savefig(f'figs/memory_{EXP}_w{WINDOW}.png', dpi=150, bbox_inches='tight')