tmp/video-velocity-model/data_processing/tools/cal_metric.py

from __future__ import annotations

import argparse
import copy
import importlib.util
import json
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Callable, Iterable

import matplotlib
import numpy as np
import torch
from tqdm import tqdm

matplotlib.use("Agg")
import matplotlib.pyplot as plt

REPO_ROOT = Path(__file__).resolve().parents[1]
RAE_SRC_ROOT = REPO_ROOT / "RAE" / "src"
if str(RAE_SRC_ROOT) not in sys.path:
    sys.path.append(str(RAE_SRC_ROOT))

from disc.lpips import LPIPS  # noqa: E402

PATH_VIDEO_PATH = RAE_SRC_ROOT / "stage2" / "transport" / "path_video.py"
PATH_VIDEO_SPEC = importlib.util.spec_from_file_location("path_video", PATH_VIDEO_PATH)
if PATH_VIDEO_SPEC is None or PATH_VIDEO_SPEC.loader is None:
    raise ImportError(f"Failed to load TrajectoryPlan from {PATH_VIDEO_PATH}")
path_video = importlib.util.module_from_spec(PATH_VIDEO_SPEC)
PATH_VIDEO_SPEC.loader.exec_module(path_video)
TrajectoryPlan = path_video.TrajectoryPlan

try:
    from .forward import (  # noqa: E402
        BACKBONE_SPECS,
        ResolvedRAEConfig,
        batched_indices,
        decode_latents,
        encode_frames,
        frames_to_tensor,
        get_device,
        load_rae,
        load_video_frames,
        resolve_rae_config,
    )
    from .forward_intrapolate import (  # noqa: E402
        collect_video_entries,
        get_interpolator,
        resolve_interpolation_mode,
    )
except ImportError:
    from forward import (  # type: ignore  # noqa: E402
        BACKBONE_SPECS,
        ResolvedRAEConfig,
        batched_indices,
        decode_latents,
        encode_frames,
        frames_to_tensor,
        get_device,
        load_rae,
        load_video_frames,
        resolve_rae_config,
    )
    from forward_intrapolate import (  # type: ignore  # noqa: E402
        collect_video_entries,
        get_interpolator,
        resolve_interpolation_mode,
    )


@dataclass(frozen=True)
class EvalJob:
    backbone: str
    protocol: str
    interp_mode: str


@dataclass(frozen=True)
class VideoEvaluation:
    sample_values: np.ndarray
    pair_count: int
    gap_indices: np.ndarray
    gap_sample_counts: np.ndarray
    gap_ppl_mean: np.ndarray
    gap_ppl_std: np.ndarray
    gap_ppl_min: np.ndarray
    gap_ppl_max: np.ndarray


@dataclass(frozen=True)
class PerGapCurve:
    backbone: str
    protocol: str
    interp_mode_used: str
    video_path: str
    gap_indices: np.ndarray
    gap_sample_counts: np.ndarray
    gap_ppl_mean: np.ndarray
    gap_ppl_std: np.ndarray
    gap_ppl_min: np.ndarray
    gap_ppl_max: np.ndarray


@dataclass(frozen=True)
class EvalResult:
    backbone: str
    protocol: str
    interp_mode_used: str
    videos_evaluated: int
    videos_skipped: int
    pairs_evaluated: int
    samples_evaluated: int
    ppl_mean_raw: float
    ppl_mean_trimmed: float
    ppl_median: float
    ppl_std: float
    ppl_min: float
    ppl_max: float
    trim_value_low: float
    trim_value_high: float
    eps: float
    samples_per_gap: int
    split_name: str | None
    video_path: str | None
    image_size: int
    seed: int


@dataclass(frozen=True)
class JobEvaluation:
    result: EvalResult
    per_gap_curves: list[PerGapCurve]


