icassp2027-experiment/test_local_turn_features.py

"""CPU tests for B1 alignment, causal pooling, fusion, and train/infer round trips."""

from __future__ import annotations

import copy
import importlib.util
import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import unittest
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np
import torch
from PIL import Image, ImageDraw

from easycom_head import (
    LOCAL_FUSION_SCOPES,
    TurnPredictionHead,
    checkpoint_head_config,
)
from extract_turn_local_features import (
    FrozenQwenLocalEncoder, ego4d_clip_boxes, extract_local,
    validate_source_timeline,
    model_config_digest,
    box_coordinate_offset, frame_crops, parse_args, preview_boxes, preview_ticks, preview_time_bounds,
    TESTED_TRANSFORMERS_VERSIONS, validate_local_transformers_version,
    main as extract_main,
)
from extract_highrate_person_av_features import (
    MODEL_FPS as HIGHRATE_MODEL_FPS,
    _chunked_audio_encode,
    _chunked_person_encode,
    _model_tick_bounds,
    _square_gray_face,
)
from local_turn_features import (
    HIGHRATE_PERSON_AV_FORMAT, LOCAL_FORMAT, FramePersonFeatures,
    causal_frame_ticks, crop_box, easycom_boxes,
    load_frame_set_features, load_highrate_person_av_features,
    load_local_features, load_track_features,
    local_head_dimensions,
    pooled_recent_visual, reference_alignment,
)
from train_turn_head import _load_examples


