#!/usr/bin/env python
"""Stage 10 — operational producer: PRODUCTION model + per-init CF-NetCDF EHD product.

Two phases (run `train` once, then `produce` for any init source):

  train    Fit the PRODUCTION model on ALL reforecast years (no held-out year — the honest
           skill estimate is the stage-06/08 LOYO cross-validation; this refit on all data is
           the deployment model). For the headline family (hgb) it fits, per lead block:
             * the HistGradientBoosting classifier on all-year block features,
             * the isotonic recalibration on the in-sample predictions,
             * the per-REGION decision probability threshold (Peirce-maximising) on the
               recalibrated in-sample predictions.
           The full-period Tx90p threshold field min(Tx90p,35) (baseline 2002-2023) is saved
           alongside. Artefact: outputs/production_model/ehd_production_model.joblib (+ the
           threshold NetCDF and a JSON manifest).

  produce  Score one or more inits through the saved model and write ONE CF-compliant NetCDF
           per init to outputs/operational/: prob_ehd(block,lat,lon) = P(>=1 EHD in block)
           and ehd_flag(block,lat,lon) = 1 where prob >= that cell's region decision
           threshold. dims (init=1, block=6, lat, lon).

REAL-TIME NOTE (see CLAUDE.md / MEMORY.md L-01): the reforecast is CONTROL-ONLY (1 member)
and this model consumes single-member features. A real-time run must feed the same single
control member (or the ensemble control/mean reduced to one field) built the same way
(daily Tx = max of four 6-hourly mx2t6, 29 Feb dropped, matched on (month,day)). Check the
operational model version matches the reforecast; document any mismatch.
"""
from __future__ import annotations
import os, sys, json, time, logging, argparse
import numpy as np
import xarray as xr
import joblib
from sklearn.isotonic import IsotonicRegression

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, calibration as cal, verify as vf

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("stage10")

HEADLINE_FAMILY = "hgb"          # production family (recalibrated at apply time)
MODEL_DIR = "production_model"
PRODUCT_DIR = "operational"


# ------------------------------------------------------------------ shared geometry
def land_geometry(cfg, ref):
    """Land-cell indices (iy,ix), their region ids, and the region mask aligned to ref grid."""
    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 iy, ix, region_of_cell


def production_threshold(cfg, da):
    """Full-period min(Tx90p,35) threshold field (doy,lat,lon) on the baseline period."""
    b0, b1 = cfg["obs"]["gridded_years"]
    baseline = list(range(b0, b1 + 1))
    tx90p = clim.tx90p_climatology(da, baseline, cfg["index"]["percentile"],
                                   cfg["index"]["percentile_window_days"])
    thr = clim.ehd_threshold(tx90p, cfg["index"]["abs_degC"], cfg["index"]["threshold_rule"])
    return thr, baseline