PPL_INTERP_MODES = ("auto", "linear", "slerp", "smooth_continuous")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Compute StyleGAN-inspired perceptual path length (PPL) on video latent interpolation paths "
            "for MAE / DINOv2 / SigLIP2 RAE backbones."
        )
    )
    parser.add_argument(
        "--config",
        type=Path,
        default=None,
        help="Optional stage-1 or stage-2 YAML. Only supported when evaluating a single backbone.",
    )
    parser.add_argument(
        "--rae-backbone",
        type=str,
        choices=sorted(BACKBONE_SPECS),
        default="dinov2",
        help="RAE encoder backbone used when --config is not supplied.",
    )
    parser.add_argument(
        "--all-backbones",
        action="store_true",
        help="Evaluate dinov2, mae, and siglip2 sequentially.",
    )
    parser.add_argument(
        "--rae-root",
        type=Path,
        default=Path("/mnt/posttrain/zhaoshitian/models/RAE-collections"),
        help="Root directory containing RAE decoder and latent normalization weights.",
    )
    parser.add_argument(
        "--encoder-path",
        type=Path,
        default=None,
        help="Optional local encoder path override for a single selected backbone.",
    )
    parser.add_argument(
        "--video-path",
        type=Path,
        default=None,
        help="Optional single video to evaluate. When set, split traversal is skipped.",
    )
    parser.add_argument(
        "--video-root",
        type=Path,
        default=Path("/mnt/posttrain/zhaoshitian/datasets/ucf101/OpenDataLab___UCF101/raw/data/UCF101"),
        help="Root directory containing extracted UCF101 class folders.",
    )
    parser.add_argument(
        "--split-zip",
        type=Path,
        default=Path(
            "/mnt/posttrain/zhaoshitian/datasets/ucf101/OpenDataLab___UCF101/raw/data/UCF101TrainTestSplits-RecognitionTask.zip"
        ),
        help="Path to the UCF101 recognition split zip.",
    )
    parser.add_argument(
        "--split-name",
        type=str,
        default="trainlist01.txt",
        help="Split file inside the recognition split zip, e.g. trainlist01.txt.",
    )
    parser.add_argument(
        "--image-size",
        type=int,
        default=None,
        help="Center-crop size before encoding. Defaults to the selected encoder input size.",
    )
    parser.add_argument("--encode-batch-size", type=int, default=16, help="Batch size used for latent extraction.")
    parser.add_argument("--decode-batch-size", type=int, default=16, help="Batch size used for latent decoding.")
    parser.add_argument("--lpips-batch-size", type=int, default=16, help="Batch size used for LPIPS computation.")
    parser.add_argument("--device", type=str, default=None, help="Torch device. Defaults to cuda when available.")
    parser.add_argument("--max-videos", type=int, default=None, help="Optional cap for debugging a subset of videos.")
    parser.add_argument(
        "--interp-mode",
        type=str,
        choices=PPL_INTERP_MODES,
        default="auto",
        help=(
            "Interpolation geometry for the default protocol. "
            "`smooth_continuous` is only implemented for PPL evaluation, while the controlled protocol always uses linear."
        ),
    )
    parser.add_argument(
        "--comparison-protocol",
        type=str,
        choices=("default", "controlled", "both"),
        default="both",
        help=(
            "How to compare backbones: default uses each backbone's requested/default mode, "
            "controlled forces linear, both reports both."
        ),
    )
    parser.add_argument(
        "--eps",
        type=float,
        default=1e-4,
        help="Small path step used in PPL finite differences. alpha is sampled in [0, 1 - eps].",
    )
    parser.add_argument(
        "--samples-per-gap",
        type=int,
        default=1,
        help="How many random alpha samples to draw from each adjacent source-latent pair.",
    )
    parser.add_argument(
        "--trim-percentile-low",
        type=float,
        default=1.0,
        help="Lower percentile for trimmed PPL mean.",
    )
    parser.add_argument(
        "--trim-percentile-high",
        type=float,
        default=99.0,
        help="Upper percentile for trimmed PPL mean.",
    )
    parser.add_argument("--seed", type=int, default=0, help="Random seed for alpha sampling.")
    parser.add_argument(
        "--save-json",
        type=Path,
        default=None,
        help="Optional JSON path to save the aggregated PPL results.",
    )
    parser.add_argument(
        "--plot-per-gap-ppl",
        action="store_true",
        help="Plot per-gap PPL means for a single input video.",
    )
    parser.add_argument(
        "--plot-output",
        type=Path,
        default=None,
        help="Optional PNG path for the per-gap PPL plot. Requires --plot-per-gap-ppl.",
    )
    parser.add_argument(
        "--plot-output-dir",
        type=Path,
        default=None,
        help="Optional output directory for the per-gap PPL plot. Requires --plot-per-gap-ppl.",
    )
    return parser.parse_args()


