"""Great Aletsch 1800-2025: how the AdaPBS particles shape the past (LIA) state.

From the regional v8 calibration (out_region_v8_pf35/RGI60-11.01450_*.nc;
every particle's area / volume / length, 1800-2025):
- prior: the particles of the first AdaPBS iteration (drawn from the prior
  around the static calibration), unweighted, 16-84% band + faint lines;
- posterior: all particles with their final weights, weighted 16-84% band;
- the top-weighted particle (the set written into the v8 L4 gdir);
- round5: the old calibration (v7 random precip, dyn melt_f only: 14.15).
Observations: RGI area 2003 and the inversion volume (both at the start of
2004, the calibration targets), Leclercq length changes (relative to 2003),
and the main LIA outline (data/inventories; for Aletsch Maisch et al. 2000,
which includes 10 glaciers that are separate in the RGI): as published, and minus the
tributaries' modelled LIA maxima (top particles), error bars 5.1% / 15.3%
(Reinthaler & Paul's regional 1-sigma).
Panel d: per particle, the LIA maximum area (1800-1875) against the 2003
area, coloured by weight.

    [RID=RGI60-11.00897] python plot_aletsch_particles.py
    -> figs/<name>_particles_v8.png (+ .csv); default Great Aletsch
"""
import os

import numpy as np
import pandas as pd
import xarray as xr
import netCDF4 as nc
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, LogNorm

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
RID = os.environ.get('RID', 'RGI60-11.01450')
NAMES = {'RGI60-11.01450': 'Great Aletsch', 'RGI60-11.00897': 'Hintereisferner'}
NAME = NAMES.get(RID, RID)
TAG = NAME.lower().replace(' ', '_')
# NOPRIOR=1: posterior only (no prior band / lines / markers), axes zoomed on
# the weighted particles; output <name>_particles_v8_noprior.png
NOPRIOR = os.environ.get('NOPRIOR', '0') == '1'
W_MIN = 1e-4          # particles below this weight are not drawn with NOPRIOR
if NOPRIOR:
    TAG += '_noprior'
TAG = TAG.replace('_noprior', '') + '_particles_v8' + ('_noprior' if NOPRIOR else '')
F = f'{HERE}/out_region_v8_pf35/{RID}_region_v8_pf35.nc'
R5 = f'{ROOT}/round5/output/sw_pf3.5_mf5_lam12_ys1500/hist_dyn_calib_reg11.nc'
TOP = f'{HERE}/out_region_v8_pf35_topparticle_lia.csv'
OUT = f'{HERE}/figs'
os.makedirs(OUT, exist_ok=True)
LIA_WIN = (1800, 1875)
ERR_A, ERR_V = 0.051, 0.153

INK, INK2, GRID, SURF = '#0b0b0b', '#52514e', '#e4e3df', '#fcfcfb'
BLUE, BLUE_D, ORANGE, PRIOR = '#2a78d6', '#184e8f', '#eb6834', '#a9a8a2'


def wq(x, w, q):
    """Weighted quantiles along axis 0 (particles), ignoring NaN."""
    out = np.full((len(q), x.shape[1]), np.nan)
    for j in range(x.shape[1]):
        v, ww = x[:, j], w
        ok = np.isfinite(v) & (ww > 0)
        if ok.sum() == 0:
            continue
        o = np.argsort(v[ok])
        cw = np.cumsum(ww[ok][o]) - 0.5 * ww[ok][o]
        out[:, j] = np.interp(q, cw / ww[ok].sum(), v[ok][o])
    return out


def leclercq(rid):
    from oggm import cfg
    from oggm.utils import get_demo_file
    cfg.initialize(logging_level='CRITICAL')
    links = pd.read_csv(get_demo_file('rgi_leclercq_links_2014_RGIV6.csv')).set_index('LID')
    d = nc.Dataset(get_demo_file('Glacier_Lengths_Leclercq.nc'))
    lids = d.variables['index'][:]
    for k, g in d.groups.items():
        lid = int(lids[int(k) - 1])
        if lid in links.index and links.loc[lid, 'RGI_ID'] == rid:
            s = pd.Series(np.asarray(g.variables['dL'][:], float),
                          index=np.asarray(g.variables['years'][:]).astype(int))
            d.close()
            return s
    d.close()
    return None