# ------------------------------------------------------------------ TRAIN
def train(cfg):
    blocks = lead_blocks(cfg)
    region_ids = cfg["regions"]["ids"]
    years = cfg["forecast"]["init_years"]
    outdir = os.path.join(cfg["paths"]["outputs"], MODEL_DIR)
    os.makedirs(outdir, exist_ok=True)

    log.info("loading obs + %d forecast years ...", len(years))
    da = clim.load_gridded_obs(cfg["paths"]["gridded_obs"])
    obs_vals, obs_lut = feat.build_obs_index(da)
    fc = {y: xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx",
                                            f"fc_daily_tx_{y}.nc")).load() for y in years}
    ref = fc[years[0]]
    iy, ix, region_of_cell = land_geometry(cfg, ref)

    thr, baseline = production_threshold(cfg, da)
    thr_i = thr.reindex(lat=ref.lat, lon=ref.lon, method="nearest")

    # assemble ALL years (leakage rule does not apply here: this is the deployment refit)
    acc = {b: {"X": [], "y": [], "cell": []} for b in blocks}
    for y in years:
        out = feat.assemble_year(fc[y], obs_vals, obs_lut, thr_i, blocks, iy, ix)
        for b in blocks:
            acc[b]["X"].append(out[b]["X"])
            acc[b]["y"].append(out[b]["y_primary"])
            acc[b]["cell"].append(out[b]["cell_id"])

    per_block = {}
    for b in blocks:
        t0 = time.time()
        X = np.concatenate(acc[b]["X"]); y = np.concatenate(acc[b]["y"])
        cell = np.concatenate(acc[b]["cell"])
        reg = region_of_cell[cell]
        mdl = cal.make_model(HEADLINE_FAMILY).fit(X, y)
        p = mdl.predict_proba(X)
        m = np.isfinite(y) & np.isfinite(p)
        iso = IsotonicRegression(out_of_bounds="clip").fit(p[m], y[m])
        p_rc = iso.transform(p)
        # per-region Peirce-optimal decision threshold on recalibrated in-sample probs
        g_thr, _ = vf.best_threshold(y[m], p_rc[m], "Peirce")
        region_thr = {}
        for r in region_ids:
            mr = m & (reg == r)
            if mr.sum() > 200 and np.nansum(y[mr]) > 10:
                region_thr[int(r)], _ = vf.best_threshold(y[mr], p_rc[mr], "Peirce")
            else:
                region_thr[int(r)] = g_thr
        base_rate = float(np.nanmean(y[m]))
        per_block[b] = dict(hgb=mdl, iso=iso, region_thr=region_thr,
                            global_thr=float(g_thr), base_rate=base_rate)
        log.info("trained %s in %.0fs | base_rate=%.3f global_thr=%.2f",
                 b, time.time() - t0, base_rate, g_thr)

    thr_i.to_dataset(name="ehd_threshold").to_netcdf(
        os.path.join(outdir, "tx90p_threshold.nc"))
    artefact = dict(
        per_block=per_block, blocks={k: list(v) for k, v in blocks.items()},
        feature_names=feat.FEATURE_NAMES, region_ids=list(region_ids),
        land_iy=iy, land_ix=ix, region_of_cell=region_of_cell,
        lat=ref.lat.values, lon=ref.lon.values,
        trained_years=list(years), baseline_years=baseline,
        index=dict(rule=cfg["index"]["threshold_rule"], abs_degC=cfg["index"]["abs_degC"],
                   percentile=cfg["index"]["percentile"],
                   window_days=cfg["index"]["percentile_window_days"]),
        family=HEADLINE_FAMILY, seed=cfg["seed"], created=time.strftime("%Y-%m-%dT%H:%M:%S"))
    joblib.dump(artefact, os.path.join(outdir, "ehd_production_model.joblib"))

    manifest = dict(created=artefact["created"], family=HEADLINE_FAMILY,
                    trained_years=list(years), baseline_years=[baseline[0], baseline[-1]],
                    index=artefact["index"], blocks=artefact["blocks"],
                    base_rate={b: per_block[b]["base_rate"] for b in blocks},
                    global_thr={b: per_block[b]["global_thr"] for b in blocks},
                    region_thr={b: per_block[b]["region_thr"] for b in blocks},
                    note="Deployment refit on ALL years; honest skill = LOYO CV (stage 06/08). "
                         "Control-only single-member MOS; recalibration fit in-sample.")
    with open(os.path.join(outdir, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)
    log.info("saved production model + manifest -> %s", outdir)


# ------------------------------------------------------------------ PRODUCE
def _cell_threshold_vector(art, block):
    """Per-land-cell decision threshold from that cell's region."""
    rthr = art["per_block"][block]["region_thr"]
    g = art["per_block"][block]["global_thr"]
    return np.array([rthr.get(int(r), g) for r in art["region_of_cell"]], dtype="float64")


def produce(cfg, year, out_subdir=PRODUCT_DIR, max_inits=None):
    outdir = os.path.join(cfg["paths"]["outputs"], MODEL_DIR)
    art = joblib.load(os.path.join(outdir, "ehd_production_model.joblib"))
    thr = xr.open_dataarray(os.path.join(outdir, "tx90p_threshold.nc"))
    blocks = {k: (int(v[0]), int(v[1])) for k, v in art["blocks"].items()}
    iy, ix = art["land_iy"], art["land_ix"]
    lat, lon = art["lat"], art["lon"]
    nlat, nlon = lat.size, lon.size
    block_names = list(blocks.keys())

    fc = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx",
                                        f"fc_daily_tx_{year}.nc")).load()
    thr_i = thr.reindex(lat=fc.lat, lon=fc.lon, method="nearest")
    feats = feat.forecast_features(fc, thr_i, blocks, iy, ix)

    # precompute per-block predictions for ALL inits, then split by init to write files
    ninit = fc.sizes["init"]
    ncell = iy.size
    prob = {b: np.full((ninit, ncell), np.nan) for b in block_names}
    flag = {b: np.zeros((ninit, ncell), dtype=np.int8) for b in block_names}
    for b in block_names:
        X = feats[b]["X"]
        p = art["per_block"][b]["hgb"].predict_proba(X)
        p = art["per_block"][b]["iso"].transform(p)
        p = p.reshape(ninit, ncell)
        cthr = _cell_threshold_vector(art, b)[None, :]
        prob[b] = p
        flag[b] = (p >= cthr).astype(np.int8)

    dest = os.path.join(cfg["paths"]["outputs"], out_subdir)
    os.makedirs(dest, exist_ok=True)
    inits = np.arange(ninit if max_inits is None else min(max_inits, ninit))
    written = []
    for it in inits:
        p_grid = np.full((len(block_names), nlat, nlon), np.nan, dtype="float32")
        f_grid = np.full((len(block_names), nlat, nlon), -1, dtype="int8")   # -1 = sea/out
        for k, b in enumerate(block_names):
            p_grid[k, iy, ix] = prob[b][it]
            f_grid[k, iy, ix] = flag[b][it]
        init_dt = np.datetime64(fc["init"].values[it], "D")
        ds = _make_dataset(cfg, art, blocks, block_names, lat, lon, p_grid, f_grid, init_dt)
        stamp = str(init_dt).replace("-", "")
        path = os.path.join(dest, f"ehd_product_{stamp}.nc")
        ds.to_netcdf(path, encoding={
            "prob_ehd": {"zlib": True, "complevel": 4, "_FillValue": np.float32(np.nan)},
            "ehd_flag": {"zlib": True, "complevel": 4, "_FillValue": np.int8(-1)}})
        written.append(path)
    log.info("wrote %d per-init products for %d -> %s", len(written), year, dest)
    return written


