training/Policy/observe_analysis.py

#!/usr/bin/env python3
"""
OBSERVE Temporal Analysis — 论文核心证据脚本

目的：证明 OBSERVE 类有真正的预警价值，即：
  1. OBSERVE 在 ALERT 之前触发（有统计显著的时间提前量）
  2. OBSERVE→ALERT 的转变顺序是可靠的（不是随机噪声）
  3. 有 OBSERVE 预警的视频比没有的更早检测到碰撞

输出内容：
  - observe_lead_time_stats.json：各类视频的 OBSERVE 提前量统计
  - transition_matrix.json：SILENT→OBSERVE→ALERT 转变频率矩阵
  - observe_analysis_plot.png：时间轴分析图（如果有 matplotlib）

使用方法：
  python -m training.Policy.observe_analysis \
      --sft_checkpoint    checkpoints/SFT/sft_v2/best \
      --policy_checkpoint checkpoints/Policy/policy_warmstart_v3/best \
      --label_dir         data/policy_labels \
      --belief_cache_dir  data/belief_cache \
      --output_dir        eval_results/observe_analysis
"""
from __future__ import annotations

import argparse
import json
import logging
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
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.Policy.policy_model import PolicyModel
from training.Policy.policy_dataset import PolicyDataset, policy_collate_fn

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

SILENT  = 0
OBSERVE = 1
ALERT   = 2
ACTION_NAMES = {0: "SILENT", 1: "OBSERVE", 2: "ALERT"}


@torch.no_grad()
def run_inference(
    model:       PolicyModel,
    loader:      DataLoader,
    device:      torch.device,
) -> List[dict]:
    """
    Run model on all val samples, return per-sample results.

    Returns list of dicts with:
      video_id, category, tta_raw, true_label, pred_label, probs [3]
    """
    model.eval()
    results = []

    for batch in tqdm(loader, desc="Inference"):
        if "beliefs" in batch:
            logits = model.forward_cached(
                batch["beliefs"].to(device),
                batch["tta_means"].to(device),
                batch["tta_vars"].to(device),
            )
        else:
            logits = model(batch["images"], batch["metadata"])

        probs = F.softmax(logits, dim=-1).cpu().numpy()
        preds = logits.argmax(dim=-1).cpu().numpy()

        for i in range(len(batch["action_labels"])):
            results.append({
                "video_id":   batch["video_ids"][i],
                "category":   batch["categories"][i],
                "tta_raw":    float(batch["tta_raws"][i]),
                "true_label": int(batch["action_labels"][i]),
                "pred_label": int(preds[i]),
                "p_silent":   float(probs[i][0]),
                "p_observe":  float(probs[i][1]),
                "p_alert":    float(probs[i][2]),
            })

    return results


def group_by_video(results: List[dict]) -> Dict[str, List[dict]]:
    """Group samples by video_id, sorted by tta_raw descending (far→near collision)."""
    by_video: Dict[str, List[dict]] = defaultdict(list)
    for r in results:
        by_video[r["video_id"]].append(r)
    # Sort each video's windows: tta_raw descending = far from collision first
    for vid in by_video:
        by_video[vid].sort(key=lambda x: -x["tta_raw"])
    return by_video


def compute_observe_lead_time(video_windows: List[dict]) -> dict:
    """
    For a single video's windows (ordered far→near collision):
    Find when OBSERVE first fires vs when ALERT first fires.

    Returns dict with timing info.
    """
    preds = [w["pred_label"] for w in video_windows]
    ttas  = [w["tta_raw"]    for w in video_windows]

    # Find first OBSERVE and first ALERT (predicted)
    first_observe_tta = None
    first_alert_tta   = None

    for pred, tta in zip(preds, ttas):
        if pred == OBSERVE and first_observe_tta is None:
            first_observe_tta = tta
        if pred == ALERT and first_alert_tta is None:
            first_alert_tta = tta

    has_observe = first_observe_tta is not None
    has_alert   = first_alert_tta   is not None

    lead_time = None
    if has_observe and has_alert and first_observe_tta > first_alert_tta:
        # OBSERVE fires before ALERT (correct temporal order)
        lead_time = first_observe_tta - first_alert_tta

    return {
        "has_observe":        has_observe,
        "has_alert":          has_alert,
        "first_observe_tta":  first_observe_tta,
        "first_alert_tta":    first_alert_tta,
        "observe_before_alert": lead_time is not None,
        "observe_lead_time_s":  lead_time,
        "n_windows":          len(video_windows),
        "category":           video_windows[0]["category"],
    }


