#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 6 — non-LLM baselines (A hedonic, B hierarchical/repeat-sales, C LightGBM [+CatBoost/XGBoost], D AVM) on the
prospective evaluation set.

The fitted baseline objects are not persisted by scripts/10_train_baselines.py, so the arms are refit here with the
identical protocol and seeds: train on `train`, tune on `validation` (same random search, 10 trials), and predict the
`prospective` rows. The feature matrix is built with features.matrix.build_feature_matrix on the prospective bundle
(snapshot ∪ prospective rows) because the as-of neighbourhood / municipality features of a prospective sale need the
snapshot history; the matrix is restricted to train, validation and prospective rows. As-of invariance is verified:
train/validation feature rows rebuilt on the union must equal those of data/processed/feature_matrix.parquet.

Usage: python3 scripts/62_prospective_baselines.py [--trials 10] [--skip-robustness]
Output: results/runs/prospective_baselines/{predictions.parquet, manifest.json, *_search.json}
"""
from __future__ import annotations

import argparse
import json
import platform
import sys
import time
from pathlib import Path

import numpy as np
import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from agentic_valuation.config import PROCESSED, RESULTS  # noqa: E402
from agentic_valuation.baselines.gbm import AVM, CatBoostModel, LightGBMModel, XGBoostModel, dump_search_log  # noqa: E402
from agentic_valuation.baselines.hedonic import HedonicRidge, HierarchicalHedonic  # noqa: E402
from agentic_valuation.features.matrix import build_feature_matrix  # noqa: E402

BUNDLE = PROCESSED / "prospective_bundle"
RUN_ID = "prospective_baselines"


def asof_invariance(fm_new: pl.DataFrame, fm_ref: pl.DataFrame) -> dict:
    """Max absolute difference of every numeric feature between the union-built matrix and the reference matrix,
    on train ∪ validation rows (must be 0: only later-dated sales were added)."""
    keep = fm_ref.filter(pl.col("split").is_in(["train", "validation"]))
    num = [c for c, t in keep.schema.items() if t == pl.Float64 and c not in ("price", "log_price")]
    j = keep.select(["tx_id", *num]).join(fm_new.select(["tx_id", *[pl.col(c).alias(f"{c}__new") for c in num]]), on="tx_id", how="inner")
    diffs = {}
    for c in num:
        a, b = j[c].to_numpy(), j[f"{c}__new"].to_numpy()
        both_nan = np.isnan(a) & np.isnan(b)
        d = np.where(both_nan, 0.0, np.abs(a - b))
        diffs[c] = float(np.nanmax(d)) if d.size else 0.0
    return {"rows_compared": j.height, "max_abs_diff": max(diffs.values()) if diffs else 0.0, "by_feature_nonzero": {k: v for k, v in diffs.items() if v > 0}}


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--trials", type=int, default=10, help="random-search trials (10 = the count used by results/runs/baselines)")
    ap.add_argument("--skip-robustness", action="store_true", help="skip C_catboost / C_xgboost")
    a = ap.parse_args(argv)
    out = RESULTS / "runs" / RUN_ID
    out.mkdir(parents=True, exist_ok=True)
    t0 = time.perf_counter()
    analytic = pl.read_parquet(BUNDLE / "analytic.parquet")
    splits = pl.read_parquet(BUNDLE / "splits.parquet")
    geo_path = BUNDLE / "geo_features.parquet"
    geo = pl.read_parquet(geo_path) if geo_path.exists() else None
    fm = build_feature_matrix(analytic, splits, geo=geo, comps_summary=True)
    fm = fm.filter(pl.col("split").is_in(["train", "validation", "prospective"]))
    t_fm = time.perf_counter() - t0
    ref_path = PROCESSED / "feature_matrix.parquet"
    invariance = asof_invariance(fm, pl.read_parquet(ref_path)) if ref_path.exists() else {"skipped": "reference feature matrix missing"}
    train = fm.filter(pl.col("split") == "train")
    valid = fm.filter(pl.col("split") == "validation")
    pros = fm.filter(pl.col("split") == "prospective")
    print(f"feature matrix {fm.height:,} rows in {t_fm:.0f}s | train {train.height:,} | valid {valid.height:,} | prospective {pros.height:,} | as-of invariance max diff {invariance.get('max_abs_diff')}")
    if pros.height == 0:
        sys.exit("no prospective rows in the bundle: run scripts/61_build_prospective.py first")

    LAGS = ("mun_", "nb_")
    models = {
        "A_hedonic": lambda: HedonicRidge(),
        "B_hierarchical": lambda: HierarchicalHedonic(),
        "C_gbm": lambda: LightGBMModel(n_trials=a.trials, exclude_prefixes=LAGS),
        "D_avm": lambda: AVM(n_trials=a.trials),
    }
    if not a.skip_robustness:
        models["C_catboost"] = lambda: CatBoostModel(exclude_prefixes=LAGS)
        models["C_xgboost"] = lambda: XGBoostModel(exclude_prefixes=LAGS)
    records, manifest = [], {
        "run_id": RUN_ID, "python": platform.python_version(), "machine": platform.machine(), "trials": a.trials,
        "data_dir": str(BUNDLE.relative_to(PROCESSED.parents[1])), "geo_joined": geo is not None,
        "n_train": train.height, "n_validation": valid.height, "n_prospective": pros.height, "feature_matrix_seconds": round(t_fm, 1),
        "asof_invariance_train_validation": invariance, "protocol": "fit on train, tune on validation, predict prospective (identical to scripts/10_train_baselines.py)",
        "arms": {},
    }
    for arm, ctor in models.items():
        t1 = time.perf_counter()
        m = ctor().fit(train, valid)
        fit_s = time.perf_counter() - t1
        t2 = time.perf_counter()
        pred = m.predict(pros)
        records.append(pred.to_records("prospective", arm, RUN_ID, latency_s=time.perf_counter() - t2))
        info = {"fit_seconds": round(fit_s, 1)}
        if isinstance(m, LightGBMModel):
            dump_search_log(m, out / f"{arm}_search.json")
            info["params"] = m.params
        if isinstance(m, AVM):
            dump_search_log(m.point, out / f"{arm}_search.json")
            info["params"] = m.point.params
            info["conformal_scale"] = m.scale
        if isinstance(m, HierarchicalHedonic):
            info["w_repeat"] = m.w_repeat
        manifest["arms"][arm] = info
        print(f"{arm}: fit {fit_s:.0f}s {info}", flush=True)
    preds = pl.concat(records)
    preds.write_parquet(out / "predictions.parquet")
    manifest["n_prediction_rows"] = preds.height
    manifest["seconds"] = round(time.perf_counter() - t0, 1)
    (out / "manifest.json").write_text(json.dumps(manifest, indent=1, default=float))
    truth = fm.select(["tx_id", "log_price"])
    s = preds.join(truth, on="tx_id").with_columns((pl.col("point").log() - pl.col("log_price")).abs().alias("ae")).group_by(["arm", "split"]).agg(pl.col("ae").mean().round(4).alias("mean_abs_log_error"), pl.col("ae").median().round(4).alias("median_abs_log_error"), pl.len()).sort(["split", "arm"])
    print(s)


if __name__ == "__main__":
    main()