# Does the NAT ensemble spread increase with time? Spread across the NAT
# members of:
# - JJAS temperature (glacier-area weighted mean over the LMR 1° grid: each
#   cell weighted by the RGI area of the glaciers of the common subset whose
#   centre falls in it): the across-member std of the
#   31-yr running mean (low-frequency spread), and the interannual
#   variability (std of the annual anomalies from an 11-yr running mean,
#   in 31-yr windows), with HIST ens183 for comparison;
# - regional glacier volume and area (the CSVs of preprocess_mira.py).
# Vertical lines: 1850 and 1900, where NAT goes from the reconstruction to
# the GCM hist-nat runs (linear blend in between).
#
# Figure 2 (nat_anomaly_*): the 30-yr mean JJAS temperature (window ending
# at year t) as an anomaly from each run's own 1800-1850 mean, NAT 5-95%
# and median, with HIST ens183.
#
# Run from round4/attribution/ with the oggm_env python. The temperature
# series are cached in figs/nat_jjas_temp_glacw_<EXP>.csv (delete it to recompute).
# Env vars: EXP (default lam8_jjas_dyn), HIST_DIR, NAT_ROUND (the unstacked
# round with the NAT temperatures; default lmr_seas_fullsplit_oldBE_newnat if
# EXP ends in _newnat, else lmr_seas_fullsplit_oldBE).
import glob
import os
from concurrent.futures import ProcessPoolExecutor

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

EXP = os.environ.get('EXP', 'lam8_jjas_dyn')
HIST_DIR = os.environ.get('HIST_DIR', 'jjas_hists_v20261005')
UNSTACKED = '/home/www/fmaussion/runs/lmr_2025/data/unstacked'
FORCING = f'{UNSTACKED}/lmr_seas_fullsplit_oldBE'        # HIST temperature
NAT_ROUND = os.environ.get('NAT_ROUND', 'lmr_seas_fullsplit_oldBE_newnat'
                           if EXP.endswith('_newnat') else 'lmr_seas_fullsplit_oldBE')
NAT_FORCING = f'{UNSTACKED}/{NAT_ROUND}'
FILTERED = f'../ensemble/filtered/{EXP}'
X0, T_END = 1000, 2021
NAT_C, HIST_C = '#2e86c1', '#c0392b'


def glacier_weights(f):
    """RGI area of the common-subset glaciers per grid cell of file f."""
    st = pd.read_csv(f'{FILTERED}/glacier_statistics_reg11.csv', index_col=0)
    with xr.open_dataset(f, use_cftime=True) as ds:
        lat, lon = ds.lat.values, ds.lon.values
    w = xr.DataArray(np.zeros((len(lat), len(lon))), coords={'lat': lat, 'lon': lon},
                     dims=('lat', 'lon'))
    ilat = np.abs(st.cenlat.values[:, None] - lat).argmin(axis=1)
    ilon = np.abs(st.cenlon.values[:, None] - lon).argmin(axis=1)
    np.add.at(w.values, (ilat, ilon), st.rgi_area_km2.values)
    return w / w.sum()


def jjas(f, w):
    with xr.open_dataset(f, use_cftime=True) as ds:
        t = ds.tas.where(ds.time.dt.month.isin([6, 7, 8, 9]), drop=True)
        t = t.groupby('time.year').mean()
        return ((t * w).sum(['lat', 'lon']) - 273.15).to_series()


os.makedirs('figs', exist_ok=True)
cache = f'figs/nat_jjas_temp_glacw_{EXP}.csv'
if not os.path.exists(cache):
    files = sorted(glob.glob(f'{NAT_FORCING}/ensemble/tp_nat_monthly_ens*.nc'))
    files.append(f'{FORCING}/{HIST_DIR}/ensemble/tp_monthly_ens183.nc')
    print('NAT temperature from', NAT_FORCING, f'({len(files) - 1} members)')
    w = glacier_weights(files[0])
    print('Glacier-weighted cells:', int((w > 0).sum()), 'of', w.size)
    with ProcessPoolExecutor(int(os.environ.get('NJOBS', 8))) as ex:
        series = list(ex.map(jjas, files, [w] * len(files)))
    names = [os.path.basename(f)[-9:-3] for f in files[:-1]] + ['hist_ens183']
    pd.DataFrame(dict(zip(names, series))).to_csv(cache)
temp = pd.read_csv(cache, index_col=0)
hist_t = temp.pop('hist_ens183')
nat_t = temp


def iav(df, win=31):
    """Interannual std: anomalies from an 11-yr running mean, std in 31-yr windows."""
    anom = df - df.rolling(11, center=True).mean()
    return anom.rolling(win, center=True, min_periods=win).std()


lf_spread = nat_t.rolling(31, center=True, min_periods=31).mean().std(axis=1)
n_t = nat_t.notnull().sum(axis=1)
# Only where most members still run (avoid the last years with few members)
lf_spread = lf_spread.where(n_t >= 100)

fig, axes = plt.subplots(4, 1, figsize=(11, 11), sharex=True,
                         gridspec_kw=dict(hspace=0.12))