def fixture(n=20, name="Session_4", dataset="easycom"):
    times = (torch.arange(n, dtype=torch.float64) + 1) / 5
    reference = {"states": torch.randn(n, 4, 8), "label_indices": torch.arange(n),
                 "times": times, "mode": "open-stream-replay", "start": 0}
    side = {"format": LOCAL_FORMAT, "dataset": dataset, "sample_id": name,
            "times": times.clone(), "label_indices": torch.arange(n), "partial": False,
            "visual": torch.randn(n, 6), "audio": torch.randn(n, 5),
            "visual_valid": torch.ones(n, dtype=torch.bool),
            "audio_valid": torch.ones(n, dtype=torch.bool),
            "config": {"visual_dim": 6, "audio_dim": 5, "encoder_id": "synthetic-test",
                       "causal_policy": "half-open-trailing-window", "uses_speech_labels": False}}
    side["visual_valid"][0] = False
    side["visual"][0] = 0
    frame_times = torch.arange(max(1, (n + 2) // 3), dtype=torch.float64) * .5
    counts = [2 if index % 3 == 0 else (0 if index % 3 == 1 else 1)
              for index in range(len(frame_times))]
    offsets = torch.tensor([0, *np.cumsum(counts).tolist()], dtype=torch.long)
    person_features = torch.randn(int(offsets[-1]), 6)
    side["persons"] = {
        "features": person_features,
        "frame_times": frame_times,
        "frame_offsets": offsets,
        "records": [{"id": str(index)} for index in range(len(person_features))],
    }
    labels = np.tile(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], dtype=np.float32), (n // 4 + 1, 1))[:n]
    return reference, side, labels


def highrate_fixture(n=20, name="Session_4", dataset="easycom"):
    reference, side, labels = fixture(n=n, name=name, dataset=dataset)
    side["format"] = HIGHRATE_PERSON_AV_FORMAT
    side["visual"] = torch.randn(n, 6)
    side["visual_valid"] = torch.ones(n, dtype=torch.bool)
    side["persons"] = {
        "features": side["visual"].clone(),
        "frame_times": reference["times"].clone(),
        "frame_offsets": torch.arange(n + 1, dtype=torch.long),
        "records": [
            {"id": "p", "asd_interval_end": float(reference["times"][row])}
            for row in range(n)
        ],
    }
    side["config"].update(
        {
            "model_fps": 25.0,
            "output_hz": 5.0,
            "person_feature_kind": (
                "muvap-asd-causal-visual-plus-activity-logit"
            ),
            "person_visual_dim": 5,
            "person_activity_logit_dim": 1,
            "audio_feature_kind": "muvap-cpc-causal-audio-state",
            "audio_source": "single-shared-traditional-cpc-stream",
            "person_tokens_duplicate_audio": False,
            "asd_audio_preprocessing": (
                "causal-float-pcm-clip-only-no-rms-normalization"
            ),
            "asd_temporal_attention": "causal",
            "causal_policy": "right-edge-causal-highrate-asd",
        }
    )
    return reference, side, labels


class LocalTests(unittest.TestCase):
    def setUp(self):
        torch.set_num_threads(1)

    def test_tested_transformers_versions(self):
        for version in ("4.52.3", "4.57.6"):
            self.assertEqual(validate_local_transformers_version(version), version)
        for version in ("4.50.0.dev0", "4.57.6.dev0", "5.0.0"):
            with self.subTest(version=version), self.assertRaisesRegex(RuntimeError, "found " + version):
                validate_local_transformers_version(version)

    def test_half_open_grid_and_pooling(self):
        self.assertEqual(list(causal_frame_ticks(.2, 2)), [0])
        self.assertEqual(list(causal_frame_ticks(1., 2)), [0, 1])
        ts = torch.tensor([0., .5, 1., 1.5], dtype=torch.float64)
        vectors = torch.tensor([[2.], [4.], [90.], [100.]])
        valid = torch.ones(4, dtype=torch.bool)
        value, present = pooled_recent_visual(ts, vectors, valid, 1., 1.)
        self.assertTrue(present)
        self.assertEqual(value.item(), 3.)  # no frame at t=1
        value, _ = pooled_recent_visual(ts, vectors, valid, 1.2, 1.)
        self.assertEqual(value.item(), 47.)  # .5 and 1.0, not 0.0
        value, present = pooled_recent_visual(ts, vectors, valid * False, 1., 1.)
        self.assertFalse(present)
        self.assertEqual(value.item(), 0.)

    def test_highrate_person_av_right_edges_are_causal_and_inclusive(self):
        reference, side, _ = highrate_fixture(n=5)
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "Session_4.pt"
            torch.save(side, path)
            local, people, contract = load_highrate_person_av_features(
                path,
                reference,
                sample_id="Session_4",
                dataset="easycom",
            )
            packed_audio, audio_contract = load_local_features(
                path,
                reference,
                sample_id="Session_4",
                dataset="easycom",
                modalities="audio",
            )
        self.assertEqual(local.shape, (5, 6))
        self.assertEqual(people.row_to_frame.tolist(), [0, 1, 2, 3, 4])
        self.assertTrue(torch.equal(people.chosen_frame_times, reference["times"]))
        self.assertEqual(contract["format"], HIGHRATE_PERSON_AV_FORMAT)
        self.assertEqual(packed_audio.shape, (5, 6))
        self.assertEqual(audio_contract["modalities"], "audio")
        values, mask = people.padded()
        self.assertEqual(values.shape, (5, 1, 6))
        self.assertTrue(mask.all())

    def test_highrate_person_av_rejects_old_or_duplicated_audio_contracts(self):
        reference, side, _ = highrate_fixture(n=5)
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "Session_4.pt"
            old = copy.deepcopy(side)
            old["format"] = "causal-highrate-person-av-v1"
            torch.save(old, path)
            with self.assertRaisesRegex(ValueError, "Invalid local feature format"):
                load_highrate_person_av_features(
                    path,
                    reference,
                    sample_id="Session_4",
                    dataset="easycom",
                )

            duplicated = copy.deepcopy(side)
            duplicated["config"]["person_tokens_duplicate_audio"] = True
            torch.save(duplicated, path)
            with self.assertRaisesRegex(ValueError, "causal ASD contract"):
                load_highrate_person_av_features(
                    path,
                    reference,
                    sample_id="Session_4",
                    dataset="easycom",
                )

    def test_highrate_face_geometry_and_chunked_causal_coverage(self):
        image = Image.new("RGB", (40, 20), "white")
        face = _square_gray_face(image)
        self.assertEqual(face.shape, (112, 112))
        self.assertEqual(face.dtype, np.uint8)
        self.assertEqual(int(face[0, 0]), 128)
        clip = SimpleNamespace(dense_start=0.0, dense_end=1.0)
        self.assertEqual(_model_tick_bounds(clip, 1.0), (0, 25))
        self.assertEqual(HIGHRATE_MODEL_FPS, 25.0)

        class FakeEncoder:
            audio_dim = 2
            feature_dim = 3

            def __init__(self):
                self.audio_lengths = []
                self.person_lengths = []

            def encode_audio(self, waveform, num_frames):
                self.audio_lengths.append(num_frames)
                return torch.stack(
                    (torch.arange(num_frames), torch.ones(num_frames)), dim=-1
                ).float()

            def encode_person(self, faces, audio_states):
                self.person_lengths.append(len(faces))
                return torch.cat((audio_states, torch.ones(len(faces), 1)), dim=-1)

        encoder = FakeEncoder()
        faces = np.zeros((9, 112, 112), dtype=np.uint8)
        waveform = np.zeros(9 * 640, dtype=np.float32)
        audio_result = _chunked_audio_encode(
            encoder,
            waveform,
            num_frames=len(faces),
            output_frames=4,
            context_frames=2,
        )
        person_result = _chunked_person_encode(
            encoder,
            faces,
            audio_result,
            output_frames=4,
            context_frames=2,
        )
        self.assertEqual(audio_result.shape, (9, 2))
        self.assertEqual(person_result.shape, (9, 3))
        self.assertEqual(encoder.audio_lengths, [4, 6, 3])
        self.assertEqual(encoder.person_lengths, [4, 6, 3])

    def test_box_formats_and_crop(self):
        person = {"Participant_ID": -1, "x1": 2, "y1": 3, "x2": 7, "y2": 9}
        faces = easycom_boxes([{"Frame_Number": 1, "Participants": [person]}], "face")
        heads = easycom_boxes([{"Frame_Number": 0, "Participants": person}], "head")
        self.assertEqual(faces, heads)
        self.assertEqual(crop_box([202., 3., 207., 9.], 1920, 1080, offset_x=200), (2, 3, 7, 9))
        self.assertIsNone(crop_box([4., 4., 3., 5.], 10, 10))
        self.assertIsNone(crop_box([float("nan"), 4., 9., 5.], 10, 10))

    def test_ego4d_clip_frames(self):
        boxes = [{"frame": 3, "clip_frame": None, "x": 1, "y": 2, "width": 3, "height": 4}]
        person = {"person_id": "x", "camera_wearer": False, "tracking_paths": [{"track": boxes}, {"track": boxes}]}
        wearer = {**person, "camera_wearer": True}
        actual = ego4d_clip_boxes({"persons": [wearer, person]})
        self.assertEqual(actual, {3: [{"id": "x", "xyxy": [1., 2., 4., 6.]}]})

    def test_dataset_specific_coordinate_defaults_and_explicit_overrides(self):
        required = ["--data-root", "/raw", "--feature-root", "/global",
                    "--target-root", "/targets", "--output-root", "/local", "--samples", "4"]
        self.assertEqual(parse_args(required).box_coordinate_space, "encoded")
        self.assertEqual(parse_args(required + ["--dataset", "ego4d"]).box_coordinate_space, "camera")
        self.assertEqual(parse_args(required + ["--box-coordinate-space", "camera"]).box_coordinate_space, "camera")
        # An appended correction overrides an old LOCAL_ARGS array in the shell.
        self.assertEqual(parse_args(required + ["--box-coordinate-space", "camera",
                                                "--box-coordinate-space", "encoded"]).box_coordinate_space, "encoded")

    def test_frame_crops_subtract_each_clips_banner_before_encoding(self):
        expected = (100, 80, 160, 140)
        image = Image.new("RGB", (1920, 1080))
        ImageDraw.Draw(image).rectangle(expected, fill="red")
        clips, annotations = [], []
        for width in (2123, 2124, 2203, 1920):
            offset = width - 1920
            clips.append(SimpleNamespace(crop=(offset, 0, 1920, 1080), dense_start=float(len(clips))))
            annotations.append({0: [{"id": "3", "xyxy": [100 + offset, 80, 160 + offset, 140]}]})
        media = SimpleNamespace(clips=clips,
                                _source_frame_for_time=lambda t: (int(t), 0, 0.),
                                frame_at_or_before=lambda t: image.copy())
        args = SimpleNamespace(box_coordinate_space="encoded", crop_expansion=1.)
        for index, clip in enumerate(clips):
            with self.subTest(crop=clip.crop):
                _, crops, records = frame_crops(media, annotations, float(index), args)
                self.assertEqual(records[0]["box"], expected)
                self.assertEqual(records[0]["coordinate_offset"], [clip.crop[0], 0])
                self.assertEqual(records[0]["annotation_box"], annotations[index][0][0]["xyxy"])
                self.assertEqual(crops[0].size, (60, 60))
                self.assertEqual(crops[0].getpixel((30, 30)), (255, 0, 0))

    def test_cropped_camera_annotations_and_uncropped_ego4d_are_not_shifted_twice(self):
        clip = SimpleNamespace(crop=(204, 12, 1920, 1080), dense_start=0.)
        self.assertEqual(box_coordinate_offset(clip, "encoded"), (204, 12))
        self.assertEqual(box_coordinate_offset(clip, "camera"), (0, 0))
        self.assertEqual(box_coordinate_offset(SimpleNamespace(crop=None), "encoded"), (0, 0))
        with self.assertRaisesRegex(ValueError, "coordinate space"):
            box_coordinate_offset(clip, "typo")
        media = SimpleNamespace(clips=[clip], _source_frame_for_time=lambda t: (0, 0, 0.),
                                frame_at_or_before=lambda t: Image.new("RGB", (1920, 1080)))
        annotations = [{0: [{"id": "3", "xyxy": [100, 80, 160, 140]}]}]
        args = SimpleNamespace(box_coordinate_space="camera", crop_expansion=1.)
        _, _, records = frame_crops(media, annotations, .1, args)
        self.assertEqual(records[0]["box"], (100, 80, 160, 140))
        # Encoded boxes translate on BOTH axes, then clip at the visible edge.
        annotations = [{0: [{"id": "3", "xyxy": [190, 4, 230, 50]}]}]
        args.box_coordinate_space = "encoded"
        _, _, records = frame_crops(media, annotations, .1, args)
        self.assertEqual(records[0]["box"], (0, 0, 26, 38))

    def test_preview_records_the_same_transformation_used_for_extraction(self):
        clip = SimpleNamespace(path=Path("source.mp4"), width=2124, height=1080,
                               crop=(204, 0, 1920, 1080), dense_start=0.)
        media = SimpleNamespace(clips=[clip], _source_frame_for_time=lambda t: (0, 0, 0.),
                                frame_at_or_before=lambda t: Image.new("RGB", (1920, 1080)))
        annotations = [{0: [{"id": "3", "xyxy": [304, 80, 364, 140]}]}]
        with tempfile.TemporaryDirectory() as directory:
            args = SimpleNamespace(box_coordinate_space="encoded", crop_expansion=1.,
                                   video_fps=2., output_root=Path(directory))
            preview_boxes(media, annotations, args, "Session_4", .2)
            folder = Path(directory) / "Session_4_preview"
            audit = json.loads((folder / "boxes.json").read_text())
            self.assertEqual(audit[0]["coordinate_offset"], [204, 0])
            self.assertEqual(audit[0]["boxes"][0]["box"], [100, 80, 160, 140])
            self.assertTrue((folder / audit[0]["image"]).is_file())

    def test_preview_can_save_more_than_four_images_in_a_later_range(self):
        clip = SimpleNamespace(path=Path("source.mp4"), width=640, height=360,
                               crop=None, dense_start=0.)
        media = SimpleNamespace(clips=[clip],
                                _source_frame_for_time=lambda t: (0, round(t * 2), t),
                                frame_at_or_before=lambda t: Image.new("RGB", (640, 360)))
        with tempfile.TemporaryDirectory() as directory:
            args = SimpleNamespace(box_coordinate_space="camera", crop_expansion=1., video_fps=2.,
                                   output_root=Path(directory), preview_count=8,
                                   preview_start=2., preview_end=6., preview_with_boxes=False)
            preview_boxes(media, [{}], args, "clip", 10.)
            folder = Path(directory) / "clip_preview"
            audit = json.loads((folder / "boxes.json").read_text())
            self.assertEqual([row["time"] for row in audit], [2., 2.5, 3., 3.5, 4., 4.5, 5., 5.5])
            self.assertEqual(audit[0]["image"], "frame_000004.jpg")
            self.assertEqual(len(list(folder.glob("*.jpg"))), 8)
            report = json.loads((folder / "preview.json").read_text())
            self.assertEqual(report["saved_count"], 8)
            self.assertEqual(report["saved_with_boxes"], 0)

    def test_boxed_preview_filters_valid_corrected_crops_without_decoding(self):
        # Encoded boxes are corrected by the crop origin, as in feature extraction.
        clip = SimpleNamespace(crop=(204, 0, 640, 360), width=844, height=360)
        media = SimpleNamespace(clips=[clip], _source_frame_for_time=lambda t: (0, round(t * 2), t))
        valid = {"id": "p", "xyxy": [304, 80, 364, 140]}
        invalid = {"id": "p", "xyxy": [4, 80, 64, 140]}  # entirely inside removed banner
        annotations = [{1: [valid], 3: [invalid], 5: [valid], 7: [valid]}]
        args = SimpleNamespace(box_coordinate_space="encoded", crop_expansion=1., video_fps=2.,
                               preview_count=2, preview_start=0., preview_end=None, preview_with_boxes=True)
        selected, grid_count, eligible_count = preview_ticks(media, annotations, args, 4.)
        self.assertEqual(selected, [1, 7])  # uniform among valid frames, not first N
        self.assertEqual((grid_count, eligible_count), (8, 3))
        # No RGB reader is needed by the scan, and annotations are not changed.
        self.assertEqual(len(annotations[0]), 4)

    def test_empty_boxed_preview_writes_an_explicit_empty_report(self):
        clip = SimpleNamespace(crop=None, width=640, height=360)
        media = SimpleNamespace(clips=[clip], _source_frame_for_time=lambda t: (0, round(t * 2), t))
        with tempfile.TemporaryDirectory() as directory:
            args = SimpleNamespace(box_coordinate_space="camera", crop_expansion=1., video_fps=2.,
                                   output_root=Path(directory), preview_count=40,
                                   preview_start=0., preview_end=None, preview_with_boxes=True)
            preview_boxes(media, [{}], args, "empty_clip", 4.)
            folder = Path(directory) / "empty_clip_preview"
            self.assertEqual(json.loads((folder / "boxes.json").read_text()), [])
            report = json.loads((folder / "preview.json").read_text())
            self.assertEqual(report["eligible_frame_count"], 0)
            self.assertEqual(report["saved_count"], 0)
            self.assertEqual(list(folder.glob("*.jpg")), [])

    def test_preview_range_clips_to_available_end_and_excludes_endpoint(self):
        args = SimpleNamespace(video_fps=2., preview_start=.1, preview_end=100., preview_count=100)
        self.assertEqual(preview_time_bounds(args, 1.), (.1, 1.))
        self.assertEqual(preview_ticks(None, [], args, 1.), ([1], 1, 1))
        args.preview_start = 2.
        with self.assertRaisesRegex(ValueError, "outside the available prefix"):
            preview_time_bounds(args, 1.)

    def test_preview_cli_validation(self):
        import contextlib
        import io

        required = ["--data-root", "/raw", "--feature-root", "/global", "--target-root", "/targets",
                    "--output-root", "/preview", "--samples", "4"]
        valid = parse_args(required + ["--preview-only", "--preview-count", "40", "--preview-start", "30",
                                       "--preview-end", "120", "--preview-with-boxes"])
        self.assertEqual((valid.preview_count, valid.preview_start, valid.preview_end), (40, 30., 120.))
        for extra in (["--preview-only", "--preview-count", "0"],
                      ["--preview-only", "--preview-start", "-1"],
                      ["--preview-only", "--preview-end", "nan"],
                      ["--preview-only", "--preview-start", "10", "--preview-end", "10"],
                      ["--preview-only", "--max-steps", "20", "--preview-end", "60"],
                      ["--preview-count", "40"], ["--preview-with-boxes"]):
            with self.subTest(extra=extra), contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
                parse_args(required + extra)

    def test_preview_end_limits_media_reader_without_loading_encoder(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            reference, _, labels = fixture()
            torch.save(reference, root / "Session_4.pt")
            np.save(root / "Session_4.npy", labels)
            args = ["--data-root", str(root), "--feature-root", str(root), "--target-root", str(root),
                    "--output-root", str(root / "preview"), "--samples", "4",
                    "--preview-only", "--preview-end", "2", "--preview-count", "20"]
            media = SimpleNamespace(clips=[], metadata=lambda: {})
            with patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])) as build, \
                 patch("extract_turn_local_features.preview_boxes") as preview, \
                 patch("extract_turn_local_features.FrozenQwenLocalEncoder") as encoder:
                extract_main(args)
                self.assertEqual(build.call_args.args[2], 2.)
                preview.assert_called_once()
                encoder.assert_not_called()

    def test_sidecar_alignment_contract_and_modalities(self):
        reference, side, _ = fixture()
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "Session_4.pt"
            for mode, width in (("visual", 7), ("audio", 6), ("av", 13)):
                torch.save(side, path)
                features, contract = load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities=mode)
                self.assertEqual(features.shape, (20, width))
                dims = local_head_dimensions(contract)
                self.assertEqual(dims["local_visual_dim"], 0 if mode == "audio" else 6)
            for field, value in (("times", side["times"] + .2), ("label_indices", torch.arange(20) + 1),
                                 ("sample_id", "Session_5"), ("partial", True),
                                 ("visual", torch.full((20, 6), float("nan"))),
                                 ("visual_valid", torch.ones(20))):
                broken = {**side, field: value}
                torch.save(broken, path)
                with self.subTest(field=field), self.assertRaises(ValueError):
                    load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities="av")
            torch.save(side, path)
            wrong_contract = copy.deepcopy(contract)
            wrong_contract["config"]["encoder_id"] = "other-weights"
            with self.assertRaisesRegex(ValueError, "contract"):
                load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities="av", expected_contract=wrong_contract)
            with self.assertRaisesRegex(ValueError, "times"):
                reference_alignment({**reference, "times": reference["times"] + .2})

    def test_frame_set_loader_is_causal_ragged_and_fail_closed(self):
        reference, side, _ = fixture()
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "Session_4.pt"
            torch.save(side, path)
            audio, people, contract = load_frame_set_features(
                path,
                reference,
                sample_id="Session_4",
                dataset="easycom",
            )
            self.assertEqual(audio.shape, (20, 6))
            self.assertEqual(contract["modalities"], "av")
            self.assertEqual(people.dimension, 6)
            self.assertTrue(
                torch.all(
                    people.chosen_frame_times[
                        torch.isfinite(people.chosen_frame_times)
                    ]
                    < reference["times"][
                        torch.isfinite(people.chosen_frame_times)
                    ]
                )
            )
            # At 0.2/0.4s, both rows reuse the 0.0s frame with two people.
            values, mask = people.padded(torch.tensor([0, 1, 2]))
            self.assertEqual(values.shape, (3, 2, 6))
            self.assertEqual(mask.tolist(), [[True, True], [True, True], [False, False]])
            self.assertTrue(torch.equal(values[0], values[1]))
            self.assertEqual(torch.count_nonzero(values[2]), 0)
            # The frame stamped exactly 1.0s is excluded at the 1.0s cutoff.
            self.assertEqual(people.chosen_frame_times[4].item(), .5)

            broken_cases = []
            broken = copy.deepcopy(side)
            broken["persons"]["frame_offsets"][-1] -= 1
            broken_cases.append(broken)
            broken = copy.deepcopy(side)
            broken["persons"]["frame_times"][1] = broken["persons"]["frame_times"][0]
            broken_cases.append(broken)
            broken = copy.deepcopy(side)
            broken["persons"]["records"] = []
            broken_cases.append(broken)
            broken = copy.deepcopy(side)
            broken["persons"]["features"][0, 0] = float("nan")
            broken_cases.append(broken)
            for index, broken in enumerate(broken_cases):
                torch.save(broken, path)
                with self.subTest(index=index), self.assertRaises(ValueError):
                    load_frame_set_features(
                        path,
                        reference,
                        sample_id="Session_4",
                        dataset="easycom",
                    )

    def test_track_loader_builds_causal_histories_and_resets_tracks(self):
        reference, side, _ = fixture(n=12)
        frame_times = torch.tensor([0., .5, 1., 1.5, 2.], dtype=torch.float64)
        side["persons"] = {
            "features": torch.arange(1, 9, dtype=torch.float32)[:, None].repeat(1, 6),
            "frame_times": frame_times,
            "frame_offsets": torch.tensor([0, 2, 4, 6, 7, 8]),
            "records": [
                {"id": "A", "source_clip_index": 0},
                {"id": "B", "source_clip_index": 0},
                {"id": "B", "source_clip_index": 0},
                {"id": "A", "source_clip_index": 0},
                {"id": "A", "source_clip_index": 0},
                {"id": -1, "source_clip_index": 0},
                # The same ID in another retained source clip is a new track.
                {"id": "A", "source_clip_index": 1},
                {"id": "A", "source_clip_index": 1},
            ],
        }
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "Session_4.pt"
            torch.save(side, path)
            pooled, people, contract = load_track_features(
                path,
                reference,
                sample_id="Session_4",
                dataset="easycom",
                track_max_gap_seconds=.75,
            )
            self.assertEqual(pooled.shape, (12, 13))
            self.assertEqual(contract["modalities"], "av")
            self.assertEqual(
                people.previous_observation.tolist(),
                [-1, -1, 1, 0, 3, -1, -1, 6],
            )
            # Cutoff 1.2s selects frame 1.0s. A has observations 1,4,5;
            # anonymous -1 remains a one-observation track.
            history, mask = people.padded_history(
                torch.tensor([5]),
                max_observations=4,
            )
            self.assertEqual(mask.tolist(), [[[True, True, True, False],
                                              [True, False, False, False]]])
            self.assertEqual(history[0, 0, :, 0].tolist(), [1., 4., 5., 0.])
            self.assertEqual(history[0, 1, :, 0].tolist(), [6., 0., 0., 0.])
            # Cutoff 1.8s selects frame 1.5s; source clip reset prevents A's
            # earlier observations from crossing the boundary.
            history, mask = people.padded_history(
                torch.tensor([8]),
                max_observations=4,
            )
            self.assertEqual(mask.tolist(), [[[True, False, False, False]]])
            self.assertEqual(history[0, 0, :, 0].tolist(), [7., 0., 0., 0.])

            broken = copy.deepcopy(side)
            broken["persons"]["records"][1]["id"] = "A"
            torch.save(broken, path)
            with self.assertRaisesRegex(ValueError, "Duplicate person/track ID"):
                load_track_features(
                    path,
                    reference,
                    sample_id="Session_4",
                    dataset="easycom",
                    track_max_gap_seconds=.75,
                )

    def test_frame_person_concatenation_selection_and_empty_padding(self):
        features = torch.tensor([[1., 2.], [3., 4.]])
        a = FramePersonFeatures(
            features,
            torch.tensor([0, 2], dtype=torch.long),
            torch.tensor([0, -1], dtype=torch.long),
            torch.tensor([0., float("nan")], dtype=torch.float64),
        )
        b = FramePersonFeatures(
            torch.empty(0, 2),
            torch.tensor([0, 0], dtype=torch.long),
            torch.tensor([0], dtype=torch.long),
            torch.tensor([.5], dtype=torch.float64),
        )
        combined = FramePersonFeatures.concatenate([a, b])
        self.assertEqual(len(combined), 3)
        self.assertEqual(combined.row_to_frame.tolist(), [0, -1, 1])
        selected = combined.select_rows(torch.tensor([2, 0]))
        values, mask = selected.padded()
        self.assertEqual(mask.tolist(), [[False, False], [True, True]])
        self.assertEqual(torch.count_nonzero(values[0]), 0)
        self.assertTrue(torch.equal(values[1], features))

    def test_frame_set_head_is_permutation_invariant_and_masks_padding(self):
        torch.manual_seed(7)
        head = TurnPredictionHead(
            hidden_size=8,
            num_input_layers=4,
            layer_fusion="concat-mlp",
            local_visual_dim=6,
            local_audio_dim=5,
            local_projection_dim=3,
            local_fusion="frame-set-attention",
            person_attention_dim=4,
            person_attention_heads=2,
            transition_auxiliary=True,
            dropout=0,
        ).eval()
        states = torch.randn(3, 4, 8)
        local = torch.randn(3, 6)
        local[:, -1] = 1
        people = torch.randn(3, 3, 6)
        mask = torch.tensor(
            [[True, True, False], [True, False, False], [False, False, False]]
        )
        people[~mask] = 0
        logits, transitions = head.forward_with_transition(
            states,
            local,
            people,
            mask,
        )
        self.assertEqual(logits.shape, (3, 10, 3))
        self.assertEqual(transitions.shape, (3, 10))
        (logits.square().mean() + transitions.square().mean()).backward()
        self.assertIsNotNone(head.person_attention.in_proj_weight.grad)

        permutation = torch.tensor([1, 0, 2])
        permuted = head(states, local, people[:, permutation], mask[:, permutation])
        self.assertTrue(torch.allclose(logits, permuted, atol=1e-6, rtol=1e-6))
        changed = people.clone()
        changed[~mask] = 1000
        self.assertTrue(torch.allclose(logits, head(states, local, changed, mask)))
        with self.assertRaisesRegex(ValueError, "person features"):
            head(states, local)

        config = {
            "hidden_size": 8,
            "num_input_layers": 4,
            "layer_fusion": "concat-mlp",
            "local_visual_dim": 6,
            "local_audio_dim": 5,
            "local_projection_dim": 3,
            "local_fusion": "frame-set-attention",
            "person_attention_dim": 4,
            "person_attention_heads": 2,
            "transition_auxiliary": True,
            "dropout": 0,
        }
        restored = TurnPredictionHead(
            **checkpoint_head_config({"head_config": config})
        )
        restored.load_state_dict(head.state_dict())
        restored.eval()
        self.assertTrue(torch.equal(restored(states, local, people, mask), logits))

    def test_highrate_person_av_starts_from_exact_audio_only_base(self):
        torch.manual_seed(19)
        audio = TurnPredictionHead(
            hidden_size=8,
            local_visual_dim=0,
            local_audio_dim=5,
            local_projection_dim=3,
            local_fusion="pooled",
            person_attention_dim=4,
            person_attention_heads=2,
            dropout=0,
        ).eval()
        highrate = TurnPredictionHead(
            hidden_size=8,
            local_visual_dim=6,
            local_audio_dim=5,
            local_projection_dim=3,
            local_fusion="highrate-person-av-attention",
            person_attention_dim=4,
            person_attention_heads=2,
            dropout=0,
        ).eval()
        highrate_state = highrate.state_dict()
        for name, value in audio.state_dict().items():
            self.assertIn(name, highrate_state)
            self.assertEqual(tuple(value.shape), tuple(highrate_state[name].shape))
            highrate_state[name] = value
        highrate.load_state_dict(highrate_state)

        states = torch.randn(4, 8)
        local = torch.randn(4, 6)
        local[:, -1] = 1
        people = torch.randn(4, 3, 6)
        mask = torch.tensor(
            [[True, True, False], [True, False, False], [False, False, False],
             [True, True, True]]
        )
        people[~mask] = 0
        reference = audio(states, local)
        actual = highrate(states, local, people, mask)
        self.assertTrue(torch.equal(reference, actual))
        diagnostics = highrate.person_adapter_diagnostics()
        self.assertEqual(diagnostics["residual_output_weight_norm"], 0.0)
        self.assertEqual(diagnostics["delta_logit_max_abs"], 0.0)

        actual.sum().backward()
        assert highrate.highrate_person_residual is not None
        self.assertGreater(
            float(highrate.highrate_person_residual.weight.grad.abs().sum()),
            0.0,
        )

    def test_highrate_person_av_head_uses_variable_person_sets(self):
        torch.manual_seed(9)
        head = TurnPredictionHead(
            hidden_size=8,
            local_visual_dim=6,
            local_audio_dim=5,
            local_projection_dim=3,
            local_fusion="highrate-person-av-attention",
            person_attention_dim=4,
            person_attention_heads=2,
            dropout=0,
        ).eval()
        states = torch.randn(3, 8)
        local = torch.randn(3, 6)
        local[:, -1] = 1
        people = torch.randn(3, 3, 6)
        mask = torch.tensor(
            [[True, True, False], [True, False, False], [False, False, False]]
        )
        people[~mask] = 0
        logits = head(states, local, people, mask)
        self.assertEqual(logits.shape, (3, 10, 3))
        permutation = torch.tensor([1, 0, 2])
        self.assertTrue(
            torch.allclose(
                logits,
                head(states, local, people[:, permutation], mask[:, permutation]),
                atol=1e-6,
                rtol=1e-6,
            )
        )
        changed = people.clone()
        changed[~mask] = 1000
        self.assertTrue(torch.equal(logits, head(states, local, changed, mask)))
        logits.square().mean().backward()
        self.assertIsNotNone(head.person_attention.in_proj_weight.grad)

    def test_f5_screening_audio_person_read_is_gated_and_permutation_invariant(self):
        torch.manual_seed(11)
        common = {
            "hidden_size": 8,
            "local_visual_dim": 6,
            "local_audio_dim": 5,
            "local_projection_dim": 3,
            "person_attention_dim": 4,
            "person_attention_heads": 2,
            "dropout": 0,
        }
        f1 = TurnPredictionHead(
            **common,
            local_fusion="frame-set-attention",
        ).eval()
        f5 = TurnPredictionHead(
            **common,
            local_fusion="audio-person-frame-set",
        ).eval()
        # Align all shared parameters.  A zero F5 gate must reproduce F1
        # exactly, independent of the newly initialized reverse-read block.
        f5_state = f5.state_dict()
        for name, value in f1.state_dict().items():
            if name in f5_state and f5_state[name].shape == value.shape:
                f5_state[name] = value
        f5.load_state_dict(f5_state)
        states = torch.randn(3, 8)
        local = torch.randn(3, 6)
        local[:, -1] = 1
        people = torch.randn(3, 3, 6)
        mask = torch.tensor(
            [[True, True, False], [True, False, False], [False, False, False]]
        )
        people[~mask] = 0
        baseline = f1(states, local, people, mask)
        self.assertTrue(torch.equal(f5(states, local, people, mask), baseline))
        self.assertEqual(f5.person_interaction_gates(), {"audio_person": 0.0})

        with torch.no_grad():
            assert f5.audio_person_gate is not None
            f5.audio_person_gate.fill_(0.25)
        logits = f5(states, local, people, mask)
        self.assertFalse(torch.equal(logits, baseline))
        permutation = torch.tensor([1, 0, 2])
        self.assertTrue(
            torch.allclose(
                logits,
                f5(states, local, people[:, permutation], mask[:, permutation]),
                atol=1e-6,
                rtol=1e-6,
            )
        )
        logits.square().mean().backward()
        self.assertIsNotNone(f5.audio_person_attention.in_proj_weight.grad)
        self.assertIsNotNone(f5.audio_person_gate.grad)

    def test_f6_screening_adapter_is_zero_gated_and_visual_only_when_valid(self):
        torch.manual_seed(13)
        head = TurnPredictionHead(
            hidden_size=8,
            local_visual_dim=6,
            local_audio_dim=5,
            local_projection_dim=3,
            local_fusion="gated-person-state-adapter",
            person_attention_dim=4,
            person_attention_heads=2,
            dropout=0,
        ).eval()
        states = torch.randn(3, 8)
        local = torch.randn(3, 6)
        local[:, -1] = 1
        people = torch.randn(3, 2, 6)
        mask = torch.tensor([[True, True], [True, False], [False, False]])
        people[~mask] = 0
        zero_gate = head(states, local, people, mask)
        changed = people.clone()
        changed[mask] += torch.arange(6, dtype=changed.dtype) * 10
        # Before training, person memory is an exact zero residual.
        self.assertTrue(torch.equal(head(states, local, changed, mask), zero_gate))
        self.assertEqual(
            head.person_interaction_gates(),
            {"person_state_adapter": 0.0},
        )

        with torch.no_grad():
            assert head.person_state_adapter_gate is not None
            head.person_state_adapter_gate.fill_(0.25)
        logits = head(states, local, people, mask)
        self.assertFalse(torch.equal(head(states, local, changed, mask), logits))
        # An empty set must remain visually inert even after opening the gate.
        empty_changed = people.clone()
        empty_changed[2] = 1000
        self.assertTrue(
            torch.equal(
                logits[2],
                head(states, local, empty_changed, mask)[2],
            )
        )
        permutation = torch.tensor([1, 0])
        self.assertTrue(
            torch.allclose(
                logits,
                head(states, local, people[:, permutation], mask[:, permutation]),
                atol=1e-6,
                rtol=1e-6,
            )
        )
        logits.square().mean().backward()
        self.assertIsNotNone(head.person_attention.in_proj_weight.grad)
        self.assertIsNotNone(head.person_state_adapter[-1].weight.grad)
        self.assertIsNotNone(head.person_state_adapter_gate.grad)

        config = {
            "hidden_size": 8,
            "local_visual_dim": 6,
            "local_audio_dim": 5,
            "local_projection_dim": 3,
            "local_fusion": "gated-person-state-adapter",
            "person_attention_dim": 4,
            "person_attention_heads": 2,
            "dropout": 0,
        }
        restored = TurnPredictionHead(
            **checkpoint_head_config({"head_config": config})
        )
        restored.load_state_dict(head.state_dict())
        restored.eval()
        self.assertTrue(torch.equal(restored(states, local, people, mask), logits))

    def test_track_gru_residual_starts_from_exact_pooled_logits(self):
        torch.manual_seed(17)
        common = {
            "hidden_size": 8,
            "num_input_layers": 4,
            "layer_fusion": "concat-mlp",
            "local_visual_dim": 6,
            "local_audio_dim": 5,
            "local_projection_dim": 3,
            "dropout": 0,
        }
        pooled = TurnPredictionHead(**common, local_fusion="pooled").eval()
        track = TurnPredictionHead(
            **common,
            local_fusion="track-gru-residual",
            person_attention_dim=4,
            person_attention_heads=2,
            person_history_frames=3,
        ).eval()
        track_state = track.state_dict()
        for name, value in pooled.state_dict().items():
            self.assertIn(name, track_state)
            self.assertEqual(track_state[name].shape, value.shape)
            track_state[name] = value
        track.load_state_dict(track_state)

        states = torch.randn(3, 4, 8)
        local = torch.randn(3, 13)
        local[:, 6] = 1
        local[:, 12] = 1
        histories = torch.randn(3, 2, 3, 6)
        mask = torch.tensor(
            [
                [[True, True, True], [True, True, False]],
                [[True, False, False], [False, False, False]],
                [[False, False, False], [False, False, False]],
            ]
        )
        histories[~mask] = 0
        baseline = pooled(states, local)
        self.assertTrue(torch.equal(track(states, local, histories, mask), baseline))
        self.assertEqual(track.person_interaction_gates(), {"person_track_residual": 0.0})

        with torch.no_grad():
            assert track.person_track_residual_gate is not None
            track.person_track_residual_gate.fill_(.25)
        logits = track(states, local, histories, mask)
        self.assertFalse(torch.equal(logits[:2], baseline[:2]))
        # With no current person the correction must be exactly zero.
        self.assertTrue(torch.equal(logits[2], baseline[2]))
        permutation = torch.tensor([1, 0])
        self.assertTrue(
            torch.allclose(
                logits,
                track(states, local, histories[:, permutation], mask[:, permutation]),
                atol=1e-6,
                rtol=1e-6,
            )
        )
        changed_padding = histories.clone()
        changed_padding[~mask] = 1000
        self.assertTrue(torch.equal(logits, track(states, local, changed_padding, mask)))
        logits.square().mean().backward()
        self.assertIsNotNone(track.person_track_gru.weight_ih_l0.grad)
        self.assertIsNotNone(track.person_track_delta[-1].weight.grad)
        self.assertIsNotNone(track.person_track_residual_gate.grad)

        config = {
            **common,
            "local_fusion": "track-gru-residual",
            "person_attention_dim": 4,
            "person_attention_heads": 2,
            "person_history_frames": 3,
        }
        restored = TurnPredictionHead(
            **checkpoint_head_config({"head_config": config})
        ).eval()
        restored.load_state_dict(track.state_dict())
        self.assertTrue(torch.equal(restored(states, local, histories, mask), logits))

    def test_causal_person_adapter_v2_learns_after_exact_pooled_start(self):
        common = {
            "hidden_size": 8,
            "num_input_layers": 4,
            "layer_fusion": "concat-mlp",
            "local_visual_dim": 6,
            "local_audio_dim": 5,
            "local_projection_dim": 3,
            "dropout": 0,
        }
        for mode in (
            "causal-person-adapter-v2",
            "causal-person-adapter-egospeak-gru",
        ):
            with self.subTest(mode=mode):
                torch.manual_seed(23)
                pooled = TurnPredictionHead(
                    **common,
                    local_fusion="pooled",
                ).eval()
                adapter = TurnPredictionHead(
                    **common,
                    local_fusion=mode,
                    person_attention_dim=4,
                    person_attention_heads=2,
                    person_history_frames=3,
                ).eval()
                adapter_state = adapter.state_dict()
                for name, value in pooled.state_dict().items():
                    self.assertIn(name, adapter_state)
                    self.assertEqual(adapter_state[name].shape, value.shape)
                    adapter_state[name] = value
                adapter.load_state_dict(adapter_state)
                assert adapter.person_track_residual_output is not None
                self.assertEqual(
                    torch.count_nonzero(
                        adapter.person_track_residual_output.weight
                    ),
                    0,
                )
                self.assertEqual(
                    torch.count_nonzero(
                        adapter.person_track_residual_output.bias
                    ),
                    0,
                )
                frozen_parameters = {}
                for name, parameter in adapter.named_parameters():
                    if name == "layer_logits" or name.startswith(
                        ("layer_fusion_net.", "local_projections.", "net.")
                    ):
                        parameter.requires_grad_(False)
                        frozen_parameters[name] = parameter.detach().clone()

                states = torch.randn(3, 4, 8)
                local = torch.randn(3, 13)
                local[:, 6] = 1
                local[:, 12] = 1
                histories = torch.randn(3, 2, 3, 6)
                mask = torch.tensor(
                    [
                        [[True, True, True], [True, True, False]],
                        [[True, False, False], [False, False, False]],
                        [[False, False, False], [False, False, False]],
                    ]
                )
                histories[~mask] = 0
                baseline = pooled(states, local)
                initial = adapter(states, local, histories, mask)
                self.assertTrue(torch.equal(initial, baseline))
                self.assertEqual(
                    adapter.person_adapter_diagnostics()["delta_logit_max_abs"],
                    0.0,
                )

                optimizer = torch.optim.SGD(adapter.parameters(), lr=0.1)
                target = torch.randn_like(initial)
                (initial - target).square().mean().backward()
                output_gradient = (
                    adapter.person_track_residual_output.weight.grad
                )
                self.assertIsNotNone(output_gradient)
                self.assertGreater(int(torch.count_nonzero(output_gradient)), 0)
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)

                trained = adapter(states, local, histories, mask)
                (trained - target).square().mean().backward()
                input_projection = (
                    adapter.person_track_input_embedding[0].weight
                    if adapter.person_track_input_embedding is not None
                    else adapter.person_input_projection[1].weight
                )
                for parameter in (
                    input_projection,
                    adapter.person_track_gru.weight_ih_l0,
                    adapter.person_track_gru.weight_hh_l0,
                    adapter.person_attention.in_proj_weight,
                    adapter.person_track_local_state[1].weight,
                    adapter.person_track_dynamic_gate[1].weight,
                ):
                    self.assertIsNotNone(parameter.grad)
                    self.assertTrue(torch.isfinite(parameter.grad).all())
                    self.assertGreater(int(torch.count_nonzero(parameter.grad)), 0)
                for name, expected in frozen_parameters.items():
                    self.assertTrue(
                        torch.equal(dict(adapter.named_parameters())[name], expected),
                        name,
                    )

                # Empty rows remain exactly on the pooled path after the
                # residual opens.
                self.assertTrue(
                    torch.equal(trained[2], pooled(states, local)[2])
                )
                permutation = torch.tensor([1, 0])
                self.assertTrue(
                    torch.allclose(
                        trained,
                        adapter(
                            states,
                            local,
                            histories[:, permutation],
                            mask[:, permutation],
                        ),
                        atol=1e-6,
                        rtol=1e-6,
                    )
                )
                changed_padding = histories.clone()
                changed_padding[~mask] = 1000
                self.assertTrue(
                    torch.equal(
                        trained,
                        adapter(states, local, changed_padding, mask),
                    )
                )

                config = {
                    **common,
                    "local_fusion": mode,
                    "person_attention_dim": 4,
                    "person_attention_heads": 2,
                    "person_history_frames": 3,
                }
                restored = TurnPredictionHead(
                    **checkpoint_head_config({"head_config": config})
                ).eval()
                restored.load_state_dict(adapter.state_dict())
                self.assertTrue(
                    torch.equal(
                        restored(states, local, histories, mask),
                        adapter(states, local, histories, mask),
                    )
                )

    def test_egospeak_style_track_gru_matches_explicit_reference(self):
        torch.manual_seed(29)
        head = TurnPredictionHead(
            hidden_size=8,
            local_visual_dim=6,
            local_audio_dim=5,
            local_projection_dim=3,
            local_fusion="causal-person-adapter-egospeak-gru",
            person_attention_dim=4,
            person_attention_heads=2,
            person_history_frames=3,
            dropout=0,
        ).eval()
        embedding = head.person_track_input_embedding
        self.assertIsNotNone(embedding)
        self.assertEqual(
            [type(module) for module in embedding],
            [
                torch.nn.Linear,
                torch.nn.LayerNorm,
                torch.nn.ReLU,
                torch.nn.Dropout,
            ],
        )
        self.assertEqual(head.person_track_gru.num_layers, 1)
        self.assertFalse(head.person_track_gru.bidirectional)
        self.assertTrue(head.person_track_gru.batch_first)
        self.assertIsInstance(
            head.person_track_output_activation,
            torch.nn.ReLU,
        )

        histories = torch.randn(2, 2, 3, 6)
        mask = torch.tensor(
            [
                [[True, True, True], [True, True, False]],
                [[True, False, False], [False, False, False]],
            ]
        )
        histories[~mask] = 1234
        actual, valid = head._encode_person_track_histories(histories, mask)
        expected = torch.zeros_like(actual)
        for batch in range(2):
            for person in range(2):
                length = int(mask[batch, person].sum())
                if not length:
                    continue
                sequence = embedding(
                    histories[batch, person, :length]
                )[None]
                initial = sequence.new_zeros(1, 1, 4)
                _, hidden = head.person_track_gru(sequence, initial)
                expected[batch, person] = torch.relu(hidden[-1, 0])
        self.assertTrue(torch.equal(valid, mask.any(dim=-1)))
        self.assertTrue(torch.allclose(actual, expected, atol=1e-7, rtol=0))
        self.assertGreaterEqual(float(actual.min()), 0.0)

        changed_padding = histories.clone()
        changed_padding[~mask] = -9876
        changed, _ = head._encode_person_track_histories(
            changed_padding,
            mask,
        )
        self.assertTrue(torch.equal(actual, changed))

    def test_head_shapes_gradients_missing_and_roundtrip(self):
        states = torch.randn(3, 4, 8)
        for v, a in ((0, 0), (6, 0), (0, 5), (6, 5)):
            for projection in (0, 3):
                head = TurnPredictionHead(hidden_size=8, num_input_layers=4, layer_fusion="concat-mlp",
                                          local_visual_dim=v, local_audio_dim=a, local_projection_dim=projection,
                                          transition_auxiliary=True, dropout=0)
                width = sum(d + 1 for d in (v, a) if d)
                local = torch.randn(3, width)
                offset = 0
                for d in (v, a):
                    if d:
                        local[:, offset + d] = torch.tensor([0., 1., 1.])
                        offset += d + 1
                logits, transition = head.forward_with_transition(states, local)
                self.assertEqual(logits.shape, (3, 10, 3))
                self.assertEqual(transition.shape, (3, 10))
                (logits.square().mean() + transition.square().mean()).backward()
                self.assertIsNotNone(head.net[1].weight.grad)
                if width:
                    # Changing invalid raw features cannot change the readout.
                    changed = local.clone()
                    offset = 0
                    for d in (v, a):
                        if d:
                            changed[0, offset:offset + d] = 1000
                            offset += d + 1
                    self.assertTrue(torch.equal(head(states, local)[0], head(states, changed)[0]))
                    with self.assertRaisesRegex(ValueError, "requires aligned"):
                        head(states)
                config = {"hidden_size": 8, "num_input_layers": 4, "layer_fusion": "concat-mlp",
                          "local_visual_dim": v, "local_audio_dim": a, "local_projection_dim": projection,
                          "transition_auxiliary": True, "dropout": 0}
                restored = TurnPredictionHead(**checkpoint_head_config({"head_config": config}))
                restored.load_state_dict(head.state_dict())
                self.assertTrue(torch.equal(restored(states, local), head(states, local)))

    def test_legacy_head_parameter_keys_and_values(self):
        torch.manual_seed(42)
        old = torch.nn.Sequential(torch.nn.LayerNorm(8), torch.nn.Linear(8, 4), torch.nn.GELU(),
                                  torch.nn.Dropout(.1), torch.nn.Linear(4, 30))
        torch.manual_seed(42)
        head = TurnPredictionHead(hidden_size=8)
        self.assertEqual(set(head.state_dict()), {f"net.{key}" for key in old.state_dict()})
        for key, value in old.state_dict().items():
            self.assertTrue(torch.equal(value, head.state_dict()[f"net.{key}"]))

    def test_sparse_selection_keeps_local_rows(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            for folder in ("features", "local", "targets"):
                (root / folder).mkdir()
            reference, side, targets = fixture()
            torch.save(reference, root / "features/Session_4.pt")
            torch.save(side, root / "local/Session_4.pt")
            np.save(root / "targets/Session_4.npy", targets)
            data = _load_examples(root / "features", root / "targets", ["Session_4"], "open-stream-replay",
                                  "easycom", {"Session_4": {1, 4}}, root / "local", "audio")
            self.assertTrue(torch.equal(data[-2][:, :5], side["audio"][[1, 4]]))
            self.assertIsNone(data[-1])
            self.assertTrue(torch.equal(data[0], reference["states"][[1, 4]]))
            self.assertTrue(torch.equal(data[1][:, 0], torch.from_numpy(targets[[2, 5]])))

    def test_source_inventory_rejects_different_clip_order(self):
        source = {"path": "/old/00-00-000.mp4", "source_frames": 1200, "dense_start": 0., "duration": 60.}
        media = SimpleNamespace(clips=[object()], metadata=lambda: {"source_clips": [source]})
        validate_source_timeline({"direct_media": {"source_clips": [source]}}, media)
        bad = {**source, "path": "/old/01-00-000.mp4"}
        with self.assertRaisesRegex(ValueError, "timeline differs"):
            validate_source_timeline({"direct_media": {"source_clips": [bad]}}, media)

    def test_config_digest_ignores_only_checkpoint_location(self):
        a = {"_name_or_path": "/old/model", "audio_config": {"_name_or_path": "/old/model", "output_dim": 8}}
        b = {"_name_or_path": "/new/model", "audio_config": {"_name_or_path": "/new/model", "output_dim": 8}}
        self.assertEqual(model_config_digest(a), model_config_digest(b))
        b["audio_config"]["output_dim"] = 16
        self.assertNotEqual(model_config_digest(a), model_config_digest(b))

    def test_encoder_image_pool_splits_people_before_mean(self):
        encoder = FrozenQwenLocalEncoder.__new__(FrozenQwenLocalEncoder)
        encoder.visual_dim, encoder.device, encoder.image_batch_size, encoder.max_pixels = 2, torch.device("cpu"), 8, 50176
        encoder.processor = SimpleNamespace(image_processor=lambda **kwargs: {
            "pixel_values": torch.ones(12, 2), "image_grid_thw": torch.tensor([[1, 2, 2], [1, 2, 4]])})
        encoder.model = SimpleNamespace(visual=SimpleNamespace(dtype=torch.float32),
            config=SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2)),
            get_image_features=lambda *_: torch.tensor([[1., 1.], [3., 3.], [5., 5.]]))
        actual = encoder.images([object(), object()])
        self.assertTrue(torch.equal(actual, torch.tensor([[1., 1.], [4., 4.]])))
        # A future API returning per-image tensors must not silently change
        # the pooling rule or crash with an uninformative AttributeError.
        encoder.model.get_image_features = lambda *_: (torch.ones(1, 2), torch.ones(2, 2))
        with self.assertRaisesRegex(RuntimeError, "Unexpected image encoder output"):
            encoder.images([object(), object()])

    def test_extraction_saves_runtime_version_and_rejects_mixed_versions(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            for folder in ("features", "targets", "local"):
                (root / folder).mkdir()
            reference, side, labels = fixture()
            torch.save(reference, root / "features/Session_4.pt")
            np.save(root / "targets/Session_4.npy", labels)
            media = SimpleNamespace(clips=[], metadata=lambda: {})
            argv = ["extract_turn_local_features.py", "--data-root", str(root),
                    "--feature-root", str(root / "features"), "--target-root", str(root / "targets"),
                    "--output-root", str(root / "local"), "--samples", "4", "--overwrite"]
            contracts = []
            for version in TESTED_TRANSFORMERS_VERSIONS:
                fake_encoder = SimpleNamespace(transformers_version=version, visual_dim=6, audio_dim=5,
                                               model_config_sha256="same-checkpoint")
                with patch.object(sys, "argv", argv), \
                     patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__=version)}), \
                     patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])), \
                     patch("extract_turn_local_features.FrozenQwenLocalEncoder", return_value=fake_encoder), \
                     patch("extract_turn_local_features.extract_local", return_value=copy.deepcopy(side)):
                    extract_main()
                output = root / "local/Session_4.pt"
                _, contract = load_local_features(output, reference, sample_id="Session_4",
                                                   dataset="easycom", modalities="av")
                self.assertEqual(contract["config"]["transformers"], version)
                contracts.append(contract)
            with self.assertRaisesRegex(ValueError, "contract differs"):
                load_local_features(output, reference, sample_id="Session_4", dataset="easycom",
                                    modalities="av", expected_contract=contracts[0])

    def test_unsupported_version_fails_before_media_and_preview_still_works(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            reference, _, labels = fixture()
            torch.save(reference, root / "Session_4.pt")
            np.save(root / "Session_4.npy", labels)
            argv = ["extract_turn_local_features.py", "--data-root", str(root),
                    "--feature-root", str(root), "--target-root", str(root),
                    "--output-root", str(root / "preview"), "--samples", "4"]
            with patch.object(sys, "argv", argv), \
                 patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="5.0.0")}), \
                 patch("extract_turn_local_features.build_media_and_boxes") as build:
                with self.assertRaisesRegex(RuntimeError, "found 5.0.0"):
                    extract_main()
                build.assert_not_called()
            media = SimpleNamespace(clips=[], metadata=lambda: {})
            with patch.object(sys, "argv", argv + ["--preview-only"]), \
                 patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="5.0.0")}), \
                 patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])), \
                 patch("extract_turn_local_features.preview_boxes") as preview, \
                 patch("extract_turn_local_features.FrozenQwenLocalEncoder") as encoder:
                extract_main()
                preview.assert_called_once()
                encoder.assert_not_called()

    def test_extraction_never_encodes_future_audio_and_saves_ragged_people(self):
        class Media:
            clips = [SimpleNamespace(crop=None, dense_start=0.)]
            def _source_frame_for_time(self, t):
                return 0, round(t * 20), t
            def frame_at_or_before(self, t):
                return Image.new("RGB", (16, 16))
            def audio_slice(self, start, end):
                return np.arange(round(start * 16000), round(end * 16000), dtype=np.float32)
        class Encoder:
            visual_dim = audio_dim = 2
            def __init__(self):
                self.ends = []
            def images(self, crops):
                return torch.ones(len(crops), 2)
            def audio(self, waveform):
                self.ends.append(float(waveform[-1]))
                return torch.tensor([float(waveform[0]), float(waveform[-1])])
        encoder = Encoder()
        args = SimpleNamespace(video_fps=2., visual_window_seconds=1., audio_window_seconds=1.,
                               box_coordinate_space="camera", crop_expansion=1., no_progress=True)
        annotations = [{0: [{"id": "1", "xyxy": [0., 0., 8., 8.]}]}]
        result = extract_local(Media(), annotations, encoder, torch.arange(1, 11, dtype=torch.float64) / 5, args)
        self.assertEqual(encoder.ends, [3200 * (i + 1) - 1 for i in range(10)])
        self.assertEqual(result["persons"]["frame_offsets"].tolist(), [0, 1, 1, 1, 1])
        self.assertTrue(result["visual_valid"][:5].all())
        self.assertFalse(result["visual_valid"][5:].any())
        self.assertTrue(result["audio_valid"].all())


