"""What AdaPBS did to the mass-balance parameters (melt_f, prcp_fac, temp_bias).

Per glacier (RGI 11-01, the 3817 glaciers with an AdaPBS calibration):
- v8 static: the informed three-step calibration of lmr_gdirs_v8_202610 (LAPrec precip),
  the prior of AdaPBS (the *_static columns of the L4 summaries);
- AdaPBS 2: the selected particle of the second round (area 5%, volume 10%,
  area_min_h), lmr_gdirs_v8_202610/adapbs/pf3.5_mf5_lam12_adapbs_a5v10;
- round5: the parameters behind round 6's attribution (v7 static, random precip, with the
  dynamic melt_f of round5/output/sw_pf3.5_mf5_lam12_ys1500), for reference.
Top: area-weighted distributions. Bottom: AdaPBS 2 minus v8 static against glacier area,
with the area-weighted median and 25-75% per size class.

    python params_compare.py -> figs/params_compare.png, figs/params_compare.csv
"""
import glob
import os

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
OUT = f'{HERE}/figs'
P = ['melt_f', 'prcp_fac', 'temp_bias']
LAB = {'melt_f': 'melt_f (kg m$^{-2}$ day$^{-1}$ K$^{-1}$)', 'prcp_fac': 'prcp_fac',
       'temp_bias': 'temp_bias (K)'}
INK, INK2, GRID, SURF = '#0b0b0b', '#52514e', '#e4e3df', '#fcfcfb'
C_STAT, C_ADA, C_R5 = '#8a8984', '#184e8f', '#eb6834'

S = pd.concat([pd.read_csv(f, index_col=0) for f in glob.glob(
    f'{ROOT}/lmr_gdirs_v8_202610/adapbs/pf3.5_mf5_lam12_adapbs_a5v10/summary/*.csv')])
S = S.loc[S.adapbs_status == 'top_particle']
v7 = pd.read_csv(f'{ROOT}/lmr_gdirs_v7_202610/exps/sw_pf3.5_mf5_lam12/RGI62/b_240/L3/summary/'
                 'glacier_statistics_11.csv', index_col=0, low_memory=False)
dyn = pd.read_csv(f'{ROOT}/round5/output/sw_pf3.5_mf5_lam12_ys1500/dyn_calib_reg11.csv',
                  index_col=0).melt_f_dynamic_calibration
R5 = v7.loc[S.index, P].copy()
R5['melt_f'] = dyn.reindex(S.index).fillna(R5['melt_f'])
a = S.rgi_area_km2


def wq(x, w, q):
    ok = np.isfinite(x) & np.isfinite(w)
    x, w = np.asarray(x)[ok], np.asarray(w)[ok]
    o = np.argsort(x)
    cw = np.cumsum(w[o]) - 0.5 * w[o]
    return np.interp(q, cw / w.sum(), x[o])


rows = {}
for name, d in [('round5 (v7 static + dyn melt_f)', R5),
                ('v8 static', S[[f'{p}_static' for p in P]].set_axis(P, axis=1)),
                ('AdaPBS 2', S[P])]:
    for p in P:
        rows[(name, p)] = {'area-weighted mean': np.average(d[p], weights=a),
                           'area-weighted median': wq(d[p], a, 0.5),
                           'median (by number)': d[p].median()}
T = pd.DataFrame(rows).T
T.to_csv(f'{OUT}/params_compare.csv')
pd.set_option('display.width', 200)
print(T.round(2).to_string())
for p in P:
    dlt = S[p] - S[f'{p}_static']
    print(f'{p}: AdaPBS 2 - v8 static, area-weighted median {wq(dlt, a, 0.5):+.2f}, '
          f'share of the area where it went down {a[dlt < 0].sum() / a.sum():.0%}')

fig, axs = plt.subplots(2, 3, figsize=(14, 7.5), facecolor=SURF)
bins_cls = [0.01, 0.1, 0.5, 1, 5, 20, 100]
for k, p in enumerate(P):
    ax, ab = axs[0, k], axs[1, k]
    for x in (ax, ab):
        x.set_facecolor(SURF)
        x.grid(color=GRID, lw=0.6)
        x.set_axisbelow(True)
        for sp in ['top', 'right']:
            x.spines[sp].set_visible(False)
        x.tick_params(colors=INK2, labelsize=9)
    st = S[f'{p}_static']
    lo, hi = np.nanpercentile(np.r_[st, S[p], R5[p]], [0.5, 99.5])
    bins = np.linspace(lo, hi, 41)
    ax.hist(st, bins=bins, weights=a, density=True, color=C_STAT, alpha=0.5,
            label=f'v8 static (median {wq(st, a, .5):.2f})')
    ax.hist(S[p], bins=bins, weights=a, density=True, histtype='step', color=C_ADA, lw=2,
            label=f'AdaPBS 2 ({wq(S[p], a, .5):.2f})')
    ax.hist(R5[p], bins=bins, weights=a, density=True, histtype='step', color=C_R5, lw=1.4,
            ls='--', label=f'round5, behind the 60% ({wq(R5[p], a, .5):.2f})')
    ax.set_xlabel(LAB[p], color=INK2)
    ax.set_ylabel('density (weighted by glacier area)' if k == 0 else '', color=INK2)
    ax.set_title(f'({"abc"[k]}) {p}', loc='left', color=INK, fontsize=10.5)
    ax.legend(fontsize=7.5, frameon=False)
    # change vs size
    dlt = S[p] - st
    ab.scatter(a, dlt, s=4, color=C_ADA, alpha=0.25, lw=0)
    cls = pd.cut(a, bins_cls)
    mid, med, q25, q75 = [], [], [], []
    for c, g in dlt.groupby(cls, observed=True):
        w = a.loc[g.index]
        mid.append(np.sqrt(c.left * c.right))
        med.append(wq(g, w, 0.5))
        q25.append(wq(g, w, 0.25))
        q75.append(wq(g, w, 0.75))
    ab.fill_between(mid, q25, q75, color=C_ADA, alpha=0.2, lw=0)
    ab.plot(mid, med, 'o-', color=C_ADA, lw=2, ms=5, label='area-weighted median, 25–75%')
    ab.axhline(0, color=INK, lw=0.8)
    ab.set_xscale('log')
    ab.set_xlabel('glacier area (km$^2$, RGI)', color=INK2)
    ab.set_ylabel(f'Δ{p}, AdaPBS 2 − v8 static', color=INK2)
    lim = np.nanpercentile(np.abs(dlt), 99)
    ab.set_ylim(-1.1 * lim, 1.1 * lim)
    ab.set_title(f'({"def"[k]}) change per glacier', loc='left', color=INK, fontsize=10.5)
    if k == 0:
        ab.legend(fontsize=7.5, frameon=False, loc='lower left')
fig.suptitle('Mass-balance parameters: v8 static calibration (the AdaPBS prior) vs AdaPBS round 2, '
             f'{len(S)} glaciers', color=INK, fontsize=12, x=0.01, ha='left')
fig.tight_layout()
fig.savefig(f'{OUT}/params_compare.png', dpi=150, facecolor=SURF)
print('saved', f'{OUT}/params_compare.png')