training/Nexar/nexar_submit.py

#!/usr/bin/env python3
"""
Generate Nexar submission CSV and optionally evaluate against solution.csv.

Modes:
  1. zero_shot  — use P(ALERT) from cached PolicyHead directly (no extra training)
  2. trained    — use fine-tuned NexarTemporalHead/NexarSimpleHead
  3. ensemble   — weighted blend of trained scores + baseline sample_submission scores

Usage:
  # Zero-shot (fastest, no training needed):
  python -m training.Nexar.nexar_submit \
      --mode zero_shot \
      --test_cache   data/nexar_cache/test.pt \
      --out_csv      submissions/nexar_zero_shot.csv

  # Trained model:
  python -m training.Nexar.nexar_submit \
      --mode trained \
      --test_cache   data/nexar_cache/test.pt \
      --model_dir    checkpoints/Nexar/nexar_v1 \
      --out_csv      submissions/nexar_trained.csv

  # Ensemble:
  python -m training.Nexar.nexar_submit \
      --mode ensemble \
      --test_cache   data/nexar_cache/test.pt \
      --model_dir    checkpoints/Nexar/nexar_v1 \
      --baseline_csv nexar-collision-prediction/sample_submission.csv \
      --ensemble_alpha 0.5 \
      --out_csv      submissions/nexar_ensemble.csv \
      --evaluate     NEXAR_COLLISION/solution.csv
"""
from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

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

from training.Nexar.nexar_dataset import NexarTestDataset, nexar_collate_test
from training.Nexar.nexar_model import build_model

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Nexar.submit")


# ── scoring functions ─────────────────────────────────────────────────────────

def scores_zero_shot(cache_file: str, n_windows: int = 3, agg: str = "max_last") -> Dict[str, float]:
    """
    Zero-shot collision scores from cached P(ALERT) values.

    agg strategies:
      last      — P(ALERT) of the last (most recent) window only
      max       — max P(ALERT) over all windows
      max_last  — 0.7 * last + 0.3 * max (emphasises recency)
      weighted  — linearly increasing weights over windows (latest = highest)
      tta       — 1 / (1 + tta_mean_last) combined with p_alert_last
    """
    cache = torch.load(cache_file, map_location="cpu", weights_only=False)
    scores = {}
    for vid_id, feat in cache["features"].items():
        p = feat["p_alert"].float()     # [T]
        tta = feat["tta_means"].float() # [T]

        if agg == "last":
            s = p[-1].item()
        elif agg == "max":
            s = p.max().item()
        elif agg == "max_last":
            s = 0.6 * p[-1].item() + 0.4 * p.max().item()
        elif agg == "weighted":
            T = p.shape[0]
            w = torch.linspace(0.5, 1.0, T)
            w = w / w.sum()
            s = (p * w).sum().item()
        elif agg == "tta":
            # Combine P(ALERT) with recency-adjusted 1/(1+tta)
            tta_score = 1.0 / (1.0 + tta[-1].item())
            s = 0.6 * p[-1].item() + 0.4 * tta_score
        else:
            s = p[-1].item()
        scores[vid_id] = float(np.clip(s, 0, 1))
    return scores


