"""Load / validate the assembled reforecast cube produced by scripts/08.

Each group lives as 104 per-model-version-date NetCDFs in
data/processed/cube/<group>/<group>_mv2023-MM-DD.nc. The full logical cube for a
group is those files stacked on a new `model_version_date` axis:

    (model_version_date[104], hindcast_year[20], lead[46], [level], lat, lon)

`load_group(group)` returns that as a lazily-concatenated x.Dataset (dask-backed,
so it does not read 18 GB into RAM). `python3 scripts/09_load_cube.py` validates
every group: file count, coordinate consistency, finite fraction, dropped vars.
"""
import glob
import os
import sys

import numpy as np
import xarray as xr

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
CUBE = os.path.join(ROOT, "data", "processed", "cube")
GROUPS = ["sfc_tx", "sfc_t2m", "sfc_dyn", "pl_hi"]
N_DATE = 104


def _mvdate(path):
    return os.path.basename(path).split("_mv")[1].replace(".nc", "")   # 2023-MM-DD


def group_files(group):
    return sorted(glob.glob(os.path.join(CUBE, group, f"{group}_mv*.nc")), key=_mvdate)


def load_group(group, chunks="auto"):
    """Lazily stack a group's per-mvdate files into one cube with a
    `model_version_date` coordinate (datetime64)."""
    files = group_files(group)
    if not files:
        raise FileNotFoundError(f"no files for group {group} under {CUBE}/{group}")
    mvdates = np.array([np.datetime64(_mvdate(f)) for f in files])
    ds = xr.open_mfdataset(files, combine="nested", concat_dim="model_version_date",
                           chunks=chunks, data_vars="all", coords="minimal",
                           compat="override")
    ds = ds.assign_coords(model_version_date=("model_version_date", mvdates))
    return ds


def validate_group(group):
    files = group_files(group)
    status = {"group": group, "n_files": len(files), "expected": N_DATE}
    if not files:
        status["ok"] = False
        status["note"] = "no files yet"
        return status
    # coordinate consistency across files (hyear/lead/lat/lon/level identical)
    ref = xr.open_dataset(files[0])
    ref_sig = {k: ref[k].values for k in ("hindcast_year", "lead", "lat", "lon")
               if k in ref.coords}
    if "level" in ref.coords:
        ref_sig["level"] = ref["level"].values
    mismatch = []
    for f in files:
        d = xr.open_dataset(f)
        for k, v in ref_sig.items():
            if d[k].values.shape != v.shape or not np.array_equal(d[k].values, v):
                mismatch.append((os.path.basename(f), k))
        d.close()
    dropped = ref.attrs.get("dropped_vars")
    status.update(
        vars=list(ref.data_vars),
        dims={k: int(v) for k, v in ref.sizes.items()},
        coord_mismatches=mismatch[:10],
        dropped_vars=dropped,
        ok=(len(files) == N_DATE and not mismatch),
    )
    ref.close()
    return status


def main():
    print(f"cube root: {CUBE}\n")
    for g in GROUPS:
        s = validate_group(g)
        flag = "OK " if s.get("ok") else "-- "
        print(f"[{flag}] {g:8s} files={s['n_files']}/{s['expected']}  "
              f"vars={s.get('vars')}  dims={s.get('dims')}")
        if s.get("dropped_vars"):
            print(f"        dropped: {s['dropped_vars']}")
        if s.get("coord_mismatches"):
            print(f"        COORD MISMATCH: {s['coord_mismatches']}")
    # deep finite-fraction check on the first fully-present group
    print("\nfinite-fraction spot check (model_version_date[0]):")
    for g in GROUPS:
        if not group_files(g):
            continue
        d = xr.open_dataset(group_files(g)[0])
        for v in d.data_vars:
            a = d[v].values
            print(f"  {g}/{v:5s} finite={100*np.isfinite(a).mean():5.1f}%  "
                  f"min/mean/max={np.nanmin(a):.2f}/{np.nanmean(a):.2f}/{np.nanmax(a):.2f}")
        d.close()


if __name__ == "__main__":
    sys.exit(main())