ds = xr.open_dataset(F).load()
t = ds.time.values
w = np.where(ds.error.values == '', np.nan_to_num(ds.weight.values), 0.)
w = w / w.sum()
itop = int(np.argmax(w))
prior = ds.iteration.values == ds.iteration.values.min()   # the first iteration: drawn from the prior
r5 = xr.open_dataset(R5).sel(rgi_id=RID).sel(time=slice(1800, 2025)).load()
r5t = r5.time.values.astype(int)
R5_MF = float(pd.read_csv(os.path.join(os.path.dirname(R5), 'dyn_calib_reg11.csv'), index_col=0)
              .loc[RID, 'melt_f_dynamic_calibration'])

# The LIA outline group and its tributaries' modelled LIA maxima
links = pd.read_csv(f'{ROOT}/data/inventories/lia_rgi6_links.csv')
lid = int(links[(links.RGIId == RID) & links.is_main_lia].LIA_ID.iloc[0])
row = pd.read_csv(f'{ROOT}/data/inventories/lia_outlines_summary.csv', index_col=0).loc[lid]
group = row.rgi_desc_ids.split()     # the RGI glaciers inside the outline
top = pd.read_csv(TOP, index_col=0)
others = [g for g in group if g != RID and g in top.index]
trib_a = float(top.loc[others, 'A_lia'].sum())
trib_v = float(top.loc[others, 'V_lia'].sum())
obs_lia_a, obs_lia_v = float(row.lia_area_km2), float(row.lia_vol_km3)
LIA_SRC = row.Citation

L = leclercq(RID)
sel = (t >= LIA_WIN[0]) & (t <= LIA_WIN[1])
amax = np.nanmax(ds.area_km2.values[:, sel], axis=1)
a03 = ds.pred.values[:, 1]

fig, axs = plt.subplots(2, 2, figsize=(12, 8.2), facecolor=SURF)
for a in axs.flat:
    a.set_facecolor(SURF)
    a.grid(color=GRID, lw=0.6)
    a.set_axisbelow(True)
    for s in ['top', 'right']:
        a.spines[s].set_visible(False)
    for s in ['left', 'bottom']:
        a.spines[s].set_color(INK2)
    a.tick_params(colors=INK2, labelsize=9)

panels = [('area_km2', 'Area (km²)', ds.obs.values[1], r5.area.values * 1e-6),
          ('volume_km3', 'Volume (km³)', ds.obs.values[2], r5.volume.values * 1e-9),
          ('length_m', 'Length change vs 2003 (km)', None, r5.length.values)]
rows = []
for a, (v, lab, ob, r5v) in zip([axs[0, 0], axs[0, 1], axs[1, 0]], panels):
    x = ds[v].values.astype(float)
    i04 = int(np.where(t == 2004)[0][0])
    if v == 'length_m':
        x = (x - x[:, [i04]]) * 1e-3
        r5v = (r5v - r5v[r5t == 2004][0]) * 1e-3
    pq = np.nanpercentile(x[prior], [16, 50, 84], axis=0)
    q = wq(x, w, [0.16, 0.5, 0.84])
    if NOPRIOR:
        # the weighted particles as faint lines, opacity by weight
        for k in np.where(w > W_MIN)[0]:
            a.plot(t, x[k], color=BLUE, lw=0.5, zorder=1,
                   alpha=float(np.clip(0.1 + 0.8 * w[k] / w.max(), 0.1, 0.9)))
    else:
        for k in np.where(prior)[0]:
            a.plot(t, x[k], color=PRIOR, lw=0.4, alpha=0.35, zorder=1)
        a.fill_between(t, pq[0], pq[2], color=PRIOR, alpha=0.35, lw=0, zorder=2,
                       label='prior particles, 16–84%')
    a.fill_between(t, q[0], q[2], color=BLUE, alpha=0.35, lw=0, zorder=3,
                   label='posterior (weighted), 16–84%')
    a.plot(t, x[itop], color=BLUE_D, lw=2, zorder=5, label='selected particle (in the gdir)')
    a.plot(r5t, r5v, color=ORANGE, lw=2, ls='--', zorder=4, label=f'round5 calibration (melt_f {R5_MF:.1f})')
    a.set_xlim(1800, 2025)
    a.set_ylabel(lab, color=INK2, fontsize=10)
    rows.append(pd.DataFrame({'year': t, 'var': v, 'prior_p16': pq[0], 'prior_p50': pq[1],
                              'prior_p84': pq[2], 'post_p16': q[0], 'post_p50': q[1],
                              'post_p84': q[2], 'top': x[itop]}))
    if ob is not None:
        a.errorbar([2004], [ob], yerr=[0.05 * ob] if v == 'area_km2' else [0.3 * ob],
                   fmt='o', color=INK, ms=6, capsize=3, zorder=6,
                   label='RGI 2003 (±5%)' if v == 'area_km2' else 'inversion 2004 (±30%)')
    # y range from the bulk of the prior (5-95%), so that a few extreme
    # particles do not squash the panel; the observations are kept in view
    if NOPRIOR:
        xw = x[w > W_MIN]
        lo95 = [np.nanmin(xw, axis=0), np.nanmax(xw, axis=0)]
    else:
        lo95 = np.nanpercentile(x[prior], [5, 95], axis=0)
    ext = [np.nanmin(lo95[0]), np.nanmax(lo95[1]), np.nanmin(r5v), np.nanmax(r5v)]
    if v == 'area_km2':
        ext += [obs_lia_a * (1 + ERR_A), ob]
    elif v == 'volume_km3':
        ext += [obs_lia_v * (1 + ERR_V), ob * 1.3]
    elif L is not None:
        l03 = np.interp(2003, L.index, L.values)
        ext += [np.nanmin(L.values - l03) * 1e-3, np.nanmax(L.values - l03) * 1e-3]
    pad = 0.08 * (max(ext) - min(ext))
    a.set_ylim(min(ext) - pad, max(ext) + pad)

