scripts/analyze_motion_alignment_3models.py

#!/usr/bin/env python3
"""Motion-stratified and correspondence-aware cache analysis for the three AR4 models.

The generation jobs deliberately remain project-native.  This file is an offline
consumer of their saved feature snapshots and RGB anchor frames.  It keeps the
feature comparison at token locations (rather than comparing only a pooled
vector), estimates RAFT flow with a forward/backward consistency mask, and
reports raw, global-transform, homography, dense-flow, and (for WorldPlay)
camera-action rotation alignment.
"""

from __future__ import annotations

import argparse
import csv
import gc
import json
import math
import os
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterable

import cv2
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F


IMG_W, IMG_H = 416, 240
ROLE_ORDER = ("early", "middle", "late", "final")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--self_root", type=Path, required=True)
    parser.add_argument("--causal_root", type=Path, required=True)
    parser.add_argument("--hy_root", type=Path, required=True)
    parser.add_argument("--hy_right_root", type=Path, default=None)
    parser.add_argument("--output_root", type=Path, required=True)
    parser.add_argument("--flow_device", default="auto")
    parser.add_argument("--flow_backend", choices=("raft", "farneback"), default="raft")
    parser.add_argument("--corr_max_tokens", type=int, default=512)
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def finite(value: Any) -> bool:
    try:
        return math.isfinite(float(value))
    except (TypeError, ValueError):
        return False


def mean_or_nan(values: Iterable[float]) -> float:
    values = [float(v) for v in values if finite(v)]
    return float(np.mean(values)) if values else float("nan")


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    fields: list[str] = []
    for row in rows:
        for key in row:
            if key not in fields:
                fields.append(key)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)