def validate_args(args: argparse.Namespace) -> None:
    if args.all_backbones and args.config is not None:
        raise ValueError("--all-backbones cannot be combined with --config.")
    if args.all_backbones and args.encoder_path is not None:
        raise ValueError("--all-backbones cannot be combined with --encoder-path.")
    if args.video_path is not None and args.max_videos is not None:
        raise ValueError("--max-videos cannot be used together with --video-path.")
    if args.eps <= 0.0 or args.eps >= 1.0:
        raise ValueError("--eps must be strictly between 0 and 1.")
    if args.samples_per_gap < 1:
        raise ValueError("--samples-per-gap must be at least 1.")
    if args.encode_batch_size < 1 or args.decode_batch_size < 1 or args.lpips_batch_size < 1:
        raise ValueError("All batch sizes must be at least 1.")
    if not (0.0 <= args.trim_percentile_low < args.trim_percentile_high <= 100.0):
        raise ValueError("Trim percentiles must satisfy 0 <= low < high <= 100.")
    if args.plot_output is not None and args.plot_output_dir is not None:
        raise ValueError("--plot-output and --plot-output-dir cannot be used together.")
    if (args.plot_output is not None or args.plot_output_dir is not None) and not args.plot_per_gap_ppl:
        raise ValueError("--plot-output and --plot-output-dir require --plot-per-gap-ppl.")
    if args.plot_per_gap_ppl and args.video_path is None:
        raise ValueError("--plot-per-gap-ppl currently requires --video-path.")


def clone_args(args: argparse.Namespace) -> argparse.Namespace:
    return argparse.Namespace(**copy.deepcopy(vars(args)))


def expand_runtime_paths(args: argparse.Namespace) -> None:
    args.video_root = args.video_root.expanduser()
    args.split_zip = args.split_zip.expanduser()
    args.rae_root = args.rae_root.expanduser()
    if args.video_path is not None:
        args.video_path = args.video_path.expanduser()
    if args.encoder_path is not None:
        args.encoder_path = args.encoder_path.expanduser()
    if args.config is not None:
        args.config = args.config.expanduser()
    if args.save_json is not None:
        args.save_json = args.save_json.expanduser()
    if args.plot_output is not None:
        args.plot_output = args.plot_output.expanduser()
    if args.plot_output_dir is not None:
        args.plot_output_dir = args.plot_output_dir.expanduser()


def build_eval_jobs(args: argparse.Namespace) -> list[EvalJob]:
    backbones = sorted(BACKBONE_SPECS) if args.all_backbones else [args.rae_backbone]
    protocol_order = (
        ["default", "controlled-linear"]
        if args.comparison_protocol == "both"
        else ["default"] if args.comparison_protocol == "default" else ["controlled-linear"]
    )

    jobs: list[EvalJob] = []
    for backbone in backbones:
        for protocol in protocol_order:
            if protocol == "default":
                interp_mode = resolve_interpolation_mode(backbone, args.interp_mode)
            else:
                interp_mode = "linear"
            jobs.append(EvalJob(backbone=backbone, protocol=protocol, interp_mode=interp_mode))
    return jobs


def make_job_args(base_args: argparse.Namespace, job: EvalJob) -> argparse.Namespace:
    job_args = clone_args(base_args)
    job_args.rae_backbone = job.backbone
    return job_args