# LIA outline
for a, og, tr, err in [(axs[0, 0], obs_lia_a, trib_a, ERR_A), (axs[0, 1], obs_lia_v, trib_v, ERR_V)]:
    a.errorbar([1850], [og], yerr=[err * og], fmt='s', mfc='none', mec=INK, ecolor=INK,
               ms=7, capsize=3, zorder=6, label=f'LIA outline ({LIA_SRC})' + (f' incl. {len(others)} tributaries' if others else ''))
    a.errorbar([1853], [og - tr], yerr=[err * og], fmt='s', color=INK, ms=7, capsize=3,
               zorder=6, label='LIA outline minus tributaries (modelled)') if others else None
# Leclercq
if L is not None:
    l03 = np.interp(2003, L.index, L.values)
    axs[1, 0].plot(L.index, (L.values - l03) * 1e-3, 'o', color=INK, ms=3.5, zorder=6,
                   label='Leclercq length record')
for a, ttl in zip([axs[0, 0], axs[0, 1], axs[1, 0]], ['a', 'b', 'c']):
    a.set_title(f'({ttl})', loc='left', color=INK, fontsize=11)
LEG = dict(fontsize=8, frameon=True, facecolor=SURF, edgecolor='none', framealpha=0.9)
axs[0, 0].legend(loc='best', **LEG)
axs[1, 0].legend(loc='best', handles=[h for h in axs[1, 0].get_legend_handles_labels()[0]][-1:],
                 **LEG)

# (d) LIA max area vs 2003 area, per particle
a = axs[1, 1]
cmap = LinearSegmentedColormap.from_list('b', ['#d6e4f5', BLUE, BLUE_D])
ok = np.isfinite(amax) & np.isfinite(a03)
if NOPRIOR:
    ok &= w > W_MIN
ws = np.where(w > 0, w, np.nan)
o = np.argsort(np.nan_to_num(ws[ok], nan=0))
sc = a.scatter(a03[ok][o], amax[ok][o], c=np.nan_to_num(ws[ok][o], nan=1e-6), s=14,
               cmap=cmap, norm=LogNorm(vmin=1e-4, vmax=max(w.max(), 1e-3)),
               edgecolor='none', zorder=3)
if not NOPRIOR:
    a.scatter(a03[prior], amax[prior], s=6, facecolor='none', edgecolor=PRIOR, lw=0.5, zorder=2,
              label='prior particles')
a.scatter([a03[itop]], [amax[itop]], s=70, facecolor='none', edgecolor=INK, lw=1.5, zorder=5,
          label='selected particle')