@unittest.skipUnless(importlib.util.find_spec("transformers"), "optional: install a tested Transformers version")
class TinyQwenEncoderTests(unittest.TestCase):
    def setUp(self):
        import transformers
        if transformers.__version__ not in TESTED_TRANSFORMERS_VERSIONS:
            self.skipTest(f"Encoder API tests require one of {TESTED_TRANSFORMERS_VERSIONS}")
        torch.set_num_threads(1)

    def tiny_model_and_processor(self):
        from transformers import (
            Qwen2_5OmniThinkerConfig, Qwen2_5OmniThinkerForConditionalGeneration,
            Qwen2VLImageProcessor, WhisperFeatureExtractor,
        )
        config = Qwen2_5OmniThinkerConfig(
            audio_config=dict(d_model=16, encoder_layers=1, encoder_attention_heads=2,
                              encoder_ffn_dim=32, output_dim=8),
            vision_config=dict(depth=1, hidden_size=16, intermediate_size=32, num_heads=2,
                               out_hidden_size=8, fullatt_block_indexes=[0]),
            text_config=dict(vocab_size=32, hidden_size=8, intermediate_size=16,
                             num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2),
        )
        model = Qwen2_5OmniThinkerForConditionalGeneration(config).eval().requires_grad_(False)
        processor = SimpleNamespace(image_processor=Qwen2VLImageProcessor(),
                                    feature_extractor=WhisperFeatureExtractor(feature_size=128))
        return model, processor

    def test_real_encoder_apis_without_downloading_weights(self):
        encoder = FrozenQwenLocalEncoder.__new__(FrozenQwenLocalEncoder)
        encoder.model, encoder.processor = self.tiny_model_and_processor()
        encoder.model.model = None
        encoder.model.lm_head = None
        encoder.visual_dim = encoder.audio_dim = 8
        encoder.device, encoder.max_pixels, encoder.image_batch_size = torch.device("cpu"), 224 * 224, 8
        crops = [Image.new("RGB", (56, 84), "red"), Image.new("RGB", (112, 56), "blue")]
        batch = encoder.images(crops)
        single = encoder.images(crops[:1])
        self.assertEqual(batch.shape, (2, 8))
        self.assertTrue(torch.allclose(batch[0], single[0], atol=1e-5))
        for n in (3200, 6400, 16000):
            audio = encoder.audio(np.sin(np.arange(n, dtype=np.float32) / 10))
            self.assertEqual(audio.shape, (8,))
            self.assertTrue(torch.isfinite(audio).all())
            self.assertFalse(audio.requires_grad)

    @unittest.skipUnless(importlib.util.find_spec("accelerate"), "optional: install accelerate for checkpoint loading")
    def test_constructor_loads_local_checkpoint_and_preserves_encoder_weights(self):
        import transformers

        model, processor = self.tiny_model_and_processor()
        expected = {key: value.clone() for key, value in model.state_dict().items()
                    if key.startswith(("visual.", "audio_tower."))}
        with tempfile.TemporaryDirectory() as directory:
            model.save_pretrained(directory)
            args = SimpleNamespace(model_id=directory, device="cpu", dtype="float32",
                                   flash_attention_2=False, max_pixels=224 * 224, image_batch_size=8)
            with patch.object(transformers.Qwen2_5OmniProcessor, "from_pretrained", return_value=processor) as load:
                encoder = FrozenQwenLocalEncoder(args)
            load.assert_called_once_with(directory, use_fast=False)
            self.assertEqual(encoder.transformers_version, transformers.__version__)
            self.assertIsNone(encoder.model.model)
            self.assertIsNone(encoder.model.lm_head)
            self.assertFalse(encoder.model.training)
            self.assertFalse(any(p.requires_grad for p in encoder.model.parameters()))
            for key, value in encoder.model.state_dict().items():
                self.assertTrue(torch.equal(value, expected[key]), key)
            self.assertEqual(encoder.images([Image.new("RGB", (56, 84), "red")]).shape, (1, 8))
            self.assertEqual(encoder.audio(np.zeros(3200, dtype=np.float32)).shape, (8,))

    def test_constructor_rejects_incomplete_encoder_weights(self):
        import transformers

        args = SimpleNamespace(model_id="unused", device="cpu", dtype="float32", flash_attention_2=False)
        for info in ({"missing_keys": ["visual.patch_embed.proj.weight"]},
                     {"missing_keys": ["audio_tower.conv1.weight"]},
                     {"mismatched_keys": [("visual.weight", (2,), (3,))]},
                     {"error_msgs": ["broken checkpoint"]}):
            with self.subTest(info=info), \
                 patch.object(transformers.Qwen2_5OmniThinkerForConditionalGeneration,
                              "from_pretrained", return_value=(None, info)), \
                 self.assertRaisesRegex(RuntimeError, "incompletely loaded weights"):
                FrozenQwenLocalEncoder(args)


