experiment-2107/scripts/temporal_scripts/run_experiment.py

"""Combo-decomposition experiment: train CTM, BiLSTM, and xLSTM on FS-Jump3D singles,
then infer on the 13 held-out Comb clips with multiple strategies.

Phase 1: Data prep (reuses pipeline.py's existing logic)
Phase 2: Train 3 architectures in parallel
Phase 3: Sliding-window inference at multiple window/stride settings
Phase 4: Per-frame scoring (timestep-level predictions from each architecture)
Phase 5: Confidence-thresholded analysis

Usage:
    python -m temporal_scripts.run_experiment
"""

from __future__ import annotations

import subprocess
import sys
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import config as dataset_config
import labels
import pipeline
import preprocessing

SINGLES_DIR = Path("../data/processed_fs_jump3d_singles_temporal")
COMBO_BUNDLE_PATH = SINGLES_DIR / "combo_holdout.pkl"
LOG_DIR = SINGLES_DIR / "logs"

WINDOW_CONFIGS = [
    {"window": 200, "stride": 10, "label": "200/10 (short, fine)"},
    {"window": 280, "stride": 20, "label": "280/20 (baseline)"},
    {"window": 400, "stride": 30, "label": "400/30 (long, coarse)"},
]

CONFIDENCE_THRESHOLD = 0.4

SCRIPTS_DIR = Path(__file__).resolve().parent.parent


# ---------------------------------------------------------------------------
# Phase 1: Data prep
# ---------------------------------------------------------------------------

def preprocess_full_sequence(sample: dict, config: dict) -> np.ndarray:
    width = float(sample.get("width", config["default_width"]))
    height = float(sample.get("height", config["default_height"]))
    fps = float(sample.get("fps", config["fps"]))
    stride = max(1, int(sample.get("extract_every_n_frames", 1) or 1))
    effective_fps = fps / stride

    skeleton = preprocessing.correct_aspect_ratio(sample["skeleton"], width, height)
    skeleton = preprocessing.filter_low_confidence(skeleton, config["confidence_threshold"])
    skeleton = preprocessing.interpolate_missing(skeleton)
    skeleton = preprocessing.smooth_skeleton(
        skeleton, config["smoothing_window"], config["smoothing_polyorder"]
    )
    skeleton = preprocessing.center_on_hips(skeleton)
    skeleton = preprocessing.normalize_scale(skeleton)
    return preprocessing.build_feature_tensor(skeleton, effective_fps)


def prepare_data(source_config: dict) -> list[dict]:
    print("=" * 72)
    print("PHASE 1: data prep")
    print("=" * 72)

    all_samples = pipeline.load_source_samples("fs_jump3d", source_config)
    print(f"loaded {len(all_samples)} fs_jump3d samples total")

    singles_features, singles_labels, singles_names = [], [], []
    combo_samples = []
    for sample in all_samples:
        label_name, _ = pipeline.map_label_for_source("fs_jump3d", sample["label"])
        if label_name == "Comb":
            combo_samples.append(sample)
            continue
        features, _, _ = pipeline.process_sample(sample, "fs_jump3d", source_config)
        singles_features.append(features)
        singles_labels.append(labels.FS_JUMP3D_SINGLES_LABEL_TO_IDX[label_name])
        singles_names.append(label_name)

    print(f"singles: {len(singles_features)}  |  held-out Comb: {len(combo_samples)}")
    assert len(combo_samples) == 13, f"expected 13 Comb clips, found {len(combo_samples)}"

    features = np.stack(singles_features, axis=0)
    y = np.asarray(singles_labels, dtype=np.int64)

    save_config = dict(source_config)
    save_config["output_dir"] = SINGLES_DIR
    pipeline.save_outputs(features, y, singles_names, save_config,
                          taxonomy=labels.FS_JUMP3D_SINGLES_TAXONOMY)
    print(f"saved singles-only train/val/test -> {SINGLES_DIR.resolve()}")

    combo_bundle = []
    for sample in combo_samples:
        full_features = preprocess_full_sequence(sample, source_config)
        combo_bundle.append({"source_file": sample["source_file"], "features": full_features})
        print(f"  {Path(sample['source_file']).name}: full sequence shape {full_features.shape}")

    SINGLES_DIR.mkdir(parents=True, exist_ok=True)
    pipeline.save_pickle(combo_bundle, COMBO_BUNDLE_PATH)
    print(f"saved {len(combo_bundle)} held-out Comb sequences -> {COMBO_BUNDLE_PATH.resolve()}")
    return combo_bundle