def json_dump(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def _as_numpy(value: Any) -> np.ndarray:
    if isinstance(value, torch.Tensor):
        return value.detach().cpu().numpy()
    return np.asarray(value)


def _key_parts(key: str) -> tuple[int, int]:
    left, right = str(key).split(":")[:2]
    return int(left), int(right)


class FeatureRun:
    """A normalized view of one prompt/action's saved features."""

    def __init__(
        self,
        model: str,
        action: str,
        prompt_id: int,
        path: Path,
        anchors: np.ndarray,
        chunks: int,
        chunk_size: int,
        stages: dict[tuple[str, int, int], tuple[np.ndarray, np.ndarray]],
        camera: tuple[np.ndarray, np.ndarray] | None = None,
        dense_path: Path | None = None,
    ):
        self.model = model
        self.action = action
        self.prompt_id = int(prompt_id)
        self.path = path
        self.anchors = anchors
        self.chunk_size = int(chunk_size)
        self.chunks = int(chunks)
        self.stages = stages
        self.camera = camera
        self.dense_path = dense_path
        self._dense_features = None
        self._dense_coords = None
        self._dense_stages = None
        self._dense_chunks = None
        self._dense_steps = None
        self.grid_h = 30
        self.grid_w = 52
        for _, (_, coords) in stages.items():
            if len(coords):
                self.grid_h = max(self.grid_h, int(np.max(coords[:, 1])) + 1)
                self.grid_w = max(self.grid_w, int(np.max(coords[:, 2])) + 1)

    @property
    def stage_names(self) -> list[str]:
        return sorted({key[0] for key in self.stages}, key=stage_sort)

    def get(self, stage: str, chunk: int, step: int):
        # HY's block_53 dense snapshots are large compressed archives.  Keep
        # them lazy so that one run is decompressed only while it is being
        # analyzed, instead of retaining ~1.2 GB per prompt/action in RAM.
        if stage == "block_53" and self.dense_path is not None and self.dense_path.exists():
            if self._dense_features is None:
                data = np.load(self.dense_path, allow_pickle=False)
                # NpzFile lazily decompresses an array on every indexing
                # operation.  Materialize each member once per run so the
                # 6.2-GB float16 feature tensor is not decompressed once per
                # chunk/timestep query.
                try:
                    self._dense_features = np.asarray(data["features"])
                    self._dense_coords = np.asarray(data["coords"], dtype=np.int32)
                    self._dense_stages = np.asarray(data["stages"]).astype(str)
                    self._dense_chunks = np.asarray(data["chunks"], dtype=np.int32)
                    self._dense_steps = np.asarray(data["steps"], dtype=np.int32)
                finally:
                    data.close()
            stages = self._dense_stages
            chunks = self._dense_chunks
            steps = self._dense_steps
            matches = np.flatnonzero(
                (stages == stage) & (chunks == int(chunk)) & (steps == int(step))
            )
            if len(matches):
                index = int(matches[0])
                return self._dense_features[index].astype(np.float32), self._dense_coords
        return self.stages.get((stage, int(chunk), int(step)))

    def release(self):
        if self._dense_features is not None:
            self._dense_features = None
            self._dense_coords = None
            self._dense_stages = None
            self._dense_chunks = None
            self._dense_steps = None
            gc.collect()


def stage_sort(stage: str):
    if stage == "projected":
        return (10000, stage)
    try:
        return (int(stage.split("_")[1]), stage)
    except (IndexError, ValueError):
        return (9000, stage)


def stage_role(run: FeatureRun, stage: str) -> str:
    if stage == "projected":
        return "final"
    block_stages = [name for name in run.stage_names if name.startswith("block_")]
    if stage in block_stages:
        index = block_stages.index(stage)
        return ROLE_ORDER[min(index, len(ROLE_ORDER) - 1)]
    return stage


def _anchor(path: Path) -> np.ndarray:
    data = np.load(path, allow_pickle=False)
    value = data["frames"]
    if value.ndim != 4:
        raise ValueError(f"Expected [T,H,W,3] anchors at {path}, got {value.shape}")
    return value.astype(np.uint8)


def load_self(root: Path) -> list[FeatureRun]:
    result: list[FeatureRun] = []
    for path in sorted((root / "runs").glob("prompt_*.pt")):
        state = torch.load(path, map_location="cpu", weights_only=False)
        anchor_path = path.with_suffix(".anchors.npz")
        if not anchor_path.exists():
            raise FileNotFoundError(anchor_path)
        coords_hidden = _as_numpy(state["sample_coords"]["hidden"]).astype(np.int32)
        stages: dict[tuple[str, int, int], tuple[np.ndarray, np.ndarray]] = {}
        for stage, values in state["records"].items():
            if not stage.endswith("_hidden"):
                continue
            for key, value in values.items():
                chunk, step = _key_parts(key)
                stages[(stage, chunk, step)] = (_as_numpy(value).astype(np.float32), coords_hidden)
        # The native recorder has a compact full-grid random projection for its
        # last block.  Keep it as an additional final-stage representation.
        full_coords = np.stack(
            np.meshgrid(np.arange(3), np.arange(30), np.arange(52), indexing="ij"), axis=-1
        ).reshape(-1, 3).astype(np.int32)
        for key, value in state.get("projected", {}).items():
            chunk, step = _key_parts(key)
            array = _as_numpy(value).astype(np.float32).reshape(-1, _as_numpy(value).shape[-1])
            stages[("projected", chunk, step)] = (array, full_coords)
        result.append(
            FeatureRun(
                "self_forcing",
                "none",
                int(state.get("run_index", len(result))),
                path,
                _anchor(anchor_path),
                int(state["num_frames"]) // int(state["num_frame_per_block"]),
                int(state["num_frame_per_block"]),
                stages,
            )
        )
        del state
    if len(result) != 10:
        print(f"[warn] Self-Forcing runs found {len(result)} prompts, expected 10")
    return result


def _flat_indices_to_coords(indices: np.ndarray, frame_count: int = 3) -> np.ndarray:
    indices = indices.astype(np.int64).reshape(-1)
    frame_size = 30 * 52
    return np.stack(
        [indices // frame_size, (indices % frame_size) // 52, indices % 52], axis=1
    ).astype(np.int32)


def load_causal(root: Path) -> list[FeatureRun]:
    result: list[FeatureRun] = []
    for run_dir in sorted((root / "runs").glob("prompt_*")):
        path = run_dir / "feature_snapshots.pt"
        if not path.exists():
            continue
        state = torch.load(path, map_location="cpu", weights_only=False)
        anchor_path = run_dir / "rgb_anchor_frames.npz"
        if not anchor_path.exists():
            raise FileNotFoundError(anchor_path)
        stages: dict[tuple[str, int, int], tuple[np.ndarray, np.ndarray]] = {}
        index_dict = state.get("feature_indices", {})
        for key, value in state["features"].items():
            layer, chunk, step = (int(v) for v in str(key).split(":"))
            indices = index_dict.get(key)
            if indices is None:
                token_count = 3 * 30 * 52
                indices = np.linspace(0, token_count - 1, int(state["max_tokens"])).round().astype(np.int64)
            coords = _flat_indices_to_coords(_as_numpy(indices))
            stages[(f"block_{layer:02d}", chunk, step)] = (_as_numpy(value).astype(np.float32), coords)
        prompt_id = int(state.get("prompt_id", run_dir.name.split("_")[-1]))
        result.append(
            FeatureRun(
                "causal_forcing",
                "none",
                prompt_id,
                path,
                _anchor(anchor_path),
                int(state["num_chunks"]),
                3,
                stages,
            )
        )
        del state
    if len(result) != 10:
        print(f"[warn] Causal-Forcing runs found {len(result)} prompts, expected 10")
    return result


def _load_hy_action(root: Path, action: str) -> list[FeatureRun]:
    result: list[FeatureRun] = []
    for case_dir in sorted((root / "runs").glob("prompt_*")):
        run_dir = case_dir / action
        if not run_dir.exists():
            continue
        anchor_path = run_dir / "rgb_anchor_frames.npz"
        dense_path = run_dir / "dense_selected_snapshots.npz"
        sampled_path = run_dir / "final_hidden_snapshots.npz"
        # The sampled archive contains all requested layers and is sufficient
        # for the ordinary token metrics.  Keep the expensive dense archive
        # only as a lazy source for HY's final block (block_53).
        snapshot_path = sampled_path if sampled_path.exists() else dense_path
        if not anchor_path.exists() or not snapshot_path.exists():
            print(f"[warn] skip incomplete HY run {run_dir}")
            continue
        data = np.load(snapshot_path, allow_pickle=False)
        features = data["features"].astype(np.float32)
        chunks = data["chunks"].astype(int)
        steps = data["steps"].astype(int)
        stages_np = data["stages"].astype(str)
        coords = data["coords"].astype(np.int32)
        grid_shape = tuple(int(v) for v in data["grid_shape"])
        stages: dict[tuple[str, int, int], tuple[np.ndarray, np.ndarray]] = {}
        for index in range(features.shape[0]):
            stages[(stages_np[index], int(chunks[index]), int(steps[index]))] = (
                features[index],
                coords,
            )
        camera = None
        camera_path = run_dir / "camera_trajectory.npz"
        if camera_path.exists():
            camera_data = np.load(camera_path, allow_pickle=False)
            camera = (camera_data["viewmats"].astype(np.float64), camera_data["intrinsics"].astype(np.float64))
        metadata = run_dir / "run_metadata.json"
        prompt_id = int(case_dir.name.split("_")[-1])
        chunks_count = int(json.loads(metadata.read_text()).get("video_length", 45) if metadata.exists() else 45)
        latent_count = (chunks_count - 1) // 4 + 1
        # Prefer the actual snapshot chunk count when available.
        chunks_count = max(1, int(np.max(chunks)) + 1)
        result.append(
            FeatureRun(
                "hy_worldplay",
                action,
                prompt_id,
                run_dir,
                _anchor(anchor_path),
                chunks_count,
                4,
                stages,
                camera,
                # Keep the cross-model comparison at the same 240-token
                # sampling budget.  The dense 6,240-token archive is retained
                # as an optional artifact, but using it for all-pairs matching
                # would make HY's final layer incomparable and unnecessarily
                # expensive (6,240^2 distances per query).
                dense_path=None,
            )
        )
    return result


class FlowEstimator:
    def __init__(self, backend: str, device: str):
        self.backend = backend
        self.device = torch.device("cuda" if device == "auto" and torch.cuda.is_available() else (device if device != "auto" else "cpu"))
        self.model = None
        self.transforms = None
        if backend == "raft":
            try:
                from torchvision.models.optical_flow import Raft_Small_Weights, raft_small

                weights = Raft_Small_Weights.DEFAULT
                self.model = raft_small(weights=weights, progress=True).eval().to(self.device)
                self.transforms = weights.transforms()
                print(f"[flow] RAFT-small on {self.device}")
            except Exception as error:
                print(f"[flow] RAFT unavailable ({error}); using Farneback")
                self.backend = "farneback"

    @staticmethod
    def _image(frame: np.ndarray) -> np.ndarray:
        if frame.shape[:2] != (IMG_H, IMG_W):
            return cv2.resize(frame, (IMG_W, IMG_H), interpolation=cv2.INTER_AREA)
        return frame

    def _farneback(self, first: np.ndarray, second: np.ndarray) -> np.ndarray:
        a = cv2.cvtColor(self._image(first), cv2.COLOR_RGB2GRAY)
        b = cv2.cvtColor(self._image(second), cv2.COLOR_RGB2GRAY)
        return cv2.calcOpticalFlowFarneback(a, b, None, 0.5, 3, 21, 5, 7, 1.5, 0)

    @torch.inference_mode()
    def pair(self, first: np.ndarray, second: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Return first->second flow, reverse flow, and a mask on first pixels."""
        first = self._image(first)
        second = self._image(second)
        if self.backend != "raft" or self.model is None:
            forward = self._farneback(first, second)
            backward = self._farneback(second, first)
        else:
            x = torch.from_numpy(first).permute(2, 0, 1).float().div_(255).unsqueeze(0).to(self.device)
            y = torch.from_numpy(second).permute(2, 0, 1).float().div_(255).unsqueeze(0).to(self.device)
            x, y = self.transforms(x, y)
            forward = self.model(x, y)[-1][0].permute(1, 2, 0).float().cpu().numpy()
            backward = self.model(y, x)[-1][0].permute(1, 2, 0).float().cpu().numpy()
        # The feature query is made at pixels in ``first`` (the current/target
        # frame), so the consistency mask must also be indexed in that frame.
        # A previous implementation used the reverse flow and a mask indexed
        # in ``second``; that silently inverted the alignment direction.
        h, w = forward.shape[:2]
        yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
        sx = xx + forward[..., 0]
        sy = yy + forward[..., 1]
        sampled_backward = cv2.remap(backward, sx, sy, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
        fb = np.linalg.norm(forward + sampled_backward, axis=-1)
        magnitude = np.linalg.norm(forward, axis=-1)
        in_bounds = (sx >= 0) & (sx < w) & (sy >= 0) & (sy < h)
        # The threshold scales mildly with motion, avoiding an overly strict
        # rejection of the fast camera actions while retaining occlusion masks.
        valid = in_bounds & (fb <= 1.5 + 0.05 * magnitude)
        return forward, backward, valid


def flow_at(flow: np.ndarray, xy: np.ndarray) -> np.ndarray:
    h, w = flow.shape[:2]
    x = np.clip(xy[:, 0], 0, w - 1).astype(np.float32)
    y = np.clip(xy[:, 1], 0, h - 1).astype(np.float32)
    return np.stack(
        [cv2.remap(flow[..., dim], x, y, cv2.INTER_LINEAR).reshape(-1) for dim in range(2)], axis=1
    )


def mask_at(mask: np.ndarray, xy: np.ndarray) -> np.ndarray:
    h, w = mask.shape[:2]
    x = np.clip(xy[:, 0], 0, w - 1).astype(np.float32)
    y = np.clip(xy[:, 1], 0, h - 1).astype(np.float32)
    value = cv2.remap(mask.astype(np.uint8), x, y, cv2.INTER_NEAREST).reshape(-1)
    return value.astype(bool)


def fit_transforms(flow: np.ndarray, valid: np.ndarray):
    h, w = flow.shape[:2]
    yy, xx = np.mgrid[0:h:8, 0:w:8].astype(np.float32)
    points = np.stack([xx.reshape(-1), yy.reshape(-1)], axis=1)
    selected = valid[::8, ::8].reshape(-1)
    points = points[selected]
    if len(points) < 6:
        fallback_y, fallback_x = np.mgrid[0:h:4, 0:w:4]
        points = np.stack([fallback_x.reshape(-1), fallback_y.reshape(-1)], axis=1).astype(np.float32)
        selected_flow = flow[::4, ::4].reshape(-1, 2)
    else:
        selected_flow = flow[::8, ::8].reshape(-1, 2)[selected]
    destinations = points + selected_flow
    affine = None
    homography = None
    if len(points) >= 3:
        affine, _ = cv2.estimateAffine2D(points, destinations, method=cv2.RANSAC, ransacReprojThreshold=3.0)
    if len(points) >= 4:
        homography, _ = cv2.findHomography(points, destinations, cv2.RANSAC, 4.0)
    median = np.median(flow[valid] if bool(np.any(valid)) else flow.reshape(-1, 2), axis=0)
    translation = np.asarray([[1.0, 0.0, median[0]], [0.0, 1.0, median[1]]], dtype=np.float32)
    return translation, affine, homography


def apply_transform(points: np.ndarray, transform: np.ndarray | None) -> np.ndarray | None:
    if transform is None or len(points) == 0:
        return None
    if transform.shape == (2, 3):
        return points @ transform[:, :2].T + transform[:, 2]
    homogeneous = np.concatenate([points, np.ones((len(points), 1), dtype=np.float32)], axis=1)
    mapped = homogeneous @ transform.T
    return mapped[:, :2] / np.clip(mapped[:, 2:3], 1e-6, None)


def nearest_features(source_xy: np.ndarray, source_features: np.ndarray, query_xy: np.ndarray):
    if len(source_xy) == 0 or len(query_xy) == 0:
        return np.empty((0, source_features.shape[-1]), dtype=np.float32), np.zeros(len(query_xy), bool), np.empty(len(query_xy))
    distances = ((query_xy[:, None, :] - source_xy[None, :, :]) ** 2).sum(axis=-1)
    indices = np.argmin(distances, axis=1)
    return source_features[indices], np.ones(len(query_xy), bool), np.sqrt(distances[np.arange(len(indices)), indices])


def cosine_mean(left: np.ndarray, right: np.ndarray, mask: np.ndarray | None = None) -> float:
    if len(left) == 0 or len(right) == 0:
        return float("nan")
    value = np.sum(left * right, axis=-1) / (
        np.linalg.norm(left, axis=-1) * np.linalg.norm(right, axis=-1) + 1e-8
    )
    if mask is not None:
        value = value[mask]
    return mean_or_nan(value)


def select_frame(features: np.ndarray, coords: np.ndarray, frame: int):
    selected = coords[:, 0] == int(frame)
    return features[selected], coords[selected, 1:3][:, ::-1].astype(np.float32)


def token_correspondence(target: np.ndarray, target_xy: np.ndarray, source: np.ndarray, source_xy: np.ndarray, max_tokens: int):
    if len(target) == 0 or len(source) == 0:
        return {"top1_cosine": float("nan"), "top5_cosine": float("nan"), "mnn_fraction": float("nan"), "match_distance": float("nan")}
    def thin(array, xy):
        if len(array) <= max_tokens:
            return array, xy
        idx = np.linspace(0, len(array) - 1, max_tokens).round().astype(int)
        return array[idx], xy[idx]
    target, target_xy = thin(target, target_xy)
    source, source_xy = thin(source, source_xy)
    target_n = target / (np.linalg.norm(target, axis=-1, keepdims=True) + 1e-8)
    source_n = source / (np.linalg.norm(source, axis=-1, keepdims=True) + 1e-8)
    similarity = target_n @ source_n.T
    topk = np.sort(similarity, axis=1)[:, -min(5, similarity.shape[1]):]
    best_source = np.argmax(similarity, axis=1)
    best_target = np.argmax(similarity, axis=0)
    target_ids = np.arange(len(best_source))
    mnn = best_target[best_source] == target_ids
    distance = np.linalg.norm(target_xy - source_xy[best_source], axis=-1)
    return {
        "top1_cosine": float(np.mean(topk[:, -1])),
        "top5_cosine": float(np.mean(topk)),
        "mnn_fraction": float(np.mean(mnn)),
        "match_distance": float(np.mean(distance)),
    }


def camera_rotation_homography(camera, source_index: int, target_index: int, image_shape=(IMG_H, IMG_W)):
    if camera is None:
        return None
    viewmats, intrinsics = camera
    if source_index >= len(viewmats) or target_index >= len(viewmats):
        return None
    source_k = intrinsics[source_index].copy()
    target_k = intrinsics[target_index].copy()
    h, w = image_shape
    for k in (source_k, target_k):
        k[0, 0] *= w
        k[0, 2] *= w
        k[1, 1] *= h
        k[1, 2] *= h
    source_r = viewmats[source_index][:3, :3]
    target_r = viewmats[target_index][:3, :3]
    try:
        return source_k @ source_r @ target_r.T @ np.linalg.inv(target_k)
    except np.linalg.LinAlgError:
        return None


def action_motion(camera, source_index: int, target_index: int):
    if camera is None:
        return {"action_translation": float("nan"), "action_rotation_deg": float("nan")}
    viewmats, _ = camera
    if source_index >= len(viewmats) or target_index >= len(viewmats):
        return {"action_translation": float("nan"), "action_rotation_deg": float("nan")}
    source_to_world = np.linalg.inv(viewmats[source_index])
    target_to_world = np.linalg.inv(viewmats[target_index])
    relative_rotation = source_to_world[:3, :3].T @ target_to_world[:3, :3]
    trace = np.clip((np.trace(relative_rotation) - 1.0) / 2.0, -1.0, 1.0)
    return {
        "action_translation": float(np.linalg.norm(target_to_world[:3, 3] - source_to_world[:3, 3])),
        "action_rotation_deg": float(np.degrees(np.arccos(trace))),
    }


def token_points(run: FeatureRun, coords: np.ndarray) -> np.ndarray:
    # coords are (frame,y,x) in the latent feature grid; convert to pixel
    # centers in the resized RGB canvas, represented as (x,y).
    return np.stack(
        [
            (coords[:, 2].astype(np.float32) + 0.5) * IMG_W / run.grid_w,
            (coords[:, 1].astype(np.float32) + 0.5) * IMG_H / run.grid_h,
        ],
        axis=1,
    )


def alignment_metrics(run: FeatureRun, stage: str, target_pair, source_pair, target_slot: int, flow: np.ndarray, valid_flow: np.ndarray, transforms, action_h, corr_max_tokens: int):
    target_features, target_coords = target_pair
    source_features, source_coords = source_pair
    target_features, target_xy = select_frame(target_features, target_coords, target_slot)
    source_boundary, source_boundary_xy = select_frame(source_features, source_coords, run.chunk_size - 1)
    source_same, source_same_xy = select_frame(source_features, source_coords, min(target_slot, run.chunk_size - 1))
    if len(target_features) == 0 or len(source_boundary) == 0:
        return None
    target_pixels = token_points(run, np.concatenate([np.full((len(target_xy), 1), target_slot), target_xy[:, ::-1]], axis=1).astype(np.int32))
    source_boundary_pixels = token_points(run, np.concatenate([np.full((len(source_boundary_xy), 1), run.chunk_size - 1), source_boundary_xy[:, ::-1]], axis=1).astype(np.int32))
    source_same_pixels = token_points(run, np.concatenate([np.full((len(source_same_xy), 1), min(target_slot, run.chunk_size - 1)), source_same_xy[:, ::-1]], axis=1).astype(np.int32))
    raw_features, _, raw_dist = nearest_features(source_boundary_pixels, source_boundary, target_pixels)
    same_features, _, _ = nearest_features(source_same_pixels, source_same, target_pixels)
    target_flow = flow_at(flow, target_pixels)
    flow_query = target_pixels + target_flow
    flow_mask = mask_at(valid_flow, target_pixels)
    flow_features, _, flow_dist = nearest_features(source_boundary_pixels, source_boundary, flow_query)
    translation, affine, homography = transforms
    def transformed_metric(transform):
        query = apply_transform(target_pixels, transform)
        if query is None:
            return float("nan"), float("nan")
        values, _, distances = nearest_features(source_boundary_pixels, source_boundary, query)
        in_bounds = (query[:, 0] >= 0) & (query[:, 0] < IMG_W) & (query[:, 1] >= 0) & (query[:, 1] < IMG_H)
        return cosine_mean(target_features, values, in_bounds), mean_or_nan(distances[in_bounds])
    trans_cos, trans_dist = transformed_metric(translation)
    affine_cos, affine_dist = transformed_metric(affine)
    homo_cos, homo_dist = transformed_metric(homography)
    action_cos, action_dist = transformed_metric(action_h)
    correspondence = token_correspondence(target_features, target_pixels, source_boundary, source_boundary_pixels, corr_max_tokens)
    return {
        "same_slot_cosine": cosine_mean(target_features, same_features),
        "boundary_raw_cosine": cosine_mean(target_features, raw_features),
        "translation_aligned_cosine": trans_cos,
        "affine_aligned_cosine": affine_cos,
        "homography_aligned_cosine": homo_cos,
        "flow_aligned_cosine": cosine_mean(target_features, flow_features, flow_mask),
        "action_rotation_cosine": action_cos,
        "raw_match_distance": mean_or_nan(raw_dist),
        "flow_match_distance": mean_or_nan(flow_dist[flow_mask]),
        "translation_match_distance": trans_dist,
        "affine_match_distance": affine_dist,
        "homography_match_distance": homo_dist,
        "action_match_distance": action_dist,
        "valid_flow_ratio": float(np.mean(flow_mask)),
        "occlusion_ratio": float(1.0 - np.mean(flow_mask)),
        **correspondence,
    }


def motion_rows(runs: list[FeatureRun], estimator: FlowEstimator, corr_max_tokens: int) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    correspondence_rows: list[dict[str, Any]] = []
    flow_cache: dict[tuple[str, int, int, int], tuple[np.ndarray, np.ndarray, np.ndarray]] = {}
    for run in runs:
        for chunk in range(1, run.chunks):
            source_index = (chunk - 1) * run.chunk_size + run.chunk_size - 1
            if source_index >= len(run.anchors):
                continue
            for slot in range(run.chunk_size):
                target_index = chunk * run.chunk_size + slot
                if target_index >= len(run.anchors):
                    continue
                cache_key = (str(run.path), chunk, source_index, target_index)
                if cache_key not in flow_cache:
                    # The first return is target->source and is therefore the
                    # displacement used to query the previous chunk's tokens.
                    forward, _, valid = estimator.pair(run.anchors[target_index], run.anchors[source_index])
                    flow_cache[cache_key] = (forward, valid, np.asarray([]))
                flow, valid, _ = flow_cache[cache_key]
                flow_mag = np.linalg.norm(flow, axis=-1)
                median = np.median(flow[valid] if bool(np.any(valid)) else flow.reshape(-1, 2), axis=0)
                residual = flow - median[None, None]
                source_img = cv2.resize(run.anchors[source_index], (IMG_W, IMG_H), interpolation=cv2.INTER_AREA)
                target_img = cv2.resize(run.anchors[target_index], (IMG_W, IMG_H), interpolation=cv2.INTER_AREA)
                scene_cut = float(np.mean(np.abs(source_img.astype(np.float32) - target_img.astype(np.float32))) / 255.0)
                action_motion_values = action_motion(run.camera, source_index, target_index)
                action_h = camera_rotation_homography(run.camera, source_index, target_index)
                transforms = fit_transforms(flow, valid)
                stage_names = [stage for stage in run.stage_names if stage.startswith("block_") or stage == "projected"]
                for stage in stage_names:
                    common_steps = sorted(
                        set(step for st, ch, step in run.stages if st == stage and ch == chunk)
                        & set(step for st, ch, step in run.stages if st == stage and ch == chunk - 1)
                    )
                    for step in common_steps:
                        target_pair = run.get(stage, chunk, step)
                        source_pair = run.get(stage, chunk - 1, step)
                        metrics = alignment_metrics(run, stage, target_pair, source_pair, slot, flow, valid, transforms, action_h, corr_max_tokens)
                        if metrics is None:
                            continue
                        row = {
                            "model": run.model,
                            "action": run.action,
                            "prompt_id": run.prompt_id,
                            "chunk": chunk,
                            "target_slot": slot,
                            "step": step,
                            "stage": stage,
                            "role": stage_role(run, stage),
                            "source_frame": source_index,
                            "target_frame": target_index,
                            "total_motion": float(np.mean(flow_mag)),
                            "camera_motion": float(np.linalg.norm(median)),
                            "object_motion": float(np.mean(np.linalg.norm(residual, axis=-1))),
                            "scene_cut": scene_cut,
                            **action_motion_values,
                            **metrics,
                        }
                        rows.append(row)
                        if row["role"] == "final":
                            correspondence_rows.append({
                                key: row[key]
                                for key in ("model", "action", "prompt_id", "chunk", "target_slot", "step", "stage", "role", "total_motion")
                            } | {key: metrics[key] for key in ("top1_cosine", "top5_cosine", "mnn_fraction", "match_distance")})
        run.release()
    values = np.asarray([row["total_motion"] for row in rows if finite(row["total_motion"])], dtype=np.float64)
    if len(values) >= 3:
        low, high = np.quantile(values, [1 / 3, 2 / 3])
        for row in rows:
            row["motion_bin"] = "low" if row["total_motion"] <= low else "high" if row["total_motion"] > high else "medium"
    for row in correspondence_rows:
        matches = [item for item in rows if all(item[k] == row[k] for k in ("model", "action", "prompt_id", "chunk", "target_slot", "step", "stage", "role"))]
        row["motion_bin"] = matches[0].get("motion_bin", "unknown") if matches else "unknown"
    metadata = {
        "flow_backend": estimator.backend,
        "flow_device": str(estimator.device),
        "flow_resolution": [IMG_W, IMG_H],
        "forward_backward_consistency": "valid = in-bounds and FB error <= 1.5 + 0.05*|flow|",
        "motion_bin_edges": [float(low), float(high)] if len(values) >= 3 else [],
    }
    return rows, correspondence_rows, metadata


def group_rows(rows: list[dict[str, Any]], keys: list[str], metrics: list[str]) -> list[dict[str, Any]]:
    groups: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        groups[tuple(row.get(key) for key in keys)].append(row)
    output: list[dict[str, Any]] = []
    for group, values in sorted(groups.items(), key=lambda item: tuple(map(str, item[0]))):
        item = {key: value for key, value in zip(keys, group)}
        item["count"] = len(values)
        for metric in metrics:
            item[metric] = mean_or_nan([row.get(metric) for row in values])
        output.append(item)
    return output


def plot_outputs(rows: list[dict[str, Any]], correspondence: list[dict[str, Any]], output: Path) -> None:
    output.mkdir(parents=True, exist_ok=True)
    if not rows:
        return
    models = sorted({str(row["model"]) for row in rows})
    colors = {model: plt.cm.tab10(index) for index, model in enumerate(models)}
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    for model in models:
        selected = [row for row in rows if row["model"] == model]
        axes[0, 0].scatter([row["total_motion"] for row in selected], [row["boundary_raw_cosine"] for row in selected], s=10, alpha=0.45, label=model, color=colors[model])
    axes[0, 0].set(xlabel="flow magnitude", ylabel="raw boundary cosine", title="Motion vs raw cross-chunk similarity")
    axes[0, 0].legend(fontsize=8)
    methods = ["boundary_raw_cosine", "translation_aligned_cosine", "affine_aligned_cosine", "homography_aligned_cosine", "flow_aligned_cosine", "action_rotation_cosine"]
    labels = ["raw", "translation", "affine", "homography", "dense flow", "action rotation"]
    positions = np.arange(len(methods))
    for index, model in enumerate(models):
        values = [mean_or_nan([row.get(method) for row in rows if row["model"] == model]) for method in methods]
        axes[0, 1].plot(positions, values, marker="o", label=model, color=colors[model])
    axes[0, 1].set_xticks(positions, labels, rotation=25, ha="right")
    axes[0, 1].set_ylim(-1, 1)
    axes[0, 1].set_title("Alignment methods")
    axes[0, 1].legend(fontsize=8)
    for index, model in enumerate(models):
        selected = [row for row in rows if row["model"] == model]
        bins = ["low", "medium", "high"]
        raw = [mean_or_nan([row["boundary_raw_cosine"] for row in selected if row.get("motion_bin") == name]) for name in bins]
        flow = [mean_or_nan([row["flow_aligned_cosine"] for row in selected if row.get("motion_bin") == name]) for name in bins]
        axes[1, 0].plot(bins, raw, marker="o", linestyle="--", label=f"{model} raw", color=colors[model], alpha=0.55)
        axes[1, 0].plot(bins, flow, marker="o", label=f"{model} flow", color=colors[model])
    axes[1, 0].set_ylim(-1, 1)
    axes[1, 0].set_title("Motion bins")
    axes[1, 0].set_ylabel("cosine")
    axes[1, 0].legend(fontsize=7)
    if correspondence:
        for model in models:
            selected = [row for row in correspondence if row["model"] == model]
            axes[1, 1].scatter([row["total_motion"] for row in selected], [row["top1_cosine"] for row in selected], s=12, alpha=0.5, label=model, color=colors[model])
        axes[1, 1].set(xlabel="flow magnitude", ylabel="top-1 feature match cosine", title="Token correspondence")
        axes[1, 1].legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(output / "motion_alignment_overview.png", dpi=180)
    plt.close(fig)

    if correspondence:
        fig, axes = plt.subplots(1, 3, figsize=(15, 4.5))
        for model in models:
            selected = [row for row in correspondence if row["model"] == model]
            for axis, metric, title in zip(axes, ("top1_cosine", "mnn_fraction", "match_distance"), ("Top-1 cosine", "MNN fraction", "Match distance")):
                axis.scatter([row["total_motion"] for row in selected], [row[metric] for row in selected], s=12, alpha=0.5, label=model, color=colors[model])
                axis.set_title(title)
                axis.set_xlabel("flow magnitude")
        axes[0].set_ylabel("value")
        axes[0].legend(fontsize=8)
        fig.tight_layout()
        fig.savefig(output / "token_correspondence_summary.png", dpi=180)
        plt.close(fig)


def markdown_table(rows: list[dict[str, Any]], columns: list[str]) -> str:
    if not rows:
        return "_No rows._"
    lines = ["| " + " | ".join(columns) + " |", "|" + "|".join("---" for _ in columns) + "|"]
    for row in rows:
        values = []
        for column in columns:
            value = row.get(column, "")
            values.append(f"{float(value):.4f}" if isinstance(value, (float, np.floating)) and finite(value) else str(value))
        lines.append("| " + " | ".join(values) + " |")
    return "\n".join(lines)


def write_report(output: Path, rows: list[dict[str, Any]], correspondence: list[dict[str, Any]], metadata: dict[str, Any], runs: list[FeatureRun]) -> None:
    stage_summary = group_rows(rows, ["model", "action", "role", "stage"], ["boundary_raw_cosine", "translation_aligned_cosine", "affine_aligned_cosine", "homography_aligned_cosine", "flow_aligned_cosine", "action_rotation_cosine", "total_motion", "valid_flow_ratio"])
    bin_summary = group_rows(rows, ["model", "action", "motion_bin"], ["total_motion", "boundary_raw_cosine", "flow_aligned_cosine", "homography_aligned_cosine", "object_motion", "occlusion_ratio"])
    json_dump(output / "motion_alignment_summary.json", {"metadata": metadata, "stage_summary": stage_summary, "motion_bin_summary": bin_summary, "run_count": len(runs), "row_count": len(rows), "correspondence_count": len(correspondence)})
    report = [
        "# Experiment 3: motion stratification and explicit alignment",
        "",
        f"Runs: {len(runs)} prompt/action runs; metric rows: {len(rows)}; correspondence rows: {len(correspondence)}.",
        "",
        "The raw baseline compares the current chunk token with the previous chunk's boundary frame at the same spatial coordinate. `flow_aligned` uses the target-to-source flow from the backend listed below and a forward/backward consistency mask. Affine and homography are global-transform fits to the valid flow. `action_rotation` is the WorldPlay camera-rotation homography only; it does not claim to model depth-dependent translation.",
        "",
        "## By model/action/stage",
        "",
        markdown_table(stage_summary, ["model", "action", "role", "stage", "count", "boundary_raw_cosine", "translation_aligned_cosine", "affine_aligned_cosine", "homography_aligned_cosine", "flow_aligned_cosine", "action_rotation_cosine"]),
        "",
        "## Motion bins",
        "",
        markdown_table(bin_summary, ["model", "action", "motion_bin", "count", "total_motion", "boundary_raw_cosine", "homography_aligned_cosine", "flow_aligned_cosine", "object_motion", "occlusion_ratio"]),
        "",
        "## Interpretation guardrails",
        "",
        "- A positive dense-flow gain with lower raw cosine supports spatial migration rather than loss of content information.",
        "- Action rotation is a physically informed control for HY; forward translation requires depth and is therefore reported separately rather than treated as an exact warp.",
        "- Correlation and alignment rows are prompt-level repeated measurements; use prompt bootstrap or a mixed-effects model for significance claims.",
        "",
        "## Run settings",
        "",
        "```json",
        json.dumps(metadata, indent=2, ensure_ascii=False),
        "```",
    ]
    (output / "REPORT.md").write_text("\n".join(report) + "\n", encoding="utf-8")
    write_csv(output / "motion_alignment_metrics.csv", rows)
    write_csv(output / "motion_alignment_summary.csv", stage_summary)
    write_csv(output / "motion_alignment_motion_bins.csv", bin_summary)
    write_csv(output / "token_correspondence.csv", correspondence)


def main() -> None:
    args = parse_args()
    output = args.output_root.resolve()
    output.mkdir(parents=True, exist_ok=True)
    if not args.overwrite and (output / "motion_alignment_metrics.csv").exists():
        print(f"[skip] existing analysis at {output}; pass --overwrite to recompute")
        return
    runs: list[FeatureRun] = []
    runs.extend(load_self(args.self_root.resolve()))
    runs.extend(load_causal(args.causal_root.resolve()))
    runs.extend(_load_hy_action(args.hy_root.resolve(), "forward"))
    runs.extend(_load_hy_action(args.hy_root.resolve(), "static"))
    if args.hy_right_root is not None:
        runs.extend(_load_hy_action(args.hy_right_root.resolve(), "right"))
    if not runs:
        raise RuntimeError("No normalized runs found")
    estimator = FlowEstimator(args.flow_backend, args.flow_device)
    rows, correspondence, metadata = motion_rows(runs, estimator, args.corr_max_tokens)
    metadata.update({
        "models": sorted({run.model for run in runs}),
        "actions": sorted({run.action for run in runs}),
        "prompt_ids": sorted({run.prompt_id for run in runs}),
        "run_descriptions": [
            {"model": run.model, "action": run.action, "prompt_id": run.prompt_id, "chunk_size": run.chunk_size, "chunks": run.chunks, "anchor_frames": len(run.anchors), "stages": run.stage_names}
            for run in runs
        ],
    })
    plot_outputs(rows, correspondence, output)
    write_report(output, rows, correspondence, metadata, runs)
    json_dump(output / "analysis_config.json", metadata)
    print(f"[complete] {output}: {len(runs)} runs, {len(rows)} motion rows, {len(correspondence)} correspondence rows")


if __name__ == "__main__":
    main()