def compute_transition_matrix(results: List[dict]) -> np.ndarray:
    """
    Compute transition matrix T[i,j] = fraction of (window_t, window_{t+1}) pairs
    where true label changes from i to j, grouped by video.

    Shows the natural progression: SILENT→OBSERVE→ALERT
    """
    counts = np.zeros((3, 3), dtype=float)
    by_video = group_by_video(results)

    for vid, windows in by_video.items():
        true_labels = [w["true_label"] for w in windows]
        for t in range(len(true_labels) - 1):
            i, j = true_labels[t], true_labels[t+1]
            counts[i, j] += 1

    # Normalise rows
    row_sums = counts.sum(axis=1, keepdims=True).clip(min=1)
    return counts / row_sums


def compute_prediction_transition_matrix(results: List[dict]) -> np.ndarray:
    """Same but for predicted labels — shows what the model actually does."""
    counts = np.zeros((3, 3), dtype=float)
    by_video = group_by_video(results)
    for vid, windows in by_video.items():
        preds = [w["pred_label"] for w in windows]
        for t in range(len(preds) - 1):
            counts[preds[t], preds[t+1]] += 1
    row_sums = counts.sum(axis=1, keepdims=True).clip(min=1)
    return counts / row_sums


def compute_tta_bins(results: List[dict], bins: List[Tuple[float, float]]) -> dict:
    """
    Per TTA bin: P(pred=OBSERVE), P(pred=ALERT) for ego_collision videos.

    Answers: "at X seconds before collision, what fraction of windows are OBSERVE vs ALERT?"
    """
    ego = [r for r in results if r["category"] in ("ego_collision", "ego_positive")]
    out = {}
    for lo, hi in bins:
        in_bin = [r for r in ego if lo <= r["tta_raw"] < hi]
        if not in_bin:
            continue
        n = len(in_bin)
        label = f"{lo:.1f}-{hi:.1f}s"
        out[label] = {
            "n":          n,
            "tta_range":  [lo, hi],
            "p_silent":   float(np.mean([r["pred_label"] == SILENT  for r in in_bin])),
            "p_observe":  float(np.mean([r["pred_label"] == OBSERVE for r in in_bin])),
            "p_alert":    float(np.mean([r["pred_label"] == ALERT   for r in in_bin])),
            "true_silent":  float(np.mean([r["true_label"] == SILENT  for r in in_bin])),
            "true_observe": float(np.mean([r["true_label"] == OBSERVE for r in in_bin])),
            "true_alert":   float(np.mean([r["true_label"] == ALERT   for r in in_bin])),
        }
    return out