# ---------------------------------------------------------------------------
# Phase 2: Parallel training
# ---------------------------------------------------------------------------

def launch_training_jobs() -> dict[str, subprocess.Popen]:
    print("\n" + "=" * 72)
    print("PHASE 2: training BiLSTM, CTM-GCN, and xLSTM in parallel")
    print("=" * 72)
    LOG_DIR.mkdir(parents=True, exist_ok=True)

    jobs = {
        "bilstm": (
            f"import sys; sys.path.insert(0, '{SCRIPTS_DIR}'); "
            f"import model_bilstm; from pathlib import Path; "
            f"model_bilstm.train(Path('{SINGLES_DIR}'), coarse=False)"
        ),
        "gcn_ctm": (
            f"import sys; sys.path.insert(0, '{SCRIPTS_DIR}'); "
            f"import model_ctm; from pathlib import Path; "
            f"model_ctm.train(Path('{SINGLES_DIR}'), coarse=False, backbone='gcn', tag='gcn_ctm')"
        ),
        "xlstm": (
            f"import sys; sys.path.insert(0, '{SCRIPTS_DIR}'); "
            f"sys.path.insert(0, '{SCRIPTS_DIR / 'temporal_scripts'}'); "
            f"from model_xlstm import train; from pathlib import Path; "
            f"train(Path('{SINGLES_DIR}'), coarse=False)"
        ),
    }

    procs = {}
    for name, code in jobs.items():
        log_path = LOG_DIR / f"{name}.log"
        log_file = open(log_path, "w")
        print(f"  launching {name} -> log: {log_path.resolve()}")
        procs[name] = subprocess.Popen(
            [sys.executable, "-c", code],
            stdout=log_file, stderr=subprocess.STDOUT,
            cwd=str(SCRIPTS_DIR),
        )
    return procs


def wait_for_jobs(procs: dict[str, subprocess.Popen]) -> dict[str, bool]:
    results = {}
    for name, proc in procs.items():
        ret = proc.wait()
        ok = ret == 0
        results[name] = ok
        status = "OK" if ok else f"FAILED (exit {ret})"
        print(f"  {name}: {status}  (see {LOG_DIR / (name + '.log')})")
        if not ok:
            tail = (LOG_DIR / f"{name}.log").read_text().splitlines()[-30:]
            print("  --- last lines ---")
            print("\n".join(f"    {line}" for line in tail))
    return results


# ---------------------------------------------------------------------------
# Phase 3: Sliding-window inference
# ---------------------------------------------------------------------------

def sliding_windows(num_frames: int, window: int, stride: int) -> list[tuple[int, int]]:
    if num_frames <= window:
        return [(0, num_frames)]
    spans = []
    start = 0
    while start + window <= num_frames:
        spans.append((start, start + window))
        start += stride
    if spans[-1][1] < num_frames:
        spans.append((num_frames - window, num_frames))
    return spans


def load_model_bilstm(device):
    from model_bilstm import SkatingBiLSTMClassifier
    ckpt = torch.load(SINGLES_DIR / "model_bilstm.pt", map_location=device, weights_only=False)
    model = SkatingBiLSTMClassifier(ckpt["in_features"], ckpt["num_classes"]).to(device)
    model.load_state_dict(ckpt["state_dict"])
    model.eval()
    mu = ckpt["feature_mean"].reshape(-1)
    sd = ckpt["feature_std"].reshape(-1)
    return model, mu, sd


def load_model_ctm(device):
    from model_ctm import SkatingCTM
    ckpt_path = SINGLES_DIR / "model_ctm_gcn_ctm_gcnbackbone.pt"
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    model = SkatingCTM(
        in_features=ckpt["in_features"], num_classes=ckpt["num_classes"],
        backbone=ckpt.get("backbone", "gcn"), max_iterations=ckpt["max_iterations"],
        memory_length=ckpt["memory_length"], use_act=ckpt["use_act"],
        node_features=ckpt.get("node_features", "full"),
        gcn_layers=ckpt.get("gcn_layers", 2),
    ).to(device)
    model.load_state_dict(ckpt["state_dict"])
    model.eval()
    mu = ckpt["feature_mean"].reshape(-1)
    sd = ckpt["feature_std"].reshape(-1)
    return model, mu, sd


