"""Checks and figure for data/unstacked/be_lmr21_v<VERSION>/tp_monthly.nc.

Glacier-area weighted (RGI6.2 11-01 glaciers, nearest 1° cell) regional
series of the new forcing. Check against the oldBE HIST ens183 JJAS
variant, which puts BE's JJAS anomaly into all 12 months from 1900 on, so
its JJAS means must equal ours. (The oldBE *monthly* HIST variant is BE
shifted by one year: its year Y holds BE year Y-1, so it is not used.)
Figure: figs/be_lmr21_forcing.png.

Usage: python check_be_lmr_forcing.py [OFFSET_PERIOD]   (e.g. 1750-1779)
"""
import os
import sys

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

from build_be_lmr_forcing import out_dir, OLDBE, SWITCH, TC

PERIOD = None
if len(sys.argv) > 1 and sys.argv[1] != 'none':
    PERIOD = tuple(int(y) for y in sys.argv[1].split('-'))
OUT = out_dir(PERIOD)

HERE = os.path.dirname(os.path.abspath(__file__))
gl = pd.read_csv(os.path.join(HERE, 'glaciers_rgi62_1101.csv'), index_col=0)


def regional(f):
    """Glacier-area weighted monthly series (°C) as a pandas Series."""
    with xr.open_dataset(f, decode_times=TC) as ds:
        lat, lon = ds.lat.values, ds.lon.values
        i = np.abs(gl.CenLat.values[:, None] - lat[None]).argmin(1)
        j = np.abs(gl.CenLon.values[:, None] - lon[None]).argmin(1)
        w = np.zeros((len(lat), len(lon)))
        np.add.at(w, (i, j), gl.Area.values)
        w = xr.DataArray(w / w.sum(), dims=('lat', 'lon'))
        s = (ds.tas * w).sum(('lat', 'lon')).values - 273.15
        yrs, mon = ds['time.year'].values, ds['time.month'].values
    return pd.Series(s, index=pd.MultiIndex.from_arrays([yrs, mon]))


def seas(s, months):
    s = s[s.index.get_level_values(1).isin(months)]
    return s.groupby(level=0).mean()


def hp(s, w=31):
    return s - s.rolling(w, center=True, min_periods=w // 2 + 1).mean()


new = regional(os.path.join(OUT, 'tp_monthly.nc'))
old = regional(os.path.join(OLDBE, 'jjas_hists_v20261005/ensemble/'
                                   'tp_monthly_ens183.nc'))

ann, jjas = seas(new, range(1, 13)), seas(new, [6, 7, 8, 9])
ann_o, jjas_o = seas(old, range(1, 13)), seas(old, [6, 7, 8, 9])

print('1961-1990 mean, new / old (annual, JJAS): %.3f / %.3f, %.3f / %.3f'
      % (ann.loc[1961:1990].mean(), ann_o.loc[1961:1990].mean(),
         jjas.loc[1961:1990].mean(), jjas_o.loc[1961:1990].mean()))
d = (jjas - jjas_o).loc[1900:]
print('JJAS new - oldBE JJAS HIST, 1900-2024: mean %.4f, max |diff| %.4f K'
      % (d.mean(), d.abs().max()))

print('\nstd of deviations from the 31-yr running mean (K):')
rows = []
for p in [(1, 500), (500, 1000), (1000, 1500), (1500, 1749), (1750, 1849),
          (1850, 1949), (1950, 2024)]:
    r = {'period': f'{p[0]}-{p[1]}'}
    r['annual'] = hp(ann).loc[p[0]:p[1]].std()
    r['jjas'] = hp(jjas).loc[p[0]:p[1]].std()
    m = new.loc[p[0]:p[1]]
    m = m - m.groupby(level=1).transform('mean')
    r['monthly'] = m.std()
    rows.append(r)
print(pd.DataFrame(rows).set_index('period').round(2))

lo = ann.rolling(31, center=True).mean()
print('\n31-yr mean annual T around the switch:',
      lo.loc[[1700, 1720, 1734, 1765, 1780, 1800]].round(2).to_dict())

# Figure
anom = lambda s: s - s.loc[1961:1990].mean()
fig, axs = plt.subplots(3, 1, figsize=(12, 10), sharex=True,
                        gridspec_kw={'height_ratios': [3, 2, 2]})
for s, so, lab, ax in [(ann, ann_o, 'Annual', axs[0])]:
    a = anom(s)
    ax.plot(a.index, a, color='#2a78d6', lw=0.4, alpha=0.5)
    ax.plot(a.index, a.rolling(31, center=True).mean(), color='#2a78d6',
            lw=2, label='31-yr running mean')
    ax.set_ylabel('Annual anomaly vs 1961–1990 (K)')
for s, c, lab in [(ann, '#2a78d6', 'annual'), (jjas, '#eb6834', 'JJAS')]:
    axs[1].plot(s.index, hp(s).rolling(30, center=True).std(), color=c,
                lw=1.5, label=lab)
axs[1].set_ylabel('30-yr std of high-pass (K)')
mm = new - new.groupby(level=1).transform('mean')
mstd = mm.groupby(level=0).std()
axs[2].plot(mstd.index, mstd.rolling(30, center=True).mean(),
            color='#1baf7a', lw=1.5)
axs[2].set_ylabel('Monthly anomaly std\nwithin year, 30-yr mean (K)')
for ax in axs:
    ax.axvline(SWITCH, color='#52514e', lw=0.8, ls=':')
    ax.grid(axis='y', color='#e6e5e0', lw=0.6)
    for sp in ['top', 'right']:
        ax.spines[sp].set_visible(False)
axs[0].legend(frameon=False, loc='upper left', fontsize=9)
axs[1].legend(frameon=False, loc='lower left', fontsize=9, ncol=2)
axs[0].set_title('BE (1750–2024) + inflated LMR v2.1 (1–1749), '
                 'glacier-area weighted over the Alps', loc='left',
                 fontsize=11)
axs[1].set_ylim(0, None)
axs[2].set_ylim(0, None)
axs[2].set_xlim(0, 2030)
fig.tight_layout()
f = os.path.join(HERE, 'figs', os.path.basename(OUT) + '.png')
fig.savefig(f, dpi=150)
print(f)