"""Check / summarise a round 6 committed-loss dataset (login node).

    python check_commit.py <commit dataset dir> <perm dataset dir>

Per member and seed: the start volume == the HIST strand's volume in the
start year (perm_state###_hist_reg11.nc), and the regional volume after 50,
100 and 150 years, on the round5 common subset.
"""
import glob
import os
import re
import sys

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

cdir, pdir = sys.argv[1], sys.argv[2]
IDS = pd.read_csv(os.path.join(os.path.dirname(__file__), '..', 'round5', 'filtered_ens',
                               'ens_pf3.5_mf5_lam12', 'glacier_statistics_reg11.csv'),
                  index_col=0).index.values
rows = []
for f in sorted(glob.glob(f'{cdir}/commit*_state*_seed*_reg11.nc')):
    start, s, seed = re.search(r'commit(\d+)_state(\d{3})_seed(\d+)', f).groups()
    start = int(start)
    with xr.open_dataset(f) as ds:
        ids = np.intersect1d(ds.rgi_id.values, IDS)
        v = ds.volume.sel(rgi_id=ids)
        n_nan = int(v.isnull().any('time').sum())
        reg = v.sum('rgi_id') * 1e-9
    with xr.open_dataset(f'{pdir}/perm_state{s}_hist_reg11.nc') as h:
        h0 = float(h.volume.sel(rgi_id=ids, time=start).sum()) * 1e-9
    t0 = int(reg.time[0])
    r = dict(state=f'ens{s}', seed=int(seed), n_glaciers=len(ids), n_nan=n_nan,
             V_hist_start=h0, V_start=float(reg.sel(time=t0)))
    for dt in [50, 100, 150]:
        if t0 + dt <= int(reg.time[-1]):
            r[f'V+{dt}'] = float(reg.sel(time=t0 + dt))
    r['committed_loss_%'] = 100 * (1 - r.get('V+150', np.nan) / r['V_start'])
    rows.append(r)
df = pd.DataFrame(rows).set_index(['state', 'seed'])
print(df.round(3).to_string())
print('\nper state, mean (min-max over seeds) of V+150 (km3):')
g = df['V+150'].groupby('state')
print(pd.DataFrame({'mean': g.mean(), 'min': g.min(), 'max': g.max()}).round(2).to_string())