def load_model_xlstm(device):
    from temporal_scripts.model_xlstm import SkatingxLSTMClassifier
    ckpt = torch.load(SINGLES_DIR / "model_xlstm.pt", map_location=device, weights_only=False)
    model = SkatingxLSTMClassifier(
        ckpt["in_features"], ckpt["num_classes"],
        num_blocks=ckpt.get("num_blocks", 4),
        context_length=ckpt.get("context_length", 128),
    ).to(device)
    model.load_state_dict(ckpt["state_dict"])
    model.eval()
    mu = ckpt["feature_mean"].reshape(-1)
    sd = ckpt["feature_std"].reshape(-1)
    return model, mu, sd


@torch.no_grad()
def classify_window_standard(model, window: np.ndarray, mu, sd, device, taxonomy) -> tuple[str, float]:
    resized = preprocessing.resample_to_length(window, 128).astype(np.float32)
    resized = (resized - mu) / sd
    xb = torch.from_numpy(resized).unsqueeze(0).to(device)
    probs = torch.softmax(model(xb), dim=1)
    pred = int(probs.argmax(1).item())
    return taxonomy[pred], float(probs.max(1).values.item())


@torch.no_grad()
def classify_window_ctm(model, window: np.ndarray, mu, sd, device, taxonomy) -> tuple[str, float]:
    resized = preprocessing.resample_to_length(window, 128).astype(np.float32)
    resized = (resized - mu) / sd
    xb = torch.from_numpy(resized).unsqueeze(0).to(device)
    preds, halt_weights, _, _ = model(xb)
    weighted = (halt_weights.unsqueeze(1) * preds).sum(dim=-1)
    probs = torch.softmax(weighted, dim=1)
    pred = int(probs.argmax(1).item())
    return taxonomy[pred], float(probs.max(1).values.item())


def run_sliding_window_inference(combo_bundle: list[dict], model_zoo: dict,
                                  taxonomy: dict) -> dict:
    print("\n" + "=" * 72)
    print("PHASE 3: sliding-window inference on held-out Comb clips")
    print("=" * 72)

    all_results = {}

    for wcfg in WINDOW_CONFIGS:
        window_size, stride, wlabel = wcfg["window"], wcfg["stride"], wcfg["label"]
        print(f"\n--- Window config: {wlabel} ---")
        config_results = {}

        for sample in combo_bundle:
            name = Path(sample["source_file"]).name
            features = sample["features"]
            spans = sliding_windows(features.shape[0], window_size, stride)
            print(f"\n  {name}  ({features.shape[0]} frames, {len(spans)} window(s))")

            header = f"  {'window':<16}"
            for mname in model_zoo:
                header += f" {mname:<10} {'conf':<6}"
            print(header)

            clip_results = []
            for start, end in spans:
                window = features[start:end]
                row = {"start": start, "end": end}
                row_str = f"  {f'{start:4d}-{end:<4d}':<16}"
                for mname, minfo in model_zoo.items():
                    classify_fn = minfo["classify_fn"]
                    pred, conf = classify_fn(
                        minfo["model"], window, minfo["mu"], minfo["sd"],
                        minfo["device"], taxonomy,
                    )
                    row[f"{mname}_pred"] = pred
                    row[f"{mname}_conf"] = conf
                    row_str += f" {pred:<10} {conf:<6.2f}"
                print(row_str)
                clip_results.append(row)
            config_results[name] = clip_results
        all_results[wlabel] = config_results

    return all_results


# ---------------------------------------------------------------------------
# Phase 4: Per-frame scoring
# ---------------------------------------------------------------------------

@torch.no_grad()
def per_frame_score_bilstm(model, features: np.ndarray, mu, sd, device,
                            taxonomy: dict) -> list[tuple[str, float]]:
    """Run BiLSTM on the full sequence and extract per-timestep hidden states for classification.

    Since the standard BiLSTM model only outputs a single pooled prediction, we access the
    LSTM hidden states before mean-pooling to get per-frame representations, then classify each.
    """
    resized = preprocessing.resample_to_length(features, 128).astype(np.float32)
    resized = (resized - mu) / sd
    xb = torch.from_numpy(resized).unsqueeze(0).to(device)

    x = xb.transpose(1, 2)
    x = torch.relu(model.stem(x))
    x = model.stem_bn(x)
    x = model.cb3(model.cb2(model.cb1(x)))
    x = x.transpose(1, 2)
    x, _ = model.lstm(x)

    B, T, D = x.shape
    flat = x.reshape(B * T, D)
    logits = model.out(model.head(flat)).reshape(B, T, -1)
    probs = torch.softmax(logits, dim=-1)

    results = []
    for t in range(T):
        pred = int(probs[0, t].argmax().item())
        conf = float(probs[0, t].max().item())
        results.append((taxonomy[pred], conf))
    return results


