#!/usr/bin/env python
"""Stage 20b — assemble the full CANDIDATE design matrix (base + teleconnection + persistence +
gridded) once per year and cache it, so the feature-screening (20c) and any retrain reuse it
instead of re-reading raw fields.

Feature blocks (28 candidates):
  base (10)   : features.FEATURE_NAMES        (forecast-Tx statistics + calendar + location)
  tele (4)    : enso_nino34, mjo_pc1/pc2/amp  (observed init-state teleconnections)
  persist (2) : f_ante_ehdfreq, f_ante_sevmean
  gridded (12): 8 local surface (soil moisture/temp, skt, cape, tcw, tcc, d2m) + 4 large-scale
                Z500/T850 indices (from stage 20a cache)

Threshold note: this cache uses the FULL-baseline diagnostic climatology `min(Tx90p, 35)` (all
2002-2023 obs). It backs the descriptive feature-importance RANKING (20c). It is NOT used for any
out-of-sample skill claim — the honest skill of a selected feature set comes from the fully nested,
per-fold LOYO retrain (fold-specific threshold). This separation is stated in the report.

Output: outputs/candidates/cand_<year>.npz with X[nrows,28], feat_names, y, sev, block_id,
region, cell, year. Row order within a year = block-major, then assemble_year (init, cell).
"""
from __future__ import annotations
import os, sys, argparse, time
import numpy as np
import xarray as xr

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.config import load_config, lead_blocks
from ehd import climatology as clim, features as feat, features_ext as fext

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

FEAT_NAMES = list(feat.FEATURE_NAMES) + list(fext.EXT_FEATURE_NAMES) + list(fext.GRID_FEATURE_NAMES)


def preload(cfg):
    da = clim.load_gridded_obs(cfg["paths"]["gridded_obs"])
    obs_vals, obs_lut = feat.build_obs_index(da)
    ref = None
    fc = {}
    for y in cfg["forecast"]["init_years"]:
        fc[y] = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx",
                                               f"fc_daily_tx_{y}.nc")).load()
        ref = fc[y]
    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))
    region_of_cell = reg.values[iy, ix].astype(int)
    return da, obs_vals, obs_lut, fc, iy, ix, region_of_cell


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_names = list(blocks)
    block_ids = {b: k + 1 for k, b in enumerate(block_names)}
    outdir = os.path.join(cfg["paths"]["outputs"], "candidates")
    os.makedirs(outdir, exist_ok=True)
    years = args.years or list(cfg["forecast"]["init_years"])

    nino34 = fext.load_nino34(os.environ.get("NINO34_TXT",
                              os.path.join(ROOT, "..", "HW", "data", "external", "nino34.txt")))
    omi = fext.load_omi(os.environ.get("OMI_TXT",
                        os.path.join(ROOT, "..", "HW", "data", "external", "omi.txt")))

    da, obs_vals, obs_lut, fc, iy, ix, region_of_cell = preload(cfg)
    # full-baseline diagnostic threshold (all obs years)
    base_years = list(range(cfg["obs"]["gridded_years"][0], cfg["obs"]["gridded_years"][1] + 1))
    thr = clim.ehd_threshold(clim.tx90p_climatology(da, base_years, 90, 15), 35.0, "min")
    ncell = iy.size

    for yr in years:
        t0 = time.time()
        thr_i = thr.reindex(lat=fc[yr].lat, lon=fc[yr].lon, method="nearest")
        out = feat.assemble_year(fc[yr], obs_vals, obs_lut, thr_i, blocks, iy, ix)
        ext = fext.ext_features_year(fc[yr], obs_vals, obs_lut, thr_i, iy, ix, nino34, omi)["X_ext"]
        eg = fext.load_extgrid(cfg["paths"]["outputs"], yr)

        Xs, ys, sevs, bids, regs, cells = [], [], [], [], [], []
        for bi, b in enumerate(block_names):
            base = out[b]["X"]
            Xg = fext.ext_grid_block(eg, bi, ncell)
            assert base.shape[0] == ext.shape[0] == Xg.shape[0], \
                f"row mismatch {b}: base {base.shape[0]} ext {ext.shape[0]} grid {Xg.shape[0]}"
            Xfull = np.column_stack([base, ext, Xg]).astype(np.float32)
            n = Xfull.shape[0]
            Xs.append(Xfull); ys.append(out[b]["y_primary"]); sevs.append(out[b]["sev"])
            cid = out[b]["cell_id"]
            bids.append(np.full(n, block_ids[b], np.int16))
            regs.append(region_of_cell[cid].astype(np.int16)); cells.append(cid.astype(np.int32))
        X = np.concatenate(Xs); assert X.shape[1] == len(FEAT_NAMES)
        np.savez_compressed(os.path.join(outdir, f"cand_{yr}.npz"),
                            X=X, y=np.concatenate(ys).astype(np.float32),
                            sev=np.concatenate(sevs).astype(np.float32),
                            block_id=np.concatenate(bids), region=np.concatenate(regs),
                            cell=np.concatenate(cells),
                            feat_names=np.array(FEAT_NAMES), year=np.int64(yr))
        print(f"[20b] {yr}: X{X.shape} feats={len(FEAT_NAMES)} "
              f"NaN frac={np.mean(np.isnan(X)):.3f} ({time.time()-t0:.0f}s)")


if __name__ == "__main__":
    main()