def _make_dataset(cfg, art, blocks, block_names, lat, lon, p_grid, f_grid, init_dt):
    bidx = np.arange(1, len(block_names) + 1, dtype="int8")
    d0 = np.array([blocks[b][0] for b in block_names], dtype="int16")
    d1 = np.array([blocks[b][1] for b in block_names], dtype="int16")
    ds = xr.Dataset(
        data_vars=dict(
            prob_ehd=(("block", "lat", "lon"), p_grid,
                      {"long_name": "probability of at least one extreme-heat day in the lead block",
                       "units": "1", "valid_range": np.array([0.0, 1.0], "float32"),
                       "coordinates": "init"}),
            ehd_flag=(("block", "lat", "lon"), f_grid,
                      {"long_name": "operational extreme-heat-day flag (prob >= region decision threshold)",
                       "units": "1", "flag_values": np.array([0, 1], "int8"),
                       "flag_meanings": "no_ehd ehd", "comment": "-1 = sea / outside domain",
                       "coordinates": "init"}),
            block_day_start=(("block",), d0, {"long_name": "first forecast lead day in block (days after init)",
                                              "units": "day"}),
            block_day_end=(("block",), d1, {"long_name": "last forecast lead day in block (days after init)",
                                            "units": "day"}),
        ),
        coords=dict(
            block=("block", bidx, {"long_name": "lead-time block index",
                                   "block_names": " ".join(block_names)}),
            lat=("lat", lat.astype("float32"),
                 {"standard_name": "latitude", "units": "degrees_north", "axis": "Y"}),
            lon=("lon", lon.astype("float32"),
                 {"standard_name": "longitude", "units": "degrees_east", "axis": "X"}),
            init=((), init_dt.astype("datetime64[ns]"),
                  {"standard_name": "forecast_reference_time", "long_name": "forecast initialisation date"}),
        ),
        attrs=dict(
            Conventions="CF-1.8",
            title="Operational S2S extreme-heat-day (EHD) probability forecast for Vietnam",
            institution="HUS (Hanoi University of Science)",
            source=("Single-member (control) S2S reforecast Tx (daily max of 6-hourly mx2t6), "
                    "calibrated by single-member MOS (histogram gradient boosting) with isotonic "
                    "recalibration; production model refit on all reforecast years."),
            references="Project QG_1; see docs/RESULTS_REPORT.md and docs/OPERATIONAL_GUIDE.md",
            comment=("EHD day iff Tx >= min(Tx90p, {:.0f} degC), Tx90p = calendar-day {}th "
                     "percentile over {}-{} with a {}-day centred window. Skilful horizon is "
                     "lead- and region-dependent (see report): significant Brier skill over "
                     "climatology to W3 (day 24) pooled; longest in the tropical south and in MAM."
                     ).format(art["index"]["abs_degC"], art["index"]["percentile"],
                              art["baseline_years"][0], art["baseline_years"][-1],
                              art["index"]["window_days"]),
            model_version=art["created"], family=art["family"],
            history=time.strftime("%Y-%m-%dT%H:%M:%S") + " scripts/10_operational_product.py produce",
        ),
    )
    return ds


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("phase", choices=["train", "produce"], help="train the model or produce products")
    ap.add_argument("--year", type=int, default=None,
                    help="produce: forecast init year to score (default: latest reforecast year)")
    ap.add_argument("--max-inits", type=int, default=None,
                    help="produce: cap number of inits (demo/testing)")
    args = ap.parse_args()
    cfg = load_config()
    if args.phase == "train":
        train(cfg)
    else:
        year = args.year or max(cfg["forecast"]["init_years"])
        produce(cfg, year, max_inits=args.max_inits)


if __name__ == "__main__":
    main()