# What does the glacier model add? The attribution SNR of attribution.py
# (15-yr window ending at t, minus 1820-1850; SNR = (HIST - NAT median) /
# NAT std, signed so that > 0 means "towards warming / ice loss") computed
# on:
# - the glacier-weighted JJAS temperature (forcing; cache of nat_spread.py);
# - the same temperature through a first-order linear response ("toy
#   glacier": x_f += (T - x_f) / tau, tau = 20 and 50 yr), i.e. a pure
#   low-pass filter with memory, no geometry, no state dependence;
# - OGGM regional volume and area.
# For each: the first year after which the SNR stays above 1 / 1.6 / 2.
#
# Run from round4/attribution/ with the oggm_env python, after nat_spread.py.
import os

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

EXP = os.environ.get('EXP', 'lam8_jjas_dyn')
MEMBER = os.environ.get('MEMBER', 'ens183')
WINDOW = int(os.environ.get('WINDOW', 15))
BASELINE = (1820, 1850)
T_END = 2021
TAUS = [20, 50]
FILTERED = f'../ensemble/filtered/{EXP}'
THRESH = [1, 1.6, 2]


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


def snr(d_nat, d_hist, sign):
    """sign = +1 if 'more' means anthropogenic (temperature), -1 for ice."""
    return sign * (d_hist - d_nat.median(axis=1)) / d_nat.std(axis=1)


def emergence(s, th):
    """First year after which s stays above th until the end."""
    below = s.index[~(s > th)]
    below = below[below < s.index[-1]]
    if len(below) == 0:
        return s.index[0]
    y = below[-1] + 1
    return y if y <= s.index[-1] and s.loc[y:].gt(th).all() else np.nan


def linear_response(df, tau):
    """First-order response to the full series from year 1 (spin-up included)."""
    x = df.values.astype(float)
    out = np.full_like(x, np.nan)
    f = np.nanmean(x[:100], axis=0)
    for i in range(len(x)):
        ok = np.isfinite(x[i])
        f = np.where(ok, f + (x[i] - f) / tau, np.nan)
        out[i] = f
    return pd.DataFrame(out, index=df.index, columns=df.columns)


temp = pd.read_csv('figs/nat_jjas_temp_glacw.csv', index_col=0)
# OGGM volume "in year t" is the state at the start of t, i.e. the result of
# the climate up to t-1: shift the temperature by one year to align them
temp.index = temp.index.astype(int) + 1
hist_t = temp.pop('hist_ens183')
nat_t = temp
if MEMBER != 'ens183':
    raise ValueError('the temperature cache only has HIST ens183')

series = {}
series['JJAS temperature'] = (nat_t, hist_t, +1, '#e67e22')
for tau, c in zip(TAUS, ['#af7ac5', '#6c3483']):
    both = linear_response(pd.concat([nat_t, hist_t.rename('h')], axis=1), tau)
    series[f'Toy glacier: T filtered, τ={tau} yr'] = (both[nat_t.columns], both['h'], +1, c)
for var, short, c in [('volume', 'km3', '#c0392b'), ('area', 'km2', '#641e16')]:
    nat = pd.read_csv(f'{FILTERED}/{var}_{short}_nat_ens_reg11.csv', index_col=0)
    hist = pd.read_csv(f'{FILTERED}/{var}_{short}_hist_ens_reg11.csv', index_col=0)[MEMBER]
    nat.index, hist.index = nat.index.astype(int), hist.index.astype(int)
    series[f'OGGM regional {var}'] = (nat, hist, -1, c)

yrs = np.arange(BASELINE[1] + 1, T_END + 1)
res, rows = {}, []
for name, (nat, hist, sign, c) in series.items():
    s = snr(delta(nat).loc[yrs], delta(hist.to_frame())[hist.name].loc[yrs], sign)
    res[name] = s
    rows.append(dict(series=name, **{f'emerge_snr>{th}': emergence(s, th) for th in THRESH},
                     **{f'snr_{y}': s.loc[y] for y in [1900, 1950, 1975, 2000, T_END]}))
tab = pd.DataFrame(rows).set_index('series')
tab.to_csv(f'figs/attr_temperature_vs_glacier_{EXP}_{MEMBER}_w{WINDOW}.csv')
with pd.option_context('display.width', 250, 'display.precision', 1):
    print(tab)

# Two versions on the same axes: with and without the toy glacier
ylim = None
for toy, suffix in ((True, ''), (False, '_notoy')):
    fig, ax = plt.subplots(figsize=(11, 5))
    for name, (nat, hist, sign, c) in series.items():
        if not toy and name.startswith('Toy'):
            continue
        lw, ls = (2.6, '-') if name.startswith('OGGM') else (1.8, '-' if 'temperature' in name else '--')
        ax.plot(yrs, res[name], color=c, lw=lw, ls=ls, label=name)
    for th, c in zip(THRESH, ['#f5b041', '#e67e22', '#c0392b']):
        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])
    if ylim is None:
        ylim = ax.get_ylim()
    ax.set_ylim(ylim)
    ax.set_ylabel('SNR  (HIST − NAT median) / NAT std')
    ax.set_xlabel('Year t (end of the 15-yr window)')
    ax.set_title(f'Same attribution, different quantities: forcing vs glacier ({EXP}, HIST {MEMBER}, '
                 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(f'figs/attr_temperature_vs_glacier_{EXP}_{MEMBER}_w{WINDOW}{suffix}.png', dpi=150,
                bbox_inches='tight')