"""One row per AdaPBS calibration in out/: posterior, Neff, target misfits.

For writing back into the gdirs and flagging bad fits (handover to-do 4).
Misfits are in units of the target sigma used in the run: z = (posterior
median prediction - obs) / obs_err. Plain numpy / xarray, no OGGM.

    python summarise.py [out_dir]     # default out/ -> <out_dir>/summary.csv (+ printed)
"""
import glob
import json
import os
import sys

import numpy as np
import pandas as pd
import xarray as xr

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.abspath(sys.argv[1]) if len(sys.argv) > 1 else f'{HERE}/out'
PARAMS = ['melt_f', 'prcp_fac', 'temp_bias']
R5 = f'{HERE}/../round5/output/{{}}_ys1500/dyn_calib_reg11.csv'   # .format(exp)
NAMES = {'RGI60-11.01450': 'Great Aletsch', 'RGI60-11.02822': 'Gorner',
         'RGI60-11.00746': 'Gepatsch', 'RGI60-11.01238': 'Rhone',
         'RGI60-11.02773': 'Findelen', 'RGI60-11.02072': 'Plaine Morte',
         'RGI60-11.00897': 'Hintereis', 'RGI60-11.03845': 'Pizzo Tambau',
         'RGI60-11.03429': "de l'Argentiere 1"}


def wq(x, w, q):
    """Weighted quantiles, ignoring NaN."""
    x, w = np.asarray(x, float), np.asarray(w, float)
    ok = np.isfinite(x) & (w > 0)
    x, w = x[ok], w[ok]
    o = np.argsort(x)
    cw = np.cumsum(w[o]) - 0.5 * w[o]
    return np.interp(np.atleast_1d(q), cw / w.sum(), x[o])


_r5 = {}
rows = []
for f in sorted(glob.glob(f'{OUT}/RGI60-*.nc')):
    ds = xr.open_dataset(f)
    rid = ds.attrs['rgi_id']
    exp = ds.attrs.get('exp', 'sw_pf3.5_mf5_lam12')
    tree = ds.attrs.get('tree', 'lmr_gdirs_v7_202610')
    if exp not in _r5:
        f5 = R5.format(exp)
        _r5[exp] = pd.read_csv(f5, index_col=0) if (tree.startswith('lmr_gdirs_v7')
                                                      and os.path.exists(f5)) else None
    r5 = _r5[exp]
    tag = os.path.basename(f)[len(rid) + 1:-3]
    w = ds.weight.values
    st = json.loads(ds.attrs['static'])
    r = {'tree': tree, 'exp': exp, 'rgi_id': rid, 'name': NAMES.get(rid, ''), 'tag': tag,
         'n_runs': len(w), 'n_iter': int(ds.iteration.max()) + 1,
         'failed': int((ds.error.values != '').sum()),
         'neff_final': float(ds.adapbs_neff.values[-1]),
         'run_time_median_s': float(np.nanmedian(ds.run_time.values)),
         'cpu_time_total_h': float(np.nansum(ds.run_time.values)) / 3600,
         'rgi_area_km2': float(ds.obs[1]),
         'area_k': float(ds.attrs.get('opt_area_k', 0) or 0),
         'terminus_step_frac': float(ds.attrs.get('terminus_step_km2', np.nan)) / float(ds.obs[1])}
    for p in PARAMS:
        lo, md, hi = wq(ds[p], w, [0.05, 0.5, 0.95])
        r[f'{p}_static'] = st[p]
        r[f'{p}'] = md
        r[f'{p}_p05'], r[f'{p}_p95'] = lo, hi
        r[f'{p}_mean'] = np.average(ds[p], weights=w)
        lo_b, hi_b = (ds.attrs['prior_min'][PARAMS.index(p)], ds.attrs['prior_max'][PARAMS.index(p)])
        # posterior mass within 2% of the prior range from a bound
        span = hi_b - lo_b
        r[f'{p}_at_bound'] = float(w[(ds[p] < lo_b + 0.02 * span) | (ds[p] > hi_b - 0.02 * span)].sum() / w.sum())
    r['melt_f_r5dyn'] = float(r5.loc[rid, 'melt_f_dynamic_calibration']) if r5 is not None else np.nan
    for k, t in enumerate(['dmdtda', 'area_km2', 'vol_km3']):
        md = wq(ds.pred[:, k], w, 0.5)[0]
        r[f'{t}_obs'] = float(ds.obs[k])
        r[f'{t}_err'] = float(ds.obs_err[k])
        r[f'{t}_post'] = md
        r[f'{t}_z'] = (md - float(ds.obs[k])) / float(ds.obs_err[k])
    r['area_1850_post'] = wq(ds.area_km2.sel(time=1850), w, 0.5)[0]
    r['vol_1850_post'] = wq(ds.volume_km3.sel(time=1850), w, 0.5)[0]
    r['flag'] = ';'.join(x for x, c in [('neff<50', r['neff_final'] < 50),
                                         ('failed>10%', r['failed'] > 0.1 * r['n_runs']),
                                         ('|z|>2', max(abs(r['dmdtda_z']), abs(r['area_km2_z']),
                                                       abs(r['vol_km3_z'])) > 2),
                                         ('at_bound', max(r[f'{p}_at_bound'] for p in PARAMS) > 0.2)]
                        if c)
    rows.append(r)
S = pd.DataFrame(rows)
S.to_csv(f'{OUT}/summary.csv', index=False)
pd.set_option('display.width', 250)
show = ['name', 'tag', 'n_runs', 'failed', 'neff_final',
        'melt_f_static', 'melt_f_r5dyn', 'melt_f', 'prcp_fac_static', 'prcp_fac', 'temp_bias_static',
        'temp_bias', 'dmdtda_z', 'area_km2_z', 'vol_km3_z', 'area_km2_obs', 'area_km2_post',
        'vol_km3_obs', 'vol_km3_post', 'area_1850_post', 'flag']
print(S[show].round(2).to_string(index=False) if len(S) < 40 else S[show].round(2).head(40).to_string(index=False))
if len(S) >= 10:
    good = S.flag == ''
    print(f'\n{len(S)} glaciers ({S.rgi_area_km2.sum():.0f} km2): {good.sum()} unflagged '
          f'({S.rgi_area_km2[good].sum() / S.rgi_area_km2.sum():.0%} of their area)')
    for fl in ['neff<50', 'failed>10%', '|z|>2', 'at_bound']:
        m = S.flag.str.contains(fl, regex=False)
        print(f'  {fl}: {m.sum()} glaciers, {S.rgi_area_km2[m].sum() / S.rgi_area_km2.sum():.0%} of the area')
    print(f'runs per glacier: median {S.n_runs.median():.0f}; run time median {S.run_time_median_s.median():.1f} s; '
          f'cpu per glacier median {S.cpu_time_total_h.median() * 60:.0f} min, total {S.cpu_time_total_h.sum():.0f} h')