#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 3 — fit the non-LLM baselines (A hedonic, B hierarchical/repeat-sales, C LightGBM [+CatBoost/XGBoost],
D AVM with conformal quantiles) on train, tune on validation, and write prediction records for the validation,
test and holdout partitions (full partitions) plus the spatial-holdout experiment.

Usage: python3 scripts/10_train_baselines.py [--synthetic] [--trials 12] [--skip-robustness]
Outputs: results/runs/baselines[_synthetic]/predictions.parquet, manifest.json, search logs.
"""
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, SYNTHETIC, load_config  # 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


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--synthetic", action="store_true")
    ap.add_argument("--trials", type=int, default=12)
    ap.add_argument("--skip-robustness", action="store_true")
    ap.add_argument("--max-train", type=int, default=0, help="subsample train rows (speed; 0 = all)")
    a = ap.parse_args(argv)
    base = SYNTHETIC if a.synthetic else PROCESSED
    run_id = "baselines_synthetic" if a.synthetic else "baselines"
    out = RESULTS / "runs" / run_id
    out.mkdir(parents=True, exist_ok=True)
    fm = pl.read_parquet(base / "feature_matrix.parquet")
    if "split" not in fm.columns:
        fm = fm.join(pl.read_parquet(base / "splits.parquet").select(["tx_id", "split"]), on="tx_id", how="left")
    spatial_regions = load_config("data")["spatial_holdout"]["regions"]
    train = fm.filter(pl.col("split") == "train")
    if a.max_train and train.height > a.max_train:
        train = train.sample(a.max_train, seed=20261027)
    valid = fm.filter(pl.col("split") == "validation")
    evals = {s: fm.filter(pl.col("split") == s) for s in ("validation", "test", "holdout")}
    print(f"train {train.height:,} | valid {valid.height:,} | test {evals['test'].height:,} | holdout {evals['holdout'].height:,}")

    # Arm C = gradient boosting on attributes, roll values and location (no as-of market / neighbourhood lag features);
    # arm D = AVM with the full spatial-temporal feature set + quantile models + conformal scaling.
    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, "synthetic": a.synthetic, "python": platform.python_version(), "machine": platform.machine(), "arms": {}}
    for arm, ctor in models.items():
        t0 = time.perf_counter()
        m = ctor().fit(train, valid)
        fit_s = time.perf_counter() - t0
        info = {"fit_seconds": round(fit_s, 1)}
        for split, df in evals.items():
            if df.height == 0:
                continue
            t1 = time.perf_counter()
            pred = m.predict(df)
            records.append(pred.to_records(split, arm, run_id, latency_s=time.perf_counter() - t1))
        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}")
    # spatial holdout: retrain A, C, D without the withheld regions; predict test rows of those regions
    sp_train = train.filter(~pl.col("region_name").is_in(spatial_regions))
    sp_valid = valid.filter(~pl.col("region_name").is_in(spatial_regions))
    sp_test = evals["test"].filter(pl.col("region_name").is_in(spatial_regions))
    if sp_test.height > 0 and sp_train.height > 1000:
        for arm in ("A_hedonic", "C_gbm", "D_avm"):
            m = models[arm]().fit(sp_train, sp_valid)
            records.append(m.predict(sp_test).to_records("test", f"{arm}_spatial", run_id))
        manifest["spatial_holdout"] = {"regions": spatial_regions, "n_test": sp_test.height, "n_train": sp_train.height}
    preds = pl.concat(records)
    preds.write_parquet(out / "predictions.parquet")
    (out / "manifest.json").write_text(json.dumps(manifest, indent=1, default=float))
    # quick in-sample-free summary (validation/test mean abs log error) for the console
    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), pl.len()).sort(["split", "arm"])
    print(s)


if __name__ == "__main__":
    main()