#!/usr/bin/env python
"""Stage 27 — build the FULL-FIELD spatial cache for the true-spatial ResU-Net (v4).

Stage 20a reduced the large-scale circulation to a handful of BOX-MEAN indices (z500_reg,
z500_sub, ...). The ResU-Net could therefore only see those boxes as broadcast constant
channels, never the actual spatial pattern (the HONEST LIMITATION in MEMORY qg1-hybrid-v3).

This stage caches the FULL gridded fields, block-MEAN over the six fixed lead blocks, on both
grids the archive provides:

  ext  — extended 71x61 domain (1 deg, 80-150E / -15-45N): Z500, T850, U/V200 (circulation).
  vn   — native 41x65 domain (0.25 deg): all multi-level (T/H/W/QUV_plev) + surface
         (Srf_daily_ave, Srf_24hour) fields, as native-resolution spatial channels.

All operations are block-mean reductions over lead-day windows — NO fitting, NO climatology,
NO cross-fold information — so the cache is leakage-free and can be built once for all years.
The channel set is declared in config.yaml `fullfield:` (single source of truth). Ocean cells
are KEPT (no land mask): the point of the extended field is the SST/circulation pattern.

Output (one file per year, row-aligned to the Stage-1 forecast init order so downstream code
can zip it against the existing per-cell caches):

  outputs/fullfield_cache/ff_<year>.npz
    init_ymd[ninit]
    ext[ninit, nblock, 61, 71, n_ext_ch]  float32,  ext_names[n_ext_ch]
    vn [ninit, nblock, 65, 41, n_vn_ch ]  float32,  vn_names [n_vn_ch]
    ext_lat[61], ext_lon[71], vn_lat[65], vn_lon[41], block_bounds[nblock,2]

Run one task per year (SLURM array) — comprehensive channel set is ~0.3-0.4 GB/yr in RAM.
"""
from __future__ import annotations
import os, sys, argparse, time, glob, warnings
import numpy as np
import xarray as xr
import pandas as pd

# block-mean over all-NaN ocean columns (land-only soil vars) -> benign "empty slice"
warnings.filterwarnings("ignore", r"Mean of empty slice", RuntimeWarning)

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.config import load_config, lead_blocks

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)


def _first_file(ec_root, d, prefix):
    """First existing file for a group, to probe which vars carry a level dim."""
    for yr in range(2003, 2023):
        fs = sorted(glob.glob(os.path.join(ec_root, d, str(yr), f"{prefix}_*.nc")))
        if fs:
            return fs[0]
    return None


def build_channel_plan(cfg):
    """Probe one file per group -> ordered channel plan per grid.

    Returns dict grid -> list of (chan_name, group_idx, var, level_or_None) and, for each
    group, the read spec (dir, prefix, path template). A var gets one channel per configured
    level ONLY if it actually has a 'level' dim in the data; otherwise a single channel.
    """
    ec_root = cfg["paths"]["ec_root"]
    groups = cfg["fullfield"]["groups"]
    plan = {"ext": [], "vn": []}
    specs = []
    for gi, g in enumerate(groups):
        specs.append({"dir": g["dir"], "prefix": g["prefix"], "grid": g["grid"]})
        probe = _first_file(ec_root, g["dir"], g["prefix"])
        if probe is None:
            print(f"[27] WARNING: no file found for group {g['dir']} — skipped")
            continue
        ds = xr.open_dataset(probe)
        cfg_levels = g.get("levels")
        for v in g["vars"]:
            if v not in ds:
                print(f"[27] WARNING: {v} absent in {g['dir']} — skipped")
                continue
            has_level = "level" in ds[v].dims
            if has_level:
                levels = cfg_levels if cfg_levels else [int(x) for x in ds["level"].values]
                for lv in levels:
                    plan[g["grid"]].append((f"{v}{int(lv)}", gi, v, int(lv)))
            else:
                plan[g["grid"]].append((v, gi, v, None))
        ds.close()
    return plan, specs