@torch.no_grad()
def per_frame_score_ctm(model, features: np.ndarray, mu, sd, device,
                         taxonomy: dict, window: int = 128, stride: int = 8,
                         ) -> list[tuple[str, float]]:
    """CTM per-frame: the CTM has no frame-indexed output head -- its cross-attention pools
    over the *whole* input window with a single query per thinking step, so one forward pass
    only ever yields one clip-level prediction (this is why the old implementation could only
    repeat a single (label, conf) 128 times: there was nothing else to report). To get a real
    per-frame signal without pretending the architecture has resolution it doesn't, run many
    small overlapping windows across the raw sequence (same sliding-window technique Phase 3
    already uses, just at frame resolution instead of clip resolution) and assign each frame
    the softmax-probability average over every window that covers it.
    """
    T = features.shape[0]
    win = min(window, T)
    spans = sliding_windows(T, win, stride)
    num_classes = len(taxonomy)

    prob_sum = np.zeros((T, num_classes), dtype=np.float64)
    count = np.zeros(T, dtype=np.int64)

    for start, end in spans:
        resized = preprocessing.resample_to_length(features[start:end], 128).astype(np.float32)
        resized = (resized - mu) / sd
        xb = torch.from_numpy(resized).unsqueeze(0).to(device)
        preds, halt_weights, _, _ = model(xb)
        weighted = (halt_weights.unsqueeze(1) * preds).sum(dim=-1)
        probs = torch.softmax(weighted, dim=1).cpu().numpy()[0]
        prob_sum[start:end] += probs
        count[start:end] += 1

    avg_probs = prob_sum / np.maximum(count, 1)[:, None]

    # Resample the raw-frame-resolution predictions to a fixed length of 128 so CTM's output
    # is directly comparable to BiLSTM/xLSTM's 128-step per_frame_score_* output.
    idxs = np.linspace(0, T - 1, 128).round().astype(int)
    results = []
    for i in idxs:
        pred = int(avg_probs[i].argmax())
        conf = float(avg_probs[i].max())
        results.append((taxonomy[pred], conf))
    return results


@torch.no_grad()
def per_frame_score_xlstm(model, features: np.ndarray, mu, sd, device,
                           taxonomy: dict) -> list[tuple[str, float]]:
    resized = preprocessing.resample_to_length(features, 128).astype(np.float32)
    resized = (resized - mu) / sd
    xb = torch.from_numpy(resized).unsqueeze(0).to(device)

    logits = model.forward_per_step(xb)
    probs = torch.softmax(logits, dim=-1)

    results = []
    for t in range(probs.shape[1]):
        pred = int(probs[0, t].argmax().item())
        conf = float(probs[0, t].max().item())
        results.append((taxonomy[pred], conf))
    return results


def run_per_frame_scoring(combo_bundle: list[dict], model_zoo: dict,
                           taxonomy: dict) -> dict:
    print("\n" + "=" * 72)
    print("PHASE 4: per-frame scoring on held-out Comb clips")
    print("=" * 72)

    per_frame_fns = {
        "BiLSTM": per_frame_score_bilstm,
        "CTM": per_frame_score_ctm,
        "xLSTM": per_frame_score_xlstm,
    }

    all_results = {}
    for sample in combo_bundle:
        name = Path(sample["source_file"]).name
        features = sample["features"]
        print(f"\n  {name}  ({features.shape[0]} frames)")

        clip_results = {}
        for mname, minfo in model_zoo.items():
            if mname not in per_frame_fns:
                continue
            fn = per_frame_fns[mname]
            frame_preds = fn(
                minfo["model"], features, minfo["mu"], minfo["sd"],
                minfo["device"], taxonomy,
            )
            clip_results[mname] = frame_preds

            segments = _compress_predictions(frame_preds)
            seg_str = " -> ".join(
                f"{pred}({s}-{e}, conf={c:.2f})"
                for pred, c, s, e in segments
            )
            print(f"    {mname}: {seg_str}")

        all_results[name] = clip_results

    return all_results