@torch.no_grad()
def scores_trained(
    cache_file: str,
    model_dir:  str,
    batch_size: int = 128,
    n_windows:  int = 3,
) -> Dict[str, float]:
    """Run fine-tuned NexarHead on test features."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    meta_path = Path(model_dir) / "best_meta.json"
    with open(meta_path) as f:
        meta = json.load(f)
    hidden_dim = meta["hidden_dim"]
    arch       = meta["arch"]
    n_windows  = meta.get("n_windows", n_windows)

    model = build_model(hidden_dim, arch).to(device)
    model.load_state_dict(torch.load(Path(model_dir) / "best_model.pt", map_location=device))
    model.eval()

    ds = NexarTestDataset(cache_file, n_windows=n_windows)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False,
                        num_workers=4, collate_fn=nexar_collate_test, pin_memory=True)

    scores = {}
    for batch in tqdm(loader, desc="Scoring"):
        beliefs   = batch["beliefs"].to(device)
        tta_means = batch["tta_means"].to(device)
        tta_vars  = batch["tta_vars"].to(device)
        p_alerts  = batch["p_alerts"].to(device)

        if hasattr(model, "lstm"):
            s = model(beliefs, tta_means, tta_vars, p_alerts)
        else:
            s = model(beliefs[:, -1, :], tta_means[:, -1], tta_vars[:, -1], p_alerts[:, -1])

        for vid_id, score in zip(batch["video_ids"], s.cpu().tolist()):
            scores[vid_id] = float(np.clip(score, 0, 1))
    return scores


def scores_ensemble(
    primary:  Dict[str, float],
    baseline: Dict[str, float],
    alpha:    float = 0.5,
) -> Dict[str, float]:
    """
    Blend primary scores with baseline (sample_submission) scores.
    score = alpha * primary + (1 - alpha) * baseline
    """
    all_ids = set(primary) | set(baseline)
    result = {}
    for vid_id in all_ids:
        p = primary.get(vid_id, 0.5)
        b = baseline.get(vid_id, 0.5)
        result[vid_id] = float(np.clip(alpha * p + (1 - alpha) * b, 0, 1))
    return result


def evaluate_submission(submission: Dict[str, float], solution_csv: str):
    """Compute mAP (Public) and mAP (Private) against solution.csv."""
    from sklearn.metrics import average_precision_score

    sol = pd.read_csv(solution_csv)
    sol["id"] = sol["id"].astype(str).str.zfill(5)
    sub_df = pd.DataFrame(list(submission.items()), columns=["id", "score"])
    sub_df["id"] = sub_df["id"].astype(str).str.zfill(5)

    for usage in ["Public", "Private"]:
        subset = sol[sol["Usage"] == usage].copy()
        merged = subset.merge(sub_df, on="id", how="left").fillna(0.5)
        ap_list = []
        for g in sorted(merged["group"].unique()):
            g_df = merged[merged["group"] == g]
            if g_df["target"].nunique() < 2:
                continue
            ap = average_precision_score(g_df["target"], g_df["score"])
            ap_list.append(ap)
        mean_ap = np.mean(ap_list) if ap_list else float("nan")
        print(f"mAP ({usage}): {mean_ap:.6f}")


def write_submission(scores: Dict[str, float], out_csv: str, test_csv: Optional[str] = None):
    """Write submission CSV. Aligns with test.csv order if provided."""
    out = Path(out_csv)
    out.parent.mkdir(parents=True, exist_ok=True)

    if test_csv and Path(test_csv).exists():
        ids = pd.read_csv(test_csv)["id"].astype(str).str.zfill(5).tolist()
    else:
        ids = sorted(scores.keys())

    rows = [{"id": vid_id, "score": scores.get(vid_id, 0.5)} for vid_id in ids]
    df = pd.DataFrame(rows)
    df.to_csv(out_csv, index=False)
    logger.info(f"Submission saved → {out_csv}  ({len(df)} rows)")

    # Print score distribution
    vals = [r["score"] for r in rows]
    logger.info(
        f"Score stats:  mean={np.mean(vals):.3f}  std={np.std(vals):.3f}  "
        f"min={np.min(vals):.3f}  max={np.max(vals):.3f}  "
        f"p50={np.median(vals):.3f}  p90={np.percentile(vals,90):.3f}"
    )


def main():
    parser = argparse.ArgumentParser("nexar_submit")
    parser.add_argument("--mode", required=True, choices=["zero_shot", "trained", "ensemble"])
    parser.add_argument("--test_cache",      required=True)
    parser.add_argument("--model_dir",       default=None)
    parser.add_argument("--baseline_csv",    default=None,
                        help="sample_submission.csv for ensemble mode")
    parser.add_argument("--ensemble_alpha",  type=float, default=0.5,
                        help="Weight for primary model (0=baseline only, 1=model only)")
    parser.add_argument("--zero_shot_agg",   default="max_last",
                        choices=["last", "max", "max_last", "weighted", "tta"])
    parser.add_argument("--n_windows",       type=int, default=3)
    parser.add_argument("--out_csv",         required=True)
    parser.add_argument("--test_csv",        default="nexar-collision-prediction/test.csv")
    parser.add_argument("--evaluate",        default=None,
                        help="Path to solution.csv for local evaluation")
    args = parser.parse_args()

    # ── compute scores ────────────────────────────────────────────────────────
    if args.mode == "zero_shot":
        logger.info(f"Zero-shot inference (agg={args.zero_shot_agg}) ...")
        primary_scores = scores_zero_shot(args.test_cache, args.n_windows, args.zero_shot_agg)

    elif args.mode == "trained":
        if not args.model_dir:
            parser.error("--model_dir required for mode=trained")
        logger.info("Running trained NexarHead ...")
        primary_scores = scores_trained(args.test_cache, args.model_dir,
                                        n_windows=args.n_windows)

    elif args.mode == "ensemble":
        if not args.model_dir:
            parser.error("--model_dir required for mode=ensemble")
        logger.info("Running trained NexarHead + ensemble ...")
        trained = scores_trained(args.test_cache, args.model_dir, n_windows=args.n_windows)

        baseline = {}
        if args.baseline_csv and Path(args.baseline_csv).exists():
            b_df = pd.read_csv(args.baseline_csv)
            for _, row in b_df.iterrows():
                vid_id = str(row["id"]).zfill(5)
                baseline[vid_id] = float(row["score"])
            logger.info(f"Baseline scores loaded: {len(baseline)} entries")

        primary_scores = scores_ensemble(trained, baseline, args.ensemble_alpha)

    # ── write output ──────────────────────────────────────────────────────────
    write_submission(primary_scores, args.out_csv, args.test_csv)

    # ── optional local evaluation ─────────────────────────────────────────────
    if args.evaluate:
        logger.info(f"\nLocal evaluation against {args.evaluate}:")
        evaluate_submission(primary_scores, args.evaluate)


if __name__ == "__main__":
    main()