def block_mean_field(da, block_ranges):
    """da: (time, lat, lon) -> (nblock, lat, lon) block-mean; NaN where block beyond horizon."""
    arr = np.asarray(da.values, dtype=np.float32)
    nt = arr.shape[0]
    out = np.full((len(block_ranges),) + arr.shape[1:], np.nan, np.float32)
    for bi, (d0, d1) in enumerate(block_ranges):
        hi = min(d1, nt - 1)
        if d0 > hi:
            continue
        out[bi] = np.nanmean(arr[d0:hi + 1], axis=0)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--years", type=int, nargs="*", default=None)
    ap.add_argument("--max-inits", type=int, default=None, help="smoke-test: cap inits/year")
    args = ap.parse_args()

    cfg = load_config()
    ec_root = cfg["paths"]["ec_root"]
    blocks = lead_blocks(cfg)
    block_ranges = list(blocks.values())
    nblock = len(block_ranges)
    ff = cfg["fullfield"]
    ne_lat, ne_lon = ff["ext_grid"]["nlat"], ff["ext_grid"]["nlon"]
    nv_lat, nv_lon = ff["vn_grid"]["nlat"], ff["vn_grid"]["nlon"]

    plan, specs = build_channel_plan(cfg)
    ext_names = [p[0] for p in plan["ext"]]
    vn_names = [p[0] for p in plan["vn"]]
    print(f"[27] channels: ext={len(ext_names)} {ext_names}")
    print(f"[27] channels: vn ={len(vn_names)} {vn_names}")

    years = args.years or list(cfg["forecast"]["init_years"])
    outdir = os.path.join(cfg["paths"]["outputs"], "fullfield_cache")
    os.makedirs(outdir, exist_ok=True)

    # group index -> list of (chan_slot, var, level) it must fill, per grid
    ext_by_group = {}
    vn_by_group = {}
    for slot, (_, gi, v, lv) in enumerate(plan["ext"]):
        ext_by_group.setdefault(gi, []).append((slot, v, lv))
    for slot, (_, gi, v, lv) in enumerate(plan["vn"]):
        vn_by_group.setdefault(gi, []).append((slot, v, lv))

    for yr in years:
        t0 = time.time()
        fc = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx",
                                            f"fc_daily_tx_{yr}.nc"))
        inits = pd.to_datetime(fc["init"].values)
        fc.close()
        if args.max_inits:
            inits = inits[:args.max_inits]
        ninit = len(inits)
        init_ymd = np.array([d.year * 10000 + d.month * 100 + d.day for d in inits], np.int64)

        ext = np.full((ninit, nblock, ne_lat, ne_lon, len(ext_names)), np.nan, np.float32)
        vn = np.full((ninit, nblock, nv_lat, nv_lon, len(vn_names)), np.nan, np.float32)
        ext_lat = ext_lon = vn_lat = vn_lon = None

        for it, t in enumerate(inits):
            ds_str = t.strftime("%Y-%m-%d")
            for gi, spec in enumerate(specs):
                grid = spec["grid"]
                slots = (ext_by_group if grid == "ext" else vn_by_group).get(gi)
                if not slots:
                    continue
                path = os.path.join(ec_root, spec["dir"], str(yr), f"{spec['prefix']}_{ds_str}.nc")
                if not os.path.exists(path):
                    continue
                d = xr.open_dataset(path)
                if grid == "ext" and ext_lat is None:
                    ext_lat = d["latitude"].values.astype(np.float32)
                    ext_lon = d["longitude"].values.astype(np.float32)
                if grid == "vn" and vn_lat is None:
                    vn_lat = d["latitude"].values.astype(np.float32)
                    vn_lon = d["longitude"].values.astype(np.float32)
                for slot, v, lv in slots:
                    if v not in d:
                        continue
                    da = d[v] if lv is None else d[v].sel(level=lv)
                    bm = block_mean_field(da, block_ranges)   # (nblock, lat, lon)
                    if grid == "ext":
                        ext[it, :, :, :, slot] = bm
                    else:
                        vn[it, :, :, :, slot] = bm
                d.close()

        outf = os.path.join(outdir, f"ff_{yr}.npz")
        np.savez_compressed(
            outf, init_ymd=init_ymd,
            ext=ext, vn=vn,
            ext_names=np.array(ext_names), vn_names=np.array(vn_names),
            ext_lat=ext_lat if ext_lat is not None else np.zeros(ne_lat, np.float32),
            ext_lon=ext_lon if ext_lon is not None else np.zeros(ne_lon, np.float32),
            vn_lat=vn_lat if vn_lat is not None else np.zeros(nv_lat, np.float32),
            vn_lon=vn_lon if vn_lon is not None else np.zeros(nv_lon, np.float32),
            block_bounds=np.array(block_ranges, np.int64),
        )
        mb = os.path.getsize(outf) / 1e6
        print(f"[27] {yr}: {ninit} inits -> {os.path.basename(outf)} ({mb:.0f} MB, "
              f"{time.time()-t0:.0f}s)  ext NaN={np.mean(np.isnan(ext)):.3f} "
              f"vn NaN={np.mean(np.isnan(vn)):.3f}", flush=True)


if __name__ == "__main__":
    main()