"""Check a round 6 dataset against the round5 ensemble (run on the login node).

    python check_perm.py <dataset dir>

- the twin runs (state i + NAT climate i) vs the round5 ensemble NAT run of
  member i: the same up to the restart noise (see
  lmr_gdirs_v7_202610/states/README.md);
- counts of merged files per state, and the regional volume of every state
  in 1850 and in 1900 (median over its NAT runs) and the HIST strand.
"""
import glob
import os
import re
import sys

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

data_dir = sys.argv[1]
ENS = '/home/www/fmaussion/runs/lmr_2025/round5/output_ens/ens_pf3.5_mf5_lam12'

files = sorted(glob.glob(f'{data_dir}/perm_state*_nat_ens*_reg11.nc'))
pairs = [tuple(re.search(r'state(\d{3})_nat_ens(\d{3})', f).groups()) for f in files]
states = sorted({s for s, _ in pairs})
print(f'{len(files)} permutation files, states {states}')

for s in states:
    f = f'{data_dir}/perm_state{s}_nat_ens{s}_reg11.nc'
    if not os.path.exists(f):
        continue
    run = xr.open_dataset(f)
    ens = xr.open_dataset(f'{ENS}/lmr_mira_nat_ens{s}_reg11.nc').volume.sel(
        rgi_id=run.rgi_id, time=run.time)
    v = run.volume
    ok = np.isfinite(v) & np.isfinite(ens)
    rel = (np.abs(v - ens) / np.maximum(np.abs(ens), 1.)).where(ok)
    reg = (v.where(ok).sum('rgi_id') - ens.where(ok).sum('rgi_id')) * 1e-9
    print(f'twin state{s} + NAT ens{s}: glaciers > 1e-4 rel: '
          f'{int((rel.max("time") > 1e-4).sum())} of {int(ok.any("time").sum())}, '
          f'max regional |diff| {float(np.abs(reg).max()):.4f} km3')

rows = []
for s in states:
    fs = [f for f, (ss, _) in zip(files, pairs) if ss == s]
    v = []
    for f in fs:
        with xr.open_dataset(f) as d:
            v.append(d.volume.sel(time=[1850, 1900]).sum('rgi_id').values * 1e-9)
    v = np.array(v)
    r = dict(state=s, n_nat=len(fs), V1850=v[:, 0].mean(),
             nat_V1900_median=np.median(v[:, 1]),
             nat_V1900_p5=np.percentile(v[:, 1], 5), nat_V1900_p95=np.percentile(v[:, 1], 95))
    h = f'{data_dir}/perm_state{s}_hist_reg11.nc'
    if os.path.exists(h):
        with xr.open_dataset(h) as d:
            r['hist_V1900'] = float(d.volume.sel(time=1900).sum()) * 1e-9
    rows.append(r)
print(pd.DataFrame(rows).set_index('state').round(2).to_string())