@dataclass(frozen=True)
class PairwisePPLInterpolator:
    trajectory: torch.Tensor
    interpolate_pair: Callable[[torch.Tensor, torch.Tensor, float], torch.Tensor]

    def interpolate_gap(self, gap_index: int, alpha: float) -> torch.Tensor:
        return self.interpolate_pair(
            self.trajectory[gap_index],
            self.trajectory[gap_index + 1],
            alpha,
        )


class SmoothContinuousPPLInterpolator:
    """Video-aware interpolator that matches TrajectoryPlan smooth_continuous semantics."""

    def __init__(self, trajectory: torch.Tensor) -> None:
        if trajectory.shape[0] < 2:
            raise ValueError("smooth_continuous interpolation requires at least 2 latent frames.")

        self.trajectory = trajectory
        self.plan = TrajectoryPlan(sampling_mode="smooth_continuous")
        self.times = torch.linspace(
            0.0,
            1.0,
            trajectory.shape[0],
            device=trajectory.device,
            dtype=trajectory.dtype,
        )
        self.gap_dt = (self.times[1] - self.times[0]).clamp_min(self.plan.eps)
        self.frame_velocities = self.plan._estimate_frame_velocities(
            trajectory.unsqueeze(0),
            self.times,
        )[0]

    def interpolate_gap(self, gap_index: int, alpha: float) -> torch.Tensor:
        if gap_index < 0 or gap_index >= self.trajectory.shape[0] - 1:
            raise IndexError(f"Gap index {gap_index} is out of range for {self.trajectory.shape[0]} frames.")
        if not 0.0 <= alpha <= 1.0:
            raise ValueError(f"alpha must lie in [0, 1], got {alpha}.")

        left_idx = gap_index
        right_idx = gap_index + 1
        query_t = self.times[left_idx] + self.gap_dt * float(alpha)
        xt, _ = self.plan._hermite_segment(
            query_t.unsqueeze(0),
            self.times[left_idx].unsqueeze(0),
            self.times[right_idx].unsqueeze(0),
            self.trajectory[left_idx].unsqueeze(0),
            self.trajectory[right_idx].unsqueeze(0),
            self.frame_velocities[left_idx].unsqueeze(0),
            self.frame_velocities[right_idx].unsqueeze(0),
        )
        return xt[0]


def build_ppl_interpolator(trajectory: torch.Tensor, interp_mode: str) -> PairwisePPLInterpolator | SmoothContinuousPPLInterpolator:
    if interp_mode == "smooth_continuous":
        return SmoothContinuousPPLInterpolator(trajectory)
    return PairwisePPLInterpolator(trajectory=trajectory, interpolate_pair=get_interpolator(interp_mode))


def compute_lpips_scores(
    lpips_model: LPIPS,
    frames_a: torch.Tensor,
    frames_b: torch.Tensor,
    batch_size: int,
    device: torch.device,
) -> np.ndarray:
    scores: list[np.ndarray] = []
    for start, end in batched_indices(frames_a.shape[0], batch_size):
        batch_a = frames_a[start:end].to(device, non_blocking=True).mul(2.0).sub(1.0)
        batch_b = frames_b[start:end].to(device, non_blocking=True).mul(2.0).sub(1.0)
        batch_scores = lpips_model(batch_a, batch_b, reduction="none")
        scores.append(batch_scores.reshape(-1).detach().cpu().numpy().astype(np.float64))
    if not scores:
        return np.empty((0,), dtype=np.float64)
    return np.concatenate(scores, axis=0)


def empty_video_evaluation(pair_count: int = 0) -> VideoEvaluation:
    empty_int = np.empty((0,), dtype=np.int32)
    empty_float = np.empty((0,), dtype=np.float64)
    return VideoEvaluation(
        sample_values=empty_float,
        pair_count=pair_count,
        gap_indices=empty_int,
        gap_sample_counts=empty_int,
        gap_ppl_mean=empty_float,
        gap_ppl_std=empty_float,
        gap_ppl_min=empty_float,
        gap_ppl_max=empty_float,
    )


