#!/usr/bin/env python
"""Stage 03 — extract daily Tx from the 6-hourly control reforecast, one file per init year.

For every initialization it derives daily Tx (46 lead days, degC) from `mx2t6` and stacks
all inits of a year into a single NetCDF:
    outputs/forecast_daily_tx/fc_daily_tx_<year>.nc
    dims (init, lead_day=46, lat, lon); coords init (datetime), valid_date(init, lead_day).

This is the deterministic predictor backbone. The observed EHD targets and the per-fold
Tx90p threshold are applied later inside cross-validation (leakage rule), so nothing here
touches the held-out-year climatology. A per-run validation summary is written to
results/tables/forecast_extraction_summary.csv.
"""
from __future__ import annotations
import os, sys, logging, argparse
import numpy as np
import pandas as pd
import xarray as xr

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

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


def build_year(cfg, year, outdir):
    inits = fio.list_inits(cfg["paths"]["forecast_tx"], year)
    if not inits:
        log.warning("no inits for %d", year); return None
    das, init_dates = [], []
    for p in inits:
        da = fio.load_forecast_daily_tx(p, var=cfg["forecast"]["tx_var"])
        das.append(da)
        init_dates.append(pd.Timestamp(da.attrs["init_date"]))
    stacked = xr.concat(das, dim="init")
    stacked = stacked.assign_coords(init=("init", pd.to_datetime(init_dates)))
    # valid_date is (init, lead_day)
    vd = np.stack([d["valid_date"].values for d in das], axis=0)
    stacked = stacked.assign_coords(valid_date=(("init", "lead_day"), vd))
    stacked = stacked.astype("float32")
    stacked.name = "tx"
    stacked.attrs.update(dict(units="C", description="daily Tx from 6-hourly mx2t6 control forecast",
                              members="1 (control only)"))
    out = os.path.join(outdir, f"fc_daily_tx_{year}.nc")
    enc = {"tx": {"zlib": True, "complevel": 4, "_FillValue": np.float32(np.nan)}}
    stacked.to_netcdf(out, encoding=enc)
    summary = dict(year=year, n_inits=len(inits),
                   tx_min=round(float(np.nanmin(stacked.values)), 2),
                   tx_max=round(float(np.nanmax(stacked.values)), 2),
                   tx_mean=round(float(np.nanmean(stacked.values)), 2),
                   nan_frac=round(float(np.isnan(stacked.values).mean()), 4),
                   path=out)
    log.info("%d: %d inits  Tx[%.1f,%.1f] mean %.1f  -> %s",
             year, len(inits), summary["tx_min"], summary["tx_max"], summary["tx_mean"], out)
    return summary


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--years", nargs="*", type=int, default=None)
    args = ap.parse_args()
    cfg = load_config()
    years = args.years or cfg["forecast"]["init_years"]
    outdir = os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx")
    os.makedirs(outdir, exist_ok=True)
    rows = [s for y in years if (s := build_year(cfg, y, outdir))]
    df = pd.DataFrame(rows)
    tbl = os.path.join(cfg["paths"]["results"], "tables", "forecast_extraction_summary.csv")
    df.to_csv(tbl, index=False)
    log.info("wrote %s\n%s", tbl, df.to_string(index=False))


if __name__ == "__main__":
    main()