class RoundTripTests(unittest.TestCase):
    def test_train_infer_easycom_and_ego4d(self):
        repo = Path(__file__).resolve().parent
        env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1")
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            for folder in ("features", "local", "highrate", "targets"):
                (root / folder).mkdir()
            train_uid = "00000000-0000-0000-0000-000000000004"
            test_uid = "00000000-0000-0000-0000-000000000001"
            manifest = root / "ego4d.tsv"
            manifest.write_text(f"split\tuid\ntrain\t{train_uid}\ntest\t{test_uid}\n")
            for name in ("Session_1", "Session_4", "Session_5", train_uid, test_uid):
                dataset = "easycom" if name.startswith("Session") else "ego4d"
                reference, side, targets = fixture(name=name, dataset=dataset)
                torch.save(reference, root / "features" / f"{name}.pt")
                torch.save(side, root / "local" / f"{name}.pt")
                _, highrate_side, _ = highrate_fixture(
                    name=name,
                    dataset=dataset,
                )
                torch.save(highrate_side, root / "highrate" / f"{name}.pt")
                np.save(root / "targets" / f"{name}.npy", targets)
            common = ["--feature-root", str(root / "features"), "--target-root", str(root / "targets"),
                      "--device", "cpu", "--no-progress", "--expected-mode", "open-stream-replay"]
            for mode in (
                "baseline",
                "audio",
                "visual",
                "av",
                "track-gru",
                "person-adapter-v2",
                "egospeak-person-gru",
                "frame-set",
                "f5-screen",
                "f6-screen",
                "traditional-audio",
                "highrate-person-av",
                "ego4d",
            ):
                checkpoint = root / f"{mode}.pth"
                train = [sys.executable, str(repo / "train_turn_head.py"), *common,
                         "--epochs", "1", "--warmup-epochs", "0", "--batch-size", "4",
                         "--layer-fusion", "concat-mlp", "--output", str(checkpoint)]
                infer = [sys.executable, str(repo / "infer_turn_head.py"), *common,
                         "--checkpoint", str(checkpoint), "--output-dir", str(root / f"result_{mode}")]
                if mode == "ego4d":
                    train += ["--dataset", "ego4d", "--manifest", str(manifest)]
                    infer += ["--dataset", "ego4d", "--manifest", str(manifest)]
                else:
                    train += ["--sessions", "4", "5"]
                    infer += ["--sessions", "1"]
                if mode == "visual":
                    train += ["--dev-sessions", "5", "--local-projection-dim", "0"]
                if mode == "av":
                    train += ["--transition-objective", "r1-supervised"]
                person_mode = {
                    "track-gru": "track-gru-residual",
                    "person-adapter-v2": "causal-person-adapter-v2",
                    "egospeak-person-gru": (
                        "causal-person-adapter-egospeak-gru"
                    ),
                    "frame-set": "frame-set-attention",
                    "f5-screen": "audio-person-frame-set",
                    "f6-screen": "gated-person-state-adapter",
                    "highrate-person-av": "highrate-person-av-attention",
                }.get(mode)
                if person_mode is not None:
                    train += [
                        "--local-fusion", person_mode,
                        "--person-attention-dim", "4",
                        "--person-attention-heads", "2",
                    ]
                if mode in (
                    "track-gru",
                    "person-adapter-v2",
                    "egospeak-person-gru",
                ):
                    train += [
                        "--person-history-frames", "3",
                        "--person-track-max-gap-seconds", "1",
                        "--initialize-pooled-checkpoint", str(root / "av.pth"),
                        "--transition-objective", "r1-supervised",
                    ]
                if mode == "highrate-person-av":
                    train += [
                        "--initialize-pooled-checkpoint",
                        str(root / "traditional-audio.pth"),
                    ]
                if mode != "baseline":
                    modality = (
                        "av"
                        if mode in (
                            "ego4d", "track-gru", "person-adapter-v2",
                            "egospeak-person-gru",
                            "frame-set", "f5-screen", "f6-screen",
                            "highrate-person-av",
                        )
                        else "audio"
                        if mode == "traditional-audio"
                        else mode
                    )
                    local_root = root / (
                        "highrate"
                        if mode in ("traditional-audio", "highrate-person-av")
                        else "local"
                    )
                    train += ["--local-feature-root", str(local_root), "--local-modalities", modality]
                    infer += ["--local-feature-root", str(local_root)]
                for command in (train, infer):
                    completed = subprocess.run(command, cwd=repo, env=env, capture_output=True, text=True, timeout=60)
                    self.assertEqual(completed.returncode, 0, f"{mode}: {completed.stdout}\n{completed.stderr}")
                metrics = json.loads((root / f"result_{mode}/metrics.json").read_text())
                self.assertEqual(metrics["checkpoint_epoch"], 1)
                self.assertEqual(bool(metrics["local_contract"]), mode != "baseline")
                if mode == "av":
                    self.assertTrue(metrics["transition_auxiliary"])
                if person_mode is not None:
                    self.assertEqual(metrics["local_fusion"], person_mode)
                    self.assertEqual(
                        metrics["local_fusion_scope"],
                        LOCAL_FUSION_SCOPES[person_mode],
                    )
            # The reporting command uses the same metric keys as inference.
            completed = subprocess.run(
                [sys.executable, str(repo / "summarize_local_baselines.py"), "--root", str(root)],
                cwd=repo, capture_output=True, text=True, timeout=20,
            )
            self.assertEqual(completed.returncode, 0, completed.stderr)
            self.assertIn("Mean ± sample standard deviation", completed.stdout)


if __name__ == "__main__":
    unittest.main()