def summarize_gap_samples(gap_sample_lists: list[list[float]]) -> tuple[np.ndarray, ...]:
    gap_indices: list[int] = []
    gap_sample_counts: list[int] = []
    gap_ppl_mean: list[float] = []
    gap_ppl_std: list[float] = []
    gap_ppl_min: list[float] = []
    gap_ppl_max: list[float] = []

    for gap_index, gap_samples in enumerate(gap_sample_lists):
        if not gap_samples:
            continue
        gap_values = np.asarray(gap_samples, dtype=np.float64)
        gap_indices.append(gap_index)
        gap_sample_counts.append(int(gap_values.size))
        gap_ppl_mean.append(float(gap_values.mean()))
        gap_ppl_std.append(float(gap_values.std()))
        gap_ppl_min.append(float(gap_values.min()))
        gap_ppl_max.append(float(gap_values.max()))

    return (
        np.asarray(gap_indices, dtype=np.int32),
        np.asarray(gap_sample_counts, dtype=np.int32),
        np.asarray(gap_ppl_mean, dtype=np.float64),
        np.asarray(gap_ppl_std, dtype=np.float64),
        np.asarray(gap_ppl_min, dtype=np.float64),
        np.asarray(gap_ppl_max, dtype=np.float64),
    )


def evaluate_video_ppl(
    *,
    rae,
    video_path: Path,
    args: argparse.Namespace,
    interp_mode: str,
    device: torch.device,
    lpips_model: LPIPS,
    rng: np.random.Generator,
) -> VideoEvaluation:
    frames_np, _ = load_video_frames(video_path, args.image_size)
    frame_tensor = frames_to_tensor(frames_np)
    source_latents = encode_frames(rae, frame_tensor, args.encode_batch_size, device)
    if source_latents.shape[0] < 2:
        return empty_video_evaluation(pair_count=0)

    interpolator = build_ppl_interpolator(source_latents, interp_mode)
    num_pairs = source_latents.shape[0] - 1
    latent_batch_a: list[torch.Tensor] = []
    latent_batch_b: list[torch.Tensor] = []
    latent_batch_gap_indices: list[int] = []
    ppl_samples: list[np.ndarray] = []
    gap_sample_lists: list[list[float]] = [[] for _ in range(num_pairs)]
    pair_count = 0

    def flush_pending_latents() -> None:
        if not latent_batch_a:
            return
        latents_a = torch.stack(latent_batch_a, dim=0)
        latents_b = torch.stack(latent_batch_b, dim=0)
        latent_batch_a.clear()
        latent_batch_b.clear()
        decoded_a = decode_latents(rae, latents_a, args.decode_batch_size, device)
        decoded_b = decode_latents(rae, latents_b, args.decode_batch_size, device)
        lpips_scores = compute_lpips_scores(lpips_model, decoded_a, decoded_b, args.lpips_batch_size, device)
        ppl_values = lpips_scores / (args.eps ** 2)
        if ppl_values.size != len(latent_batch_gap_indices):
            raise RuntimeError("Internal mismatch between decoded PPL samples and gap indices.")
        ppl_samples.append(ppl_values)
        for gap_index, sample_value in zip(latent_batch_gap_indices, ppl_values.tolist()):
            gap_sample_lists[gap_index].append(float(sample_value))
        latent_batch_gap_indices.clear()

    for idx in range(num_pairs):
        pair_count += 1
        alphas = rng.uniform(0.0, 1.0 - args.eps, size=args.samples_per_gap)
        for alpha in alphas:
            alpha_value = float(alpha)
            latent_batch_a.append(interpolator.interpolate_gap(idx, alpha_value))
            latent_batch_b.append(interpolator.interpolate_gap(idx, float(alpha_value + args.eps)))
            latent_batch_gap_indices.append(idx)
            if len(latent_batch_a) >= args.decode_batch_size:
                flush_pending_latents()

    flush_pending_latents()
    if not ppl_samples:
        return empty_video_evaluation(pair_count=pair_count)

    (
        gap_indices,
        gap_sample_counts,
        gap_ppl_mean,
        gap_ppl_std,
        gap_ppl_min,
        gap_ppl_max,
    ) = summarize_gap_samples(gap_sample_lists)
    return VideoEvaluation(
        sample_values=np.concatenate(ppl_samples, axis=0),
        pair_count=pair_count,
        gap_indices=gap_indices,
        gap_sample_counts=gap_sample_counts,
        gap_ppl_mean=gap_ppl_mean,
        gap_ppl_std=gap_ppl_std,
        gap_ppl_min=gap_ppl_min,
        gap_ppl_max=gap_ppl_max,
    )