def print_report(
    results:     List[dict],
    lead_stats:  dict,
    trans_true:  np.ndarray,
    trans_pred:  np.ndarray,
    tta_bins:    dict,
):
    n_total = len(results)
    n_ego   = sum(1 for r in results if r["category"] in ("ego_collision", "ego_positive"))

    print("\n" + "="*60)
    print("  OBSERVE TEMPORAL ANALYSIS REPORT")
    print("="*60)

    print(f"\n[Dataset]")
    print(f"  Total windows    : {n_total}")
    print(f"  Ego-collision    : {n_ego}  ({100*n_ego/max(n_total,1):.1f}%)")

    # Prediction distribution
    preds = [r["pred_label"] for r in results]
    for k, name in ACTION_NAMES.items():
        frac = np.mean([p == k for p in preds])
        print(f"  pred={name:<8}: {100*frac:.1f}%")

    print(f"\n[OBSERVE Lead Time — ego-collision videos with ≥2 windows]")
    ego_v = lead_stats.get("ego_videos", {})
    n_v = len(ego_v)
    has_obs   = sum(1 for v in ego_v.values() if v["has_observe"])
    obs_first = sum(1 for v in ego_v.values() if v["observe_before_alert"])
    lead_times = [v["observe_lead_time_s"] for v in ego_v.values()
                  if v["observe_lead_time_s"] is not None]

    print(f"  Ego videos       : {n_v}")
    print(f"  Has OBSERVE      : {has_obs} ({100*has_obs/max(n_v,1):.1f}%)")
    print(f"  OBSERVE before ALERT: {obs_first} ({100*obs_first/max(n_v,1):.1f}%)")
    if lead_times:
        print(f"  Lead time  mean  : {np.mean(lead_times):.2f}s")
        print(f"  Lead time  median: {np.median(lead_times):.2f}s")
        print(f"  Lead time  p75   : {np.percentile(lead_times, 75):.2f}s")
        print(f"  Lead time  max   : {np.max(lead_times):.2f}s")
        print(f"  ★ On average, OBSERVE fires {np.mean(lead_times):.2f}s BEFORE ALERT")
    else:
        print("  (no valid lead-time observations)")

    print(f"\n[True Label Transition Matrix — P(label_t+1 | label_t)]")
    print(f"  Rows = current state, Cols = next state")
    print(f"  {'':10} SILENT   OBSERVE   ALERT")
    for i, name in ACTION_NAMES.items():
        row = "  ".join([f"{trans_true[i,j]:.3f}" for j in range(3)])
        print(f"  {name:<10} {row}")

    print(f"\n[Predicted Transition Matrix — what the model does]")
    print(f"  {'':10} SILENT   OBSERVE   ALERT")
    for i, name in ACTION_NAMES.items():
        row = "  ".join([f"{trans_pred[i,j]:.3f}" for j in range(3)])
        print(f"  {name:<10} {row}")

    print(f"\n[OBSERVE Rate vs TTA (ego-collision windows)]")
    print(f"  {'TTA range':<12} {'n':>5} {'P(SILENT)':>10} {'P(OBSERVE)':>11} {'P(ALERT)':>9}")
    for label, d in sorted(tta_bins.items(), key=lambda x: -x[1]["tta_range"][0]):
        print(f"  {label:<12} {d['n']:>5} {d['p_silent']:>10.3f} {d['p_observe']:>11.3f} {d['p_alert']:>9.3f}")

    print("="*60)


