#!/usr/bin/env python3
"""
train_xgb_boundary.py — SHIPPED scene-boundary detector.

An XGBoost regressor over a ±WIN-second window of delta features predicts a soft
Gaussian proximity-to-boundary target; a per-film KNEE threshold on the predicted
peak heights selects the boundaries (self-calibrates the count without a magic
rate). Evaluated with NMS + P/R/F1 at ±20 s tolerance (X-Ray scenes are ~170 s,
so ±20 s placement is what flood-fill actually needs).

Why this shape (all measured, see docs/scene-detector):
  - DELTA features, not raw histogram/PSD: the raw content dilutes; |Δ| separates
    boundaries 4-5x. Audio is weak but included (XGBoost ignores what it can't use).
  - SOFT target exp(-(d/σ)²), σ=10s: a near-miss is trained as near-correct, not a
    hard negative. Regression → smooth score surface → NMS peaks.
  - KNEE threshold per film: peak-height curve has a knee where real boundaries
    give way to noise; picking it matches the true scene count without a global
    threshold that's wrong for every grade.
  - Café Society + Scarface (low-contrast grades) MUST be in training; held out,
    the model can't generalize to them. The shipped model trains on ALL 9.

Honest generalization: leave-one-out CV ≈ 26% F1 @±10s / ~34% @±20s. The shipped
all-9 model is what deployment uses (max grade coverage); LOO is the number to
quote for a brand-new film.

Usage (train on all 9 + save shipped model):
  .venv-rocm/bin/python scripts/scene_detector/train_xgb_boundary.py --train-all
Usage (held-out eval):
  ... --holdout Sound_of_Metal The_Many_Saints_of_Newark Valerian_...
"""
from __future__ import annotations
import argparse, json, sys
from pathlib import Path
import numpy as np
import h5py

sys.path.insert(0, "scripts/scene_detector")
from train_scene_boundary import nms_peaks, load_xray_boundaries, SCENE_TAU, TOL_SEC
from train_scene_boundary import video_features, audio_features, build_film
from scipy.signal import find_peaks
import xgboost as xgb

WIN = 3          # ±WIN-second context window
SIGMA = 10.0     # soft-target Gaussian width (seconds)


def per_second_matrix(dump, xray, audio_dir, win=None):
    """Windowed delta features + debounce clock → (X[T,F], y_binary[T], is_cut[T])."""
    V, A, y, is_cut, ts = build_film(dump, xray, audio_dir)
    base = np.concatenate([V] + ([A] if A is not None else []), 1)
    T, d = base.shape
    sig = V[:, 0]
    thr = np.percentile(sig, 90)
    idx = np.arange(T); peak = np.where(sig > thr, idx, -1)
    last = np.maximum.accumulate(peak)
    dt = (idx - last).astype(np.float32); dt[last < 0] = SCENE_TAU
    clock = np.stack([dt, np.minimum(1, dt/SCENE_TAU), np.exp(-dt/SCENE_TAU)], 1)
    W = WIN if win is None else win
    padded = np.pad(base, ((W, W), (0, 0)), mode="edge")
    wf = np.concatenate([padded[i:i+T] for i in range(2*W+1)], 1)
    return np.concatenate([wf, clock], 1).astype(np.float32), y, is_cut


def soft_target(dump, xray):
    ts = h5py.File(dump)["frames/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)
    return y


def knee_boundaries(prob, min_gap=5):
    """Per-film knee threshold on peak heights → selected peak indices.

    Peaks sorted by height form a convex-decreasing curve; the knee (max drop
    below the endpoints chord) is where real boundaries give way to noise. Returns
    the timestamps (indices) of peaks at or above the knee height."""
    pk, _ = find_peaks(prob, distance=min_gap)
    if len(pk) < 5:
        return list(pk)
    heights = np.sort(prob[pk])[::-1]
    n = len(heights); x = np.arange(n)/(n-1); yv = heights/(heights[0]+1e-9)
    chord = yv[0] + (yv[-1]-yv[0])*x
    k = int(np.argmax(chord - yv))
    thr = heights[k]
    return [int(i) for i in pk if prob[i] >= thr]


def train(films, audio_dir):
    X = np.concatenate([per_second_matrix(f["dump"], f["xray"], audio_dir)[0] for f in films])
    y = np.concatenate([soft_target(f["dump"], f["xray"]) for f in films])
    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(X, y)
    return reg


def prf(peaks, Tset, tol=20):
    if not peaks or len(Tset) == 0:
        return 0., 0., 0., 0, 0, len(Tset)
    tp_p = sum(any(abs(p-t) <= tol for t in Tset) for p in peaks)
    tp_t = sum(any(abs(p-t) <= tol for p in peaks) for t in Tset)
    P = tp_p/len(peaks); R = tp_t/len(Tset)
    return (P, R, (2*P*R/(P+R) if P+R else 0.),
            tp_p, len(peaks)-tp_p, len(Tset)-tp_t)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", default="experiments/manifests/films_LVFace_opencv5.json")
    ap.add_argument("--audio-dir", default="experiments/dumps/audio_features")
    ap.add_argument("--holdout", nargs="*", default=[])
    ap.add_argument("--train-all", action="store_true", help="train on all 9 + save shipped model")
    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]
    reg = train(tr, args.audio_dir)
    print(f"[xgb] 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} boundary detection (knee, NMS, ±{args.tol}s) ===")
    print(f"{'film':26s} {'TP':>4}{'FP':>5}{'FN':>5} {'P':>5}{'R':>5}{'F1':>5}  {'gray F1':>7}")
    rep = {"win": WIN, "sigma": SIGMA, "tol": args.tol, "train_all": args.train_all,
           "holdout": args.holdout, "films": {}}
    f1s, gf1s = [], []
    for f in ev:
        X, yb, ic = per_second_matrix(f["dump"], f["xray"], args.audio_dir)
        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)
        gpk = nms_peaks(ic.astype(float)); _, _, gF, *_ = prf(gpk, Tset, args.tol)
        f1s.append(F); gf1s.append(gF)
        rep["films"][f["slug"]] = {"TP": tp, "FP": fp, "FN": fn, "P": P, "R": R, "F1": F,
                                   "n_pred": len(peaks), "n_true": len(Tset), "gray_F1": gF}
        print(f"{f['slug'][:26]:26s} {tp:>4}{fp:>5}{fn:>5} {P*100:4.0f}%{R*100:4.0f}%"
              f"{F*100:4.0f}%  {gF*100:5.0f}%")
    print(f"\nmacro-F1: detector {np.mean(f1s)*100:.1f}%   grayscale {np.mean(gf1s)*100:.1f}%")
    rep["macro_f1"] = {"detector": float(np.mean(f1s)), "grayscale": float(np.mean(gf1s))}
    if args.train_all:
        reg.save_model(str(Path(args.out) / "xgb_boundary_shipped.json"))
        print(f"[xgb] shipped model → {args.out}/xgb_boundary_shipped.json", file=sys.stderr)
    json.dump(rep, open(Path(args.out) / "xgb_report.json", "w"), indent=2)


if __name__ == "__main__":
    main()