def summarize_samples(
    *,
    samples: np.ndarray,
    resolved: ResolvedRAEConfig,
    protocol: str,
    interp_mode: str,
    args: argparse.Namespace,
    videos_evaluated: int,
    videos_skipped: int,
    pairs_evaluated: int,
) -> EvalResult:
    if samples.size == 0:
        raise ValueError(f"No valid PPL samples collected for backbone={resolved.backbone}, protocol={protocol}.")

    low_value = float(np.percentile(samples, args.trim_percentile_low))
    high_value = float(np.percentile(samples, args.trim_percentile_high))
    trimmed_mask = (samples >= low_value) & (samples <= high_value)
    trimmed_samples = samples[trimmed_mask]
    if trimmed_samples.size == 0:
        trimmed_samples = samples

    return EvalResult(
        backbone=resolved.backbone,
        protocol=protocol,
        interp_mode_used=interp_mode,
        videos_evaluated=videos_evaluated,
        videos_skipped=videos_skipped,
        pairs_evaluated=pairs_evaluated,
        samples_evaluated=int(samples.size),
        ppl_mean_raw=float(samples.mean()),
        ppl_mean_trimmed=float(trimmed_samples.mean()),
        ppl_median=float(np.median(samples)),
        ppl_std=float(samples.std()),
        ppl_min=float(samples.min()),
        ppl_max=float(samples.max()),
        trim_value_low=low_value,
        trim_value_high=high_value,
        eps=float(args.eps),
        samples_per_gap=int(args.samples_per_gap),
        split_name=None if args.video_path is not None else args.split_name,
        video_path=str(args.video_path) if args.video_path is not None else None,
        image_size=int(args.image_size),
        seed=int(args.seed),
    )


def print_job_header(
    *,
    resolved: ResolvedRAEConfig,
    protocol: str,
    interp_mode: str,
    args: argparse.Namespace,
    device: torch.device,
    video_entries: list[tuple[Path, Path]],
) -> None:
    print("=" * 88)
    print(f"Backbone: {resolved.backbone}")
    print(f"Protocol: {protocol}")
    print(f"Interpolation mode: {interp_mode}")
    print(f"Encoder path/id: {resolved.encoder_reference}")
    if resolved.decoder_path is not None:
        print(f"Decoder checkpoint: {resolved.decoder_path}")
    if resolved.stats_path is not None:
        print(f"Normalization stats: {resolved.stats_path}")
    print(f"Frame crop size: {args.image_size}")
    print(f"Input videos: {len(video_entries)}")
    if args.video_path is not None:
        print(f"Input video: {args.video_path}")
    else:
        print(f"Input split: {args.split_name}")
    print(f"Samples per gap: {args.samples_per_gap}")
    print(f"Epsilon: {args.eps}")
    print(f"Trim percentiles: [{args.trim_percentile_low}, {args.trim_percentile_high}]")
    print(f"Device: {device}")


def resolve_plot_output_path(args: argparse.Namespace, video_path: Path) -> Path:
    if args.plot_output is not None:
        return args.plot_output
    if args.plot_output_dir is not None:
        return args.plot_output_dir / f"{video_path.stem}_per_gap_ppl.png"
    return Path.cwd() / f"{video_path.stem}_per_gap_ppl.png"


def build_curve_label(curve: PerGapCurve) -> str:
    return f"{curve.backbone} | {curve.protocol} | {curve.interp_mode_used}"