ax = axes[0]
ax.plot(lf_spread.index, lf_spread, color=NAT_C, lw=2)
ax.set_ylabel('K')
ax.set_title('JJAS temperature (glacier-weighted): spread across NAT members of the 31-yr running mean',
             fontsize=11, loc='left')

ax = axes[1]
nat_iav = iav(nat_t)
q = nat_iav.quantile([0.05, 0.5, 0.95], axis=1).T.where(n_t >= 100)
ax.fill_between(q.index, q[0.05], q[0.95], color=NAT_C, alpha=0.2, lw=0,
                label='NAT members, 5–95%')
ax.plot(q.index, q[0.5], color=NAT_C, lw=2, label='NAT median')
ax.plot(hist_t.index, iav(hist_t.to_frame())['hist_ens183'], color=HIST_C, lw=2,
        label='HIST ens183 (obs. from 1900)')
ax.set_ylabel('K')
ax.set_title('JJAS temperature (glacier-weighted): interannual variability (std of annual anomalies, '
             '31-yr windows)', fontsize=11, loc='left')
ax.legend(loc='upper left', fontsize=9, frameon=False)

for ax, (var, short, unit) in zip(axes[2:], [('volume', 'km3', 'km$^3$'),
                                              ('area', 'km2', 'km$^2$')]):
    nat = pd.read_csv(f'{FILTERED}/{var}_{short}_nat_ens_reg11.csv', index_col=0)
    nat.index = nat.index.astype(int)
    n = nat.notnull().sum(axis=1)
    sd = nat.std(axis=1)
    ax.plot(sd.index, sd.where(n >= 100), color=NAT_C, lw=2, label='all members available')
    full = nat.columns[nat.loc[T_END].notnull()]
    ax.plot(sd.index, nat[full].std(axis=1), color=NAT_C, lw=1, ls='--',
            label=f'only the {len(full)} members that reach {T_END}')
    ax.set_ylabel(unit)
    ax.set_title(f'Regional glacier {var}: std across NAT members', fontsize=11, loc='left')
axes[2].legend(loc='upper left', fontsize=9, frameon=False)

for ax in axes:
    for y in (1850, 1900):
        ax.axvline(y, color='0.4', lw=1, ls=':')
    ax.grid(alpha=0.3)
    ax.set_ylim(bottom=0)
axes[0].text(1852, axes[0].get_ylim()[1] * 0.92, 'blend\n1850–1900', fontsize=8, color='0.4')
axes[-1].set_xlim(X0, T_END)
axes[-1].set_xlabel('Year')
fig.suptitle(f'Does the NAT ensemble spread increase with time? ({EXP}, {nat_t.shape[1]} NAT members)',
             y=0.92)
fig.savefig(f'figs/nat_spread_{EXP}.png', dpi=150, bbox_inches='tight')

# Numbers
P = [(1601, 1700), (1701, 1800), (1801, 1850), (1901, 1950), (1951, 2000)]
for name, s in [('T low-freq spread (K)', lf_spread), ('T IAV NAT median (K)', q[0.5]),
                ('T IAV HIST (K)', iav(hist_t.to_frame())['hist_ens183'])]:
    print(f'{name:24s}', '  '.join(f'{a}-{b}: {s.loc[a:b].mean():.2f}' for a, b in P))

# %% Figure 2: 30-yr mean anomaly from 1800-1850
BASE = (1800, 1850)


def anom30(df):
    return (df.rolling(30, min_periods=30).mean()
            - df.loc[BASE[0]:BASE[1]].mean())


na = anom30(nat_t)
ha = anom30(hist_t.to_frame())['hist_ens183']
q = na.quantile([0.05, 0.5, 0.95], axis=1).T.where(n_t >= 100)
fig, ax = plt.subplots(figsize=(11, 4.2))
ax.fill_between(q.index, q[0.05], q[0.95], color=NAT_C, alpha=0.2, lw=0,
                label='NAT members, 5–95%')
ax.plot(q.index, q[0.5], color=NAT_C, lw=2, label='NAT median')
ax.plot(ha.index, ha, color=HIST_C, lw=2, label='HIST ens183 (obs. from 1900)')
ax.axhline(0, color='0.4', lw=0.8)
for y in (1850, 1900):
    ax.axvline(y, color='0.4', lw=1, ls=':')
ax.set_xlim(X0, ha.index.max())
ax.set_ylabel('K')
ax.set_xlabel('Year (end of the 30-yr window)')
ax.set_title(f'JJAS temperature (glacier-weighted): 30-yr mean, anomaly from each '
             f'run\'s {BASE[0]}–{BASE[1]} mean', fontsize=11, loc='left')
ax.legend(loc='upper left', fontsize=9, frameon=False)
ax.grid(alpha=0.3)
fig.savefig(f'figs/nat_anomaly_{EXP}.png', dpi=150, bbox_inches='tight')
for y in (1700, 1900, 1950, 2000, 2020):
    r = q.loc[y] if y in q.index else None
    print(f'{y}: HIST {ha.get(y, np.nan):+.2f}  NAT median {r[0.5]:+.2f} '
          f'[{r[0.05]:+.2f}, {r[0.95]:+.2f}]' if r is not None and np.isfinite(r[0.5])
          else f'{y}: HIST {ha.get(y, np.nan):+.2f}  NAT n/a')