def main():
    parser = argparse.ArgumentParser("observe_analysis")
    parser.add_argument("--sft_checkpoint",    required=True)
    parser.add_argument("--policy_checkpoint", required=True)
    parser.add_argument("--label_dir",         default="data/policy_labels")
    parser.add_argument("--belief_cache_dir",  default=None)
    parser.add_argument("--split",             default="val")
    parser.add_argument("--output_dir",        default="eval_results/observe_analysis")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Load model
    model = PolicyModel(args.sft_checkpoint, use_bf16=True)
    model.load_policy_checkpoint(args.policy_checkpoint)
    model.eval()

    # Load data
    cache_path = None
    if args.belief_cache_dir:
        p = Path(args.belief_cache_dir) / f"{args.split}.pt"
        if p.exists():
            cache_path = p

    ds = PolicyDataset(
        manifests=[Path(args.label_dir) / f"{args.split}.json"],
        split=args.split,
        belief_cache_path=cache_path,
    )
    loader = DataLoader(ds, batch_size=512, shuffle=False,
                        num_workers=4, collate_fn=policy_collate_fn)

    # Run inference
    results = run_inference(model, loader, device)
    logger.info(f"Inference done: {len(results)} samples")

    # ── OBSERVE lead time analysis ─────────────────────────────────────────────
    by_video = group_by_video(results)
    ego_videos = {vid: compute_observe_lead_time(windows)
                  for vid, windows in by_video.items()
                  if windows[0]["category"] in ("ego_collision", "ego_positive")
                  and len(windows) >= 2}

    lead_stats = {"ego_videos": ego_videos}

    # ── Transition matrices ────────────────────────────────────────────────────
    trans_true = compute_transition_matrix(results)
    trans_pred = compute_prediction_transition_matrix(results)

    # ── TTA bins: how predictions change as collision approaches ──────────────
    tta_bins_def = [(i, i+1) for i in range(0, 10)] + [(0, 2), (2, 5), (5, 10)]
    tta_bins = compute_tta_bins(results, tta_bins_def)

    # ── Print + save ───────────────────────────────────────────────────────────
    print_report(results, lead_stats, trans_true, trans_pred, tta_bins)

    lead_times = [v["observe_lead_time_s"] for v in ego_videos.values()
                  if v["observe_lead_time_s"] is not None]

    summary = {
        "n_samples":          len(results),
        "n_ego_videos":       len(ego_videos),
        "observe_fires_pct":  float(np.mean([v["has_observe"] for v in ego_videos.values()])),
        "observe_before_alert_pct": float(np.mean([v["observe_before_alert"] for v in ego_videos.values()])),
        "lead_time_mean_s":   float(np.mean(lead_times)) if lead_times else 0.0,
        "lead_time_median_s": float(np.median(lead_times)) if lead_times else 0.0,
        "lead_time_p75_s":    float(np.percentile(lead_times, 75)) if lead_times else 0.0,
        "transition_true":    trans_true.tolist(),
        "transition_pred":    trans_pred.tolist(),
        "tta_bins":           tta_bins,
    }
    out_json = out_dir / "observe_analysis.json"
    with open(out_json, "w") as f:
        json.dump(summary, f, indent=2)
    logger.info(f"\nResults saved → {out_json}")

    # ── Optional plot ──────────────────────────────────────────────────────────
    try:
        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches

        # Plot 1: P(OBSERVE) and P(ALERT) vs TTA
        tta_sorted = sorted(tta_bins.items(), key=lambda x: x[1]["tta_range"][0])
        labels_x = [d["tta_range"][0] for _, d in tta_sorted]
        p_obs   = [d["p_observe"] for _, d in tta_sorted]
        p_alert = [d["p_alert"]   for _, d in tta_sorted]
        p_sil   = [d["p_silent"]  for _, d in tta_sorted]

        fig, axes = plt.subplots(1, 2, figsize=(14, 5))

        ax = axes[0]
        ax.fill_between(labels_x, p_sil,    alpha=0.4, color="blue",   label="P(SILENT)")
        ax.fill_between(labels_x, p_obs,    alpha=0.4, color="orange", label="P(OBSERVE)")
        ax.fill_between(labels_x, p_alert,  alpha=0.4, color="red",    label="P(ALERT)")
        ax.plot(labels_x, p_obs,    "o-", color="orange", lw=2)
        ax.plot(labels_x, p_alert,  "s-", color="red",    lw=2)
        ax.set_xlabel("Time to Collision (seconds)", fontsize=12)
        ax.set_ylabel("Prediction Probability", fontsize=12)
        ax.set_title("LKAlert Policy: Prediction Distribution vs TTA\n(ego-collision videos)", fontsize=12)
        ax.legend(fontsize=11)
        if labels_x:
            ax.set_xlim(max(labels_x), 0)   # right side = near collision
        ax.set_ylim(0, 1)
        ax.axvline(x=0, color="red", ls="--", alpha=0.5, label="Collision")
        ax.grid(True, alpha=0.3)

        # Plot 2: OBSERVE lead time histogram
        ax2 = axes[1]
        if lead_times:
            ax2.hist(lead_times, bins=15, color="steelblue", edgecolor="white", alpha=0.8)
            ax2.axvline(np.mean(lead_times), color="red", ls="--", lw=2,
                        label=f"Mean={np.mean(lead_times):.2f}s")
            ax2.axvline(np.median(lead_times), color="orange", ls="--", lw=2,
                        label=f"Median={np.median(lead_times):.2f}s")
            ax2.set_xlabel("OBSERVE Lead Time Before ALERT (seconds)", fontsize=12)
            ax2.set_ylabel("Count", fontsize=12)
            ax2.set_title("OBSERVE Pre-Warning Lead Time\n(ego-collision videos)", fontsize=12)
            ax2.legend(fontsize=11)
            ax2.grid(True, alpha=0.3)

        plt.tight_layout()
        plot_path = out_dir / "observe_analysis.png"
        plt.savefig(plot_path, dpi=150, bbox_inches="tight")
        logger.info(f"Plot saved → {plot_path}")
        plt.close()
    except ImportError:
        logger.warning("matplotlib not available — skipping plot generation")


if __name__ == "__main__":
    main()