def plot_per_gap_curves(
    curves: list[PerGapCurve],
    output_path: Path,
    video_path: Path,
    args: argparse.Namespace,
) -> None:
    if not curves:
        raise ValueError("No per-gap PPL curves available to plot.")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    plt.figure(figsize=(12, 6))
    for curve in curves:
        plt.plot(
            curve.gap_indices,
            curve.gap_ppl_mean,
            linewidth=1.5,
            marker="o",
            markersize=2.5,
            label=build_curve_label(curve),
        )

    plt.xlabel("frame index (left frame of gap)")
    plt.ylabel("PPL")
    plt.title(
        f"Adjacent-frame PPL | {video_path.stem} | "
        f"samples_per_gap={args.samples_per_gap} | eps={args.eps}"
    )
    plt.grid(True, alpha=0.3, linewidth=0.5)
    plt.legend(loc="best")
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()
    print(f"Saved per-gap PPL plot to {output_path}")


def evaluate_job(
    *,
    base_args: argparse.Namespace,
    job: EvalJob,
    device: torch.device,
    lpips_model: LPIPS,
) -> JobEvaluation:
    args = make_job_args(base_args, job)
    resolved = resolve_rae_config(args)
    if args.image_size is None:
        args.image_size = resolved.encoder_input_size
    interp_mode = resolve_interpolation_mode(resolved.backbone, args.interp_mode) if job.protocol == "default" else "linear"

    video_entries = collect_video_entries(args)
    print_job_header(
        resolved=resolved,
        protocol=job.protocol,
        interp_mode=interp_mode,
        args=args,
        device=device,
        video_entries=video_entries,
    )

    rae = load_rae(resolved.config, device)
    rng = np.random.default_rng(args.seed)

    all_samples: list[np.ndarray] = []
    videos_evaluated = 0
    videos_skipped = 0
    pairs_evaluated = 0
    per_gap_curves: list[PerGapCurve] = []

    progress = tqdm(video_entries, desc=f"PPL {resolved.backbone} {job.protocol}")
    for relative_video_path, video_path in progress:
        try:
            video_eval = evaluate_video_ppl(
                rae=rae,
                video_path=video_path,
                args=args,
                interp_mode=interp_mode,
                device=device,
                lpips_model=lpips_model,
                rng=rng,
            )
        except Exception as exc:
            videos_skipped += 1
            tqdm.write(f"[WARN] Skipping {video_path}: {exc}")
            continue

        if video_eval.sample_values.size == 0:
            videos_skipped += 1
            tqdm.write(f"[WARN] Skipping {video_path}: no valid adjacent latent pairs.")
            continue

        all_samples.append(video_eval.sample_values)
        videos_evaluated += 1
        pairs_evaluated += video_eval.pair_count
        if args.plot_per_gap_ppl and video_eval.gap_indices.size > 0:
            per_gap_curves.append(
                PerGapCurve(
                    backbone=resolved.backbone,
                    protocol=job.protocol,
                    interp_mode_used=interp_mode,
                    video_path=str(video_path),
                    gap_indices=video_eval.gap_indices,
                    gap_sample_counts=video_eval.gap_sample_counts,
                    gap_ppl_mean=video_eval.gap_ppl_mean,
                    gap_ppl_std=video_eval.gap_ppl_std,
                    gap_ppl_min=video_eval.gap_ppl_min,
                    gap_ppl_max=video_eval.gap_ppl_max,
                )
            )
        progress.set_postfix_str(relative_video_path.name)

    del rae
    if device.type == "cuda":
        torch.cuda.empty_cache()

    return JobEvaluation(
        result=summarize_samples(
            samples=np.concatenate(all_samples, axis=0) if all_samples else np.empty((0,), dtype=np.float64),
            resolved=resolved,
            protocol=job.protocol,
            interp_mode=interp_mode,
            args=args,
            videos_evaluated=videos_evaluated,
            videos_skipped=videos_skipped,
            pairs_evaluated=pairs_evaluated,
        ),
        per_gap_curves=per_gap_curves,
    )


def format_float(value: float) -> str:
    return f"{value:.6f}"


