"""HISTALP vs Berkeley Earth: anomalies and rolling interannual std.

Top: anomaly vs 1961-1990 (thin: yearly, thick: centered 11-yr running
mean, full windows only).
Bottom: centered rolling std of the yearly values (needs the full window).

Usage: python plot_histalp_be_std.py [annual|jjas] [window]
       (defaults: annual 30)
"""
import os
import sys

import matplotlib.pyplot as plt
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
LABELS = {'berkeley-earth': 'Berkeley Earth', 'histalp': 'HISTALP'}
COLORS = {'berkeley-earth': '#eb6834', 'histalp': '#1baf7a'}

season = sys.argv[1] if len(sys.argv) > 1 else 'annual'
window = int(sys.argv[2]) if len(sys.argv) > 2 else 30

df = pd.read_csv(os.path.join(HERE, 'series',
                              f'tas_{season}_anom_1961-1990.csv'),
                 index_col=0)[list(LABELS)].dropna(how='all')
df = df.reindex(range(df.index[0], df.index[-1] + 1))

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), sharex=True,
                               gridspec_kw={'height_ratios': [3, 2]})
ax1.axhline(0, color='#52514e', lw=0.8)
ax1.axvspan(1961, 1990, color='#d0cfc8', alpha=0.35, lw=0)
for name in LABELS:
    s = df[name]
    c = COLORS[name]
    ok = s.dropna()
    lab = f'{LABELS[name]} ({ok.index[0]}–{ok.index[-1]})'
    ax1.plot(s.index, s, color=c, lw=0.7, alpha=0.45)
    ax1.plot(s.index, s.interpolate(limit=1, limit_area='inside')
             .rolling(11, center=True, min_periods=11).mean(),
             color=c, lw=2, label=lab)
    sd = s.rolling(window, center=True, min_periods=window).std()
    ax2.plot(sd.index, sd, color=c, lw=2, label=lab)

what = 'Annual' if season == 'annual' else season.upper()
ax1.set_ylabel('Anomaly vs 1961–1990 (K)')
ax1.set_title(f'{what} temperature, glacier-area weighted over the Alps '
              f'(RGI6.2 11-01). Thin: yearly, thick: 11-yr running mean',
              loc='left', fontsize=11)
ax2.set_ylabel(f'{window}-yr running std (K)')
ax2.set_title(f'Interannual variability: centered {window}-yr rolling '
              f'standard deviation of the yearly values (not detrended)',
              loc='left', fontsize=11)
for ax in (ax1, ax2):
    ax.grid(axis='y', color='#e6e5e0', lw=0.6)
    for sp in ['top', 'right']:
        ax.spines[sp].set_visible(False)
    ax.legend(frameon=False, loc='upper left', fontsize=9)
ax2.set_ylim(0, None)
ax2.set_xlim(df.index[0] - 2, df.index[-1] + 2)
fig.tight_layout()
os.makedirs(os.path.join(HERE, 'figs'), exist_ok=True)
f = os.path.join(HERE, 'figs', f'histalp_be_{season}_std{window}.png')
fig.savefig(f, dpi=150)
print(f)