base = ok if NOPRIOR else prior
xa = np.nanpercentile(a03[base], [0, 100] if NOPRIOR else [2, 98]).tolist() + \
    [a03[itop], ds.obs.values[1] * 0.95, ds.obs.values[1] * 1.05]
ya = np.nanpercentile(amax[base], [0, 100] if NOPRIOR else [2, 98]).tolist() + \
    [amax[itop], (obs_lia_a - trib_a) * (1 - ERR_A), (obs_lia_a - trib_a) * (1 + ERR_A)]
a.set_xlim(min(xa) * 0.9, max(xa) * 1.05)
a.set_ylim(min(ya) * 0.9, max(ya) * 1.05)
a.axvspan(ds.obs.values[1] * 0.95, ds.obs.values[1] * 1.05, color=GRID, alpha=0.8, zorder=0)
a.axvline(ds.obs.values[1], color=INK, lw=1, zorder=1)
a.axhspan((obs_lia_a - trib_a) * (1 - ERR_A), (obs_lia_a - trib_a) * (1 + ERR_A),
          color=GRID, alpha=0.8, zorder=0)
a.axhline(obs_lia_a - trib_a, color=INK, lw=1, zorder=1)
a.text(ds.obs.values[1], a.get_ylim()[1], ' RGI 2003', va='top', fontsize=8, color=INK2)
a.text(a.get_xlim()[0], obs_lia_a - trib_a, ' LIA outline' + (' − tributaries' if others else ''), va='bottom',
       fontsize=8, color=INK2)
a.set_xlabel('Area at the RGI date (km²)', color=INK2, fontsize=10)
a.set_ylabel(f'LIA maximum area {LIA_WIN[0]}–{LIA_WIN[1]} (km²)', color=INK2, fontsize=10)
a.set_title('(d) each particle' + (f' (weight > {W_MIN:g})' if NOPRIOR else ''), loc='left', color=INK, fontsize=11)
cb = fig.colorbar(sc, ax=a, pad=0.01)
cb.set_label('posterior weight', color=INK2, fontsize=9)
cb.ax.tick_params(colors=INK2, labelsize=8)
a.legend(fontsize=8, frameon=False, loc='lower right')

fig.suptitle(f'{NAME} ({RID}): AdaPBS particles, v8 (LAPrec precipitation), '
             'HIST ens010' + (', posterior only' if NOPRIOR else ''), color=INK, fontsize=12, x=0.01, ha='left')
fig.tight_layout()
fig.savefig(f'{OUT}/{TAG}.png', dpi=150, facecolor=SURF)
pd.concat(rows).to_csv(f'{OUT}/{TAG}.csv', index=False)

# Numbers for the text
pa, pw = amax[prior], amax[ok]
print(f'particles {len(w)} (prior {prior.sum()}), top weight {w[itop]:.3f}')
print(f'LIA max area: prior 16/50/84 {np.nanpercentile(pa, [16, 50, 84]).round(1)}, '
      f'posterior {wq(amax[:, None], w, [0.16, 0.5, 0.84]).ravel().round(1)}, top {amax[itop]:.1f}, '
      f'round5 {np.nanmax(r5.area.values[(r5t >= 1800) & (r5t <= 1875)]) * 1e-6:.1f}')
print(f'LIA outline {lid} ({LIA_SRC}) {obs_lia_a:.1f} km2 / {obs_lia_v:.2f} km3; tributaries modelled '
      f'{trib_a:.1f} / {trib_v:.2f} -> {NAME} {obs_lia_a - trib_a:.1f} / {obs_lia_v - trib_v:.2f}')
print(f'2003 area: prior {np.nanpercentile(a03[prior], [16, 50, 84]).round(1)}, top {a03[itop]:.1f}, obs {ds.obs.values[1]:.1f}')
print('r(2003 area, LIA max area) over all particles:', round(np.corrcoef(a03[ok], amax[ok])[0, 1], 2))
if L is not None:
    print('Leclercq', L.index.min(), L.index.max(), 'dL 1850-2003 obs',
          round((np.interp(1850, L.index, L.values) - np.interp(2003, L.index, L.values)) * 1e-3, 2),
          'km; model top', round((ds.length_m.values[itop][t == 1850][0] -
                                  ds.length_m.values[itop][t == 2004][0]) * 1e-3, 2))
print('saved', f'{OUT}/{TAG}.png')