"""The _length_smb figure of conditional_temperature.py with the proposed smoothing.

- cumulative glacier quantities (regional volume, area, length): annual values, no bar;
- rates (glacier-weighted JJAS temperature, regional specific mass balance): 31-yr
  trailing mean (years t-30..t), the same on HIST and NAT;
- dashed: the current definition (15-yr trailing mean on everything).
Everything else as conditional_temperature.py: NAT conditioned on the 1850 state STATE,
anomalies to 1820-1850, temperature indexed like the glacier (year t = climate year t-1),
SMB = (V(t) - V(t-1)) / A(t-1) (900 kg m-3). ToE = SNR > 2 for good, to 2021.

    python conditional_temperature_newdef.py
    -> figs/attr_temperature_vs_glacier_conditional_<NAME>_<STATE>_newdef_length_smb.{png,csv}
"""
import os

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

NAME = os.environ.get('NAME', 'perm_a5v10')
STATE = os.environ.get('STATE', 'ens083')
W_RATE, W_OLD = 31, 15
BASELINE = (1820, 1850)
T_END = 2021
THRESH = [1, 1.65, 2, 2.58]
THRESH_C = ['#f5b041', '#e67e22', '#c0392b', '#7b241c']
YRS = np.arange(BASELINE[1] + 1, T_END + 1)
FILT = f'../perm/filtered/{NAME}'
RHO = 0.9


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


def delta(df, w):
    base = df.loc[BASELINE[0]:BASELINE[1]].mean()
    x = df if w is None else df.rolling(w, min_periods=w).mean()
    return x - base


def snr(nat, hist, w, sign):
    dn, dh = delta(nat, w).loc[YRS], delta(hist, w).loc[YRS]
    return sign * (dh - dn.median(axis=1)) / dn.std(axis=1)


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


# glacier: the state's own history before 1850, then HIST / every NAT climate
s_ = STATE[3:]
pre = read(f'{FILT}/pre1850_reg11.csv')
G = {}
for var in ['volume_km3', 'area_km2', 'length_km']:
    perm = read(f'{FILT}/{var}_perm_reg11.csv')
    p = pre[f'{var}_state{s_}']
    cols = [c for c in perm.columns if c.startswith(f'state{s_}_nat_')]
    G[var] = (pd.concat([pd.DataFrame({c: p for c in cols}), perm[cols].loc[BASELINE[1]:]]),
              pd.concat([p, perm[f'state{s_}_hist'].loc[BASELINE[1]:]]))
smb = [(G['volume_km3'][i].diff() / G['area_km2'][i].shift(1)) * 1e3 * RHO for i in (0, 1)]

# temperature: STATE's climate before 1850, each NAT climate from 1850
temp = read('../../climate_analysis/series/nat_jjas_glacw_newnat_r7.csv')
hist_t = temp.pop('hist')
nat_t = temp.copy()
nat_t.loc[:1849] = np.repeat(temp[STATE].loc[:1849].values[:, None], nat_t.shape[1], axis=1)
hist_t.loc[:1849] = temp[STATE].loc[:1849].values
nat_t.index, hist_t.index = nat_t.index + 1, hist_t.index + 1

# name: (nat, hist, sign, new window, color, lw)
series = {
    'JJAS temperature': (nat_t, hist_t, 1, W_RATE, '#e67e22', 1.8),
    'regional volume': (*G['volume_km3'], -1, None, '#c0392b', 2.6),
    'regional area': (*G['area_km2'], -1, None, '#641e16', 2.6),
    'regional length': (*G['length_km'], -1, None, '#1f618d', 2.6),
    'specific mass balance': (*smb, -1, W_RATE, '#27ae60', 2.2),
}
fig, ax = plt.subplots(figsize=(11, 5.4))
rows = []
for name, (nat, hist, sign, w, c, lw) in series.items():
    new, old = snr(nat, hist, w, sign), snr(nat, hist, W_OLD, sign)
    tn, to = toe(new), toe(old)
    lab = 'annual' if w is None else f'{w}-yr mean'
    ax.plot(YRS, old, color=c, lw=1, ls='--', alpha=0.6)
    ax.plot(YRS, new, color=c, lw=lw, label=f'{name}, {lab}: ToE {tn} (15-yr: {to})')
    rows.append(dict(series=name, smoothing=lab, toe=tn, toe_15yr=to,
                     **{f'snr_{y}': new.loc[y] for y in [1900, 1950, 1975, 2000, 2020, T_END]},
                     **{f'snr15_{y}': old.loc[y] for y in [1900, 1950, 1975, 2000, 2020, T_END]}))
ax.plot([], [], color='0.4', lw=1, ls='--', label='dashed: 15-yr trailing mean (current)')
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 (trailing windows end at t)')
ax.set_title(f'Forcing vs glacier, NAT conditioned on the 1850 state (HIST {STATE}): glacier '
             f'annual, temperature and SMB {W_RATE}-yr trailing mean; ToE = SNR > 2 for good',
             fontsize=10, loc='left')
ax.legend(loc='upper left', fontsize=8.5, frameon=False)
ax.grid(alpha=0.3)
os.makedirs('figs', exist_ok=True)
out = f'figs/attr_temperature_vs_glacier_conditional_{NAME}_{STATE}_newdef_length_smb'
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, 'display.max_columns', 20):
    print(tab.round(2).to_string())
print('wrote', out + '.png')