#!/usr/bin/env python3
"""
train_xgb_cpp.py — train the scene-boundary XGBoost on the C++-EXTRACTED feature
matrices (experiments/dumps/cpp_features/<slug>.h5, written by scene_features_dump).

This is the parity-by-construction path: the model is fit on exactly the features
the C++ XGBSceneBoundary produces at inference, so C++ boundaries match by
construction — no numpy-vs-C++ feature drift to chase. Same soft Gaussian target,
knee threshold, and ±20s eval as train_xgb_boundary.py.

Usage (train all 9 + save shipped model):
  .venv-rocm/bin/python scripts/scene_detector/train_xgb_cpp.py --train-all
"""
from __future__ import annotations
import argparse, json, sys
from pathlib import Path
import numpy as np, h5py
sys.path.insert(0, "scripts/scene_detector")
from train_scene_boundary import load_xray_boundaries, nms_peaks
from train_xgb_boundary import knee_boundaries, prf, SIGMA
import xgboost as xgb

CPP_DIR = "experiments/dumps/cpp_features"


def load(slug, xray):
    with h5py.File(f"{CPP_DIR}/{slug}.h5") as f:
        X = f["features"][:].astype(np.float32)
        ts = f["timestamp_sec"][:]
    b = np.array(load_xray_boundaries(xray))
    y = np.zeros(len(ts), np.float32)
    if len(b):
        for i, t in enumerate(ts):
            y[i] = np.exp(-((np.min(np.abs(b - t)))/SIGMA)**2)
    yb = np.zeros(len(ts), np.float32)
    for bb in b:
        yb[np.abs(ts - bb) <= 2.0] = 1.0
    return X, y, yb, ts


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", default="experiments/manifests/films_LVFace_opencv5.json")
    ap.add_argument("--holdout", nargs="*", default=[])
    ap.add_argument("--train-all", action="store_true")
    ap.add_argument("--tol", type=int, default=20)
    ap.add_argument("--out", default="experiments/results/scene_boundary")
    args = ap.parse_args()
    films = json.load(open(args.manifest))
    Path(args.out).mkdir(parents=True, exist_ok=True)

    tr = films if args.train_all else [f for f in films if f["slug"] not in args.holdout]
    Xtr = np.concatenate([load(f["slug"], f["xray"])[0] for f in tr])
    ytr = np.concatenate([load(f["slug"], f["xray"])[1] for f in tr])
    reg = xgb.XGBRegressor(n_estimators=400, max_depth=5, learning_rate=0.05,
                           subsample=0.8, colsample_bytree=0.8,
                           objective="reg:squarederror", n_jobs=8, tree_method="hist")
    reg.fit(Xtr, ytr)
    print(f"[xgb-cpp] trained on {len(tr)} films", file=sys.stderr)

    ev = films if args.train_all else [f for f in films if f["slug"] in args.holdout]
    tag = "TRAIN-FIT (all 9)" if args.train_all else "HELD-OUT"
    print(f"\n=== {tag} (C++ features, knee, ±{args.tol}s) ===")
    print(f"{'film':26s} {'TP':>4}{'FP':>5}{'FN':>5} {'P':>5}{'R':>5}{'F1':>5}")
    f1s = []
    for f in ev:
        X, y, yb, ts = load(f["slug"], f["xray"])
        prob = np.clip(reg.predict(X), 0, 1)
        peaks = knee_boundaries(prob)
        Tset = np.where(yb > 0.5)[0]
        P, R, F, tp, fp, fn = prf(peaks, Tset, args.tol)
        f1s.append(F)
        print(f"{f['slug'][:26]:26s} {tp:>4}{fp:>5}{fn:>5} {P*100:4.0f}%{R*100:4.0f}%{F*100:4.0f}%")
    print(f"\nmacro-F1: {np.mean(f1s)*100:.1f}%")
    if args.train_all:
        reg.save_model(str(Path(args.out) / "xgb_boundary_cpp.json"))
        print(f"[xgb-cpp] shipped model → {args.out}/xgb_boundary_cpp.json", file=sys.stderr)


if __name__ == "__main__":
    main()