def print_results_table(results: list[EvalResult]) -> None:
    headers = (
        "backbone",
        "protocol",
        "mode",
        "videos",
        "pairs",
        "samples",
        "ppl_mean_raw",
        "ppl_mean_trimmed",
        "ppl_median",
        "ppl_std",
    )
    rows = [
        (
            result.backbone,
            result.protocol,
            result.interp_mode_used,
            str(result.videos_evaluated),
            str(result.pairs_evaluated),
            str(result.samples_evaluated),
            format_float(result.ppl_mean_raw),
            format_float(result.ppl_mean_trimmed),
            format_float(result.ppl_median),
            format_float(result.ppl_std),
        )
        for result in results
    ]
    widths = [len(header) for header in headers]
    for row in rows:
        for idx, value in enumerate(row):
            widths[idx] = max(widths[idx], len(value))

    def join_row(values: Iterable[str]) -> str:
        return "  ".join(value.ljust(widths[idx]) for idx, value in enumerate(values))

    print("=" * 88)
    print(join_row(headers))
    print(join_row("-" * width for width in widths))
    for row in rows:
        print(join_row(row))


def per_gap_curve_to_payload(curve: PerGapCurve) -> dict[str, object]:
    return {
        "backbone": curve.backbone,
        "protocol": curve.protocol,
        "interp_mode_used": curve.interp_mode_used,
        "video_path": curve.video_path,
        "gap_indices": curve.gap_indices.tolist(),
        "gap_sample_counts": curve.gap_sample_counts.tolist(),
        "gap_ppl_mean": curve.gap_ppl_mean.tolist(),
        "gap_ppl_std": curve.gap_ppl_std.tolist(),
        "gap_ppl_min": curve.gap_ppl_min.tolist(),
        "gap_ppl_max": curve.gap_ppl_max.tolist(),
    }


def maybe_save_results(args: argparse.Namespace, results: list[EvalResult], curves: list[PerGapCurve]) -> None:
    if args.save_json is None:
        return
    args.save_json.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "results": [asdict(result) for result in results],
        "config": {
            "all_backbones": bool(args.all_backbones),
            "comparison_protocol": args.comparison_protocol,
            "requested_interp_mode": args.interp_mode,
            "eps": float(args.eps),
            "samples_per_gap": int(args.samples_per_gap),
            "trim_percentile_low": float(args.trim_percentile_low),
            "trim_percentile_high": float(args.trim_percentile_high),
            "seed": int(args.seed),
            "split_name": None if args.video_path is not None else args.split_name,
            "video_path": str(args.video_path) if args.video_path is not None else None,
            "max_videos": args.max_videos,
            "device": str(get_device(args.device)),
            "plot_per_gap_ppl": bool(args.plot_per_gap_ppl),
        },
    }
    if curves:
        payload["per_gap_curves"] = [per_gap_curve_to_payload(curve) for curve in curves]
    args.save_json.write_text(json.dumps(payload, indent=2))
    print(f"Saved JSON results to {args.save_json}")


def main() -> None:
    args = parse_args()
    validate_args(args)
    expand_runtime_paths(args)

    device = get_device(args.device)
    jobs = build_eval_jobs(args)
    print(f"Using device: {device}")
    print("Loading LPIPS model (first run may download VGG/LPIPS weights)...")
    lpips_model = LPIPS().eval().to(device)

    results: list[EvalResult] = []
    per_gap_curves: list[PerGapCurve] = []
    for job in jobs:
        job_evaluation = evaluate_job(base_args=args, job=job, device=device, lpips_model=lpips_model)
        result = job_evaluation.result
        print(
            f"Finished {result.backbone} / {result.protocol}: "
            f"raw={result.ppl_mean_raw:.6f}, trimmed={result.ppl_mean_trimmed:.6f}, "
            f"samples={result.samples_evaluated}"
        )
        results.append(result)
        per_gap_curves.extend(job_evaluation.per_gap_curves)

    print_results_table(results)
    if args.plot_per_gap_ppl:
        plot_output = resolve_plot_output_path(args, args.video_path)
        plot_per_gap_curves(per_gap_curves, plot_output, args.video_path, args)
    maybe_save_results(args, results, per_gap_curves)


if __name__ == "__main__":
    main()