#!/usr/bin/env python
"""Stage 20a — pre-assemble the GRIDDED Stage-2 candidate predictors into a compact per-year cache.

The raw S2S fields (soil moisture/temperature, surface thermodynamics, and the large-scale
circulation) are per-init NetCDFs; block-aggregating them on the fly inside the 20-fold LOYO would
re-read thousands of files per run. This stage does that aggregation ONCE per year and caches the
block-reduced arrays, row-aligned to the Stage-1 (init, cell) order, so the LOYO runner just loads
the small cache.

Two predictor families are produced (both leakage-free — no fitting, block-mean reductions only):

  LOCAL, per (init, block, cell)  — from data/EC/Srf_daily_ave/ (already on the 41x65 obs grid):
    sm20, sm100   soil moisture (top 20 cm / 100 cm) — land-surface memory
    st100         deep soil temperature — land-surface memory
    skt           skin temperature
    cape          convective available potential energy
    tcw           total column water (moisture)
    tcc           total cloud cover
    d2m           2 m dewpoint (low-level humidity)

  LARGE-SCALE, per (init, block)  — from data/EC/H500 (Z500) and T850 (larger 71x61 domain):
    z500_reg      mean Z500 over the Vietnam/South-China-Sea box (ridge strength)
    z500_sub      mean Z500 over the subtropical box to the north
    z500_grad     z500_sub - z500_reg (meridional height gradient)
    t850_reg      mean 850-hPa temperature over the Vietnam box (low-level warmth/advection)

Output: outputs/ext_predictors/extgrid_<year>.npz with
    init_ymd[ninit], local_names, ls_names,
    local[ninit, nblock, ncell, nlocal], large[ninit, nblock, nls]
Cell order and init order match scripts feat.assemble_year exactly.

Run one task per year (parallelised on SLURM); ~104 inits x 3 groups per year.
"""
from __future__ import annotations
import os, sys, argparse, time, glob
import numpy as np
import xarray as xr
import pandas as pd

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

LOCAL_VARS = ["sm20", "sm100", "st100", "skt", "cape", "tcw", "tcc", "d2m"]
# large-scale boxes (lon_min, lon_max, lat_min, lat_max) on the 71x61 domain
BOX_REG = (100.0, 115.0, 8.0, 24.0)     # Vietnam + South China Sea
BOX_SUB = (100.0, 120.0, 20.0, 35.0)    # subtropical, to the north
LS_NAMES = ["z500_reg", "z500_sub", "z500_grad", "t850_reg"]

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


def _box_mean(da, box):
    lo0, lo1, la0, la1 = box
    lat = da["latitude"]; lon = da["longitude"]
    sub = da.where((lon >= lo0) & (lon <= lo1) & (lat >= la0) & (lat <= la1), drop=True)
    return sub.mean(dim=("latitude", "longitude"), skipna=True)


def land_cells(cfg, ref):
    reg = xr.open_dataset(cfg["paths"]["mask"])["reg"].reindex(
        lat=ref.lat, lon=ref.lon, method="nearest")
    iy, ix = np.where(np.isfinite(reg.values))
    return iy, ix


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--years", type=int, nargs="*", default=None)
    args = ap.parse_args()
    cfg = load_config()
    blocks = lead_blocks(cfg)
    block_ranges = list(blocks.values())          # [(d0,d1), ...] inclusive lead-day
    nblock = len(block_ranges)
    years = args.years or list(cfg["forecast"]["init_years"])
    outdir = os.path.join(cfg["paths"]["outputs"], "ext_predictors")
    os.makedirs(outdir, exist_ok=True)

    for yr in years:
        t0 = time.time()
        # init order = the Stage-1 forecast file's init coord (guarantees row alignment)
        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)
        iy, ix = land_cells(cfg, fc)
        ncell = iy.size; ninit = len(inits)
        fc.close()

        local = np.full((ninit, nblock, ncell, len(LOCAL_VARS)), np.nan, np.float32)
        large = np.full((ninit, nblock, len(LS_NAMES)), np.nan, np.float32)
        init_ymd = np.array([d.year * 10000 + d.month * 100 + d.day for d in inits], np.int64)

        for it, t in enumerate(inits):
            ds = t.strftime("%Y-%m-%d")
            # --- LOCAL surface fields (41x65 grid) ---
            sf = os.path.join(ROOT, "data", "EC", "Srf_daily_ave", str(yr), f"Srf_daily_cf_{ds}.nc")
            if os.path.exists(sf):
                d = xr.open_dataset(sf)
                nt = d.sizes["time"]
                for vi, v in enumerate(LOCAL_VARS):
                    if v not in d:
                        continue
                    arr = d[v].values                      # (time, lat, lon)
                    for bi, (d0, d1) in enumerate(block_ranges):
                        hi = min(d1, nt - 1)
                        if d0 > hi:
                            continue
                        bm = np.nanmean(arr[d0:hi + 1], axis=0)     # (lat, lon)
                        local[it, bi, :, vi] = bm[iy, ix]
                d.close()
            # --- LARGE-SCALE circulation (71x61 domain) ---
            hf = os.path.join(ROOT, "data", "EC", "H500", str(yr), f"H500_cf_{ds}.nc")
            tf = os.path.join(ROOT, "data", "EC", "T850", str(yr), f"T850_cf_{ds}.nc")
            if os.path.exists(hf) and os.path.exists(tf):
                dh = xr.open_dataset(hf); dt = xr.open_dataset(tf)
                gh = dh["gh"]; t8 = dt["t"]
                z_reg = _box_mean(gh, BOX_REG).values      # (time,)
                z_sub = _box_mean(gh, BOX_SUB).values
                t_reg = _box_mean(t8, BOX_REG).values
                nt = gh.sizes["time"]
                for bi, (d0, d1) in enumerate(block_ranges):
                    hi = min(d1, nt - 1)
                    if d0 > hi:
                        continue
                    zr = np.nanmean(z_reg[d0:hi + 1]); zs = np.nanmean(z_sub[d0:hi + 1])
                    large[it, bi, 0] = zr
                    large[it, bi, 1] = zs
                    large[it, bi, 2] = zs - zr
                    large[it, bi, 3] = np.nanmean(t_reg[d0:hi + 1])
                dh.close(); dt.close()

        outf = os.path.join(outdir, f"extgrid_{yr}.npz")
        np.savez_compressed(outf, init_ymd=init_ymd, local=local, large=large,
                            local_names=np.array(LOCAL_VARS), ls_names=np.array(LS_NAMES))
        print(f"[20a] {yr}: {ninit} inits x {ncell} cells -> {os.path.basename(outf)} "
              f"({time.time()-t0:.0f}s)  local NaN frac={np.mean(np.isnan(local)):.3f} "
              f"large NaN frac={np.mean(np.isnan(large)):.3f}")


if __name__ == "__main__":
    main()