def _compress_predictions(preds: list[tuple[str, float]],
                          ) -> list[tuple[str, float, int, int]]:
    """Compress a per-frame prediction list into segments of consecutive same-class frames."""
    if not preds:
        return []
    segments = []
    current_label, current_conf, start = preds[0][0], preds[0][1], 0
    confs = [preds[0][1]]
    for i in range(1, len(preds)):
        if preds[i][0] != current_label:
            segments.append((current_label, sum(confs) / len(confs), start, i))
            current_label = preds[i][0]
            start = i
            confs = []
        confs.append(preds[i][1])
    segments.append((current_label, sum(confs) / len(confs), start, len(preds)))
    return segments


# ---------------------------------------------------------------------------
# Phase 5: Confidence-thresholded analysis
# ---------------------------------------------------------------------------

def run_confidence_analysis(all_sw_results: dict, threshold: float = CONFIDENCE_THRESHOLD):
    print("\n" + "=" * 72)
    print(f"PHASE 5: confidence-thresholded analysis (threshold={threshold})")
    print("=" * 72)

    for wlabel, config_results in all_sw_results.items():
        print(f"\n--- {wlabel} ---")
        for clip_name, rows in config_results.items():
            print(f"\n  {clip_name}:")
            for row in rows:
                for mname_key in [k for k in row if k.endswith("_pred")]:
                    mname = mname_key.replace("_pred", "")
                    pred = row[mname_key]
                    conf = row[f"{mname}_conf"]
                    label = pred if conf >= threshold else "???"
                    marker = " " if conf >= threshold else "*"
                    if marker == "*":
                        print(f"    {row['start']:4d}-{row['end']:<4d} {mname:<8} "
                              f"{label:<10} {conf:.2f} {marker}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> int:
    t0 = time.perf_counter()

    source_config = dataset_config.get_source_config("fs_jump3d", dict(pipeline.CONFIG))

    # Phase 1
    if COMBO_BUNDLE_PATH.exists() and (SINGLES_DIR / "train_features.pkl").exists():
        print("Phase 1: reusing existing data prep")
        combo_bundle = pipeline.load_pickle(COMBO_BUNDLE_PATH)
    else:
        combo_bundle = prepare_data(source_config)

    # Phase 2
    procs = launch_training_jobs()
    job_results = wait_for_jobs(procs)

    if not any(job_results.values()):
        print("\nAll training jobs failed. Aborting.")
        return 1

    # Phase 3-5
    device = "cpu"
    if torch.cuda.is_available():
        device = "cuda"
    taxonomy = labels.FS_JUMP3D_SINGLES_TAXONOMY

    model_zoo = {}
    if job_results.get("bilstm"):
        model, mu, sd = load_model_bilstm(device)
        model_zoo["BiLSTM"] = {
            "model": model, "mu": mu, "sd": sd, "device": device,
            "classify_fn": classify_window_standard,
        }
    if job_results.get("gcn_ctm"):
        model, mu, sd = load_model_ctm(device)
        model_zoo["CTM"] = {
            "model": model, "mu": mu, "sd": sd, "device": device,
            "classify_fn": classify_window_ctm,
        }
    if job_results.get("xlstm"):
        model, mu, sd = load_model_xlstm(device)
        model_zoo["xLSTM"] = {
            "model": model, "mu": mu, "sd": sd, "device": device,
            "classify_fn": classify_window_standard,
        }

    if not model_zoo:
        print("\nNo models loaded. Cannot run inference.")
        return 1

    print(f"\nLoaded {len(model_zoo)} model(s): {list(model_zoo.keys())}")

    # Phase 3
    all_sw_results = run_sliding_window_inference(combo_bundle, model_zoo, taxonomy)

    # Phase 4
    per_frame_results = run_per_frame_scoring(combo_bundle, model_zoo, taxonomy)

    # Phase 5
    run_confidence_analysis(all_sw_results)

    elapsed = time.perf_counter() - t0
    print(f"\n{'=' * 72}")
    print(f"Experiment complete in {elapsed:.1f}s")
    print(f"Models trained: {[k for k, v in job_results.items() if v]}")
    print(f"Models failed: {[k for k, v in job_results.items() if not v]}")
    print(f"Window configs tested: {len(WINDOW_CONFIGS)}")
    print(f"Comb clips analyzed: {len(combo_bundle)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())