#!/usr/bin/env python
"""Train a compact offline camera-directing baseline from exported Unity timelines."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import random
import re
import sys
import time
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Sequence

import librosa
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset

import camera_runtime_data as runtime_data


PIPELINE_VERSION = 1
AUDIO_RATE = 24_000
AUDIO_HOP = 400  # Exactly 60 audio feature frames per second.
MEL_BANDS = 32
REVISION_DATASET_PATTERN = re.compile(
    r"^(?P<base>.+)_rev_(?P<fingerprint>[0-9a-f]{16})$"
)
RAW_TRAINING_FILE_FIELDS = (
    "jointsFile",
    "cameraFile",
    "rootFile",
    "shotIndexFile",
    "timeFile",
    "audioFile",
)
LFS_POINTER_PREFIX = b"version https://git-lfs.github.com/spec/v1"
FRAME_MASK_TRAINING_PURPOSES = ("jointTcn", "trajectoryModel", "cutModel")
SAFE_CATALOG_BOOLEAN_FIELDS = (
    "trainingEligible",
    "singleSubjectTrainingEligible",
    "selectedAnimatorTargetMatchKnown",
    "allTargetsMatchSelectedExporterAnimator",
    "exportedCharacterPathMatchKnown",
    "exportedCharacterPathMatchesSelectedAnimator",
)

# A shared semantic skeleton for the Humanoid and Generic rigs in the current data.
CANONICAL_JOINT_ALIASES = {
    "pelvis": ("hips", "pelvis"),
    "spine": ("spine",),
    "chest": ("chest", "upperchest", "spine1"),
    "neck": ("neck",),
    "head": ("head",),
    "left_shoulder": ("leftshoulder", "lclavicle"),
    "left_upper_arm": ("leftupperarm", "lupperarm"),
    "left_lower_arm": ("leftlowerarm", "lforearm"),
    "left_hand": ("lefthand", "lhand"),
    "right_shoulder": ("rightshoulder", "rclavicle"),
    "right_upper_arm": ("rightupperarm", "rupperarm"),
    "right_lower_arm": ("rightlowerarm", "rforearm"),
    "right_hand": ("righthand", "rhand"),
    "left_upper_leg": ("leftupperleg", "lthigh"),
    "left_lower_leg": ("leftlowerleg", "lcalf"),
    "left_foot": ("leftfoot", "lfoot"),
    "left_toes": ("lefttoes", "ltoe0"),
    "right_upper_leg": ("rightupperleg", "rthigh"),
    "right_lower_leg": ("rightlowerleg", "rcalf"),
    "right_foot": ("rightfoot", "rfoot"),
    "right_toes": ("righttoes", "rtoe0"),
}
CANONICAL_JOINTS = tuple(CANONICAL_JOINT_ALIASES)


def configure_utf8_stdio() -> None:
    for stream in (sys.stdout, sys.stderr):
        reconfigure = getattr(stream, "reconfigure", None)
        if reconfigure is not None:
            reconfigure(encoding="utf-8", errors="backslashreplace")


@dataclass
class SongRecord:
    manifest_path: Path
    dataset_dir: Path
    project_root: Path
    manifest: dict
    song: dict

    @property
    def name(self) -> str:
        return self.song["songName"]

    @property
    def id(self) -> str:
        scene_path = str(self.manifest.get("scenePath", ""))
        folder_name = str(self.song["folderName"])
        value = f"{scene_path}\0{self.name}\0{folder_name}"
        if self.is_revision:
            value += f"\0revision:{self.dataset_id}"
        value = value.encode("utf-8")
        return hashlib.sha256(value).hexdigest()[:16]

    @property
    def dataset_id(self) -> str:
        return self.dataset_dir.name

    @property
    def revision_of_dataset_id(self) -> str | None:
        match = REVISION_DATASET_PATTERN.fullmatch(self.dataset_id)
        return match.group("base") if match is not None else None

    @property
    def is_revision(self) -> bool:
        return self.revision_of_dataset_id is not None

    @property
    def song_dir(self) -> Path:
        return self.dataset_dir / self.song["folderName"]

    @property
    def audio_path(self) -> Path:
        portable_audio = self.song.get("audioFile")
        if portable_audio:
            return self.song_dir / portable_audio
        project_root_override = os.environ.get("CWAI_UNITY_PROJECT_ROOT")
        if project_root_override:
            return Path(project_root_override).resolve() / self.song["audioAssetPath"]
        return self.project_root / self.song["audioAssetPath"]

    @property
    def portable_prepared_path(self) -> Path:
        file_name = self.song.get("preparedFile", "prepared_v1.npz")
        return self.song_dir / file_name

    @property
    def is_generation_input(self) -> bool:
        """Whether this record is a camera-free generation target.

        Generation manifests deliberately live outside the authored-camera
        training catalog.  Accept the fields at either manifest or song scope
        so the single-song Unity exporter can keep common metadata at the
        document root while older tooling can still emit per-song metadata.
        """

        input_kind = self.song.get(
            "inputKind",
            self.manifest.get("inputKind"),
        )
        has_authored_camera = self.song.get(
            "hasAuthoredCamera",
            self.manifest.get("hasAuthoredCamera"),
        )
        return (
            str(input_kind or "").strip().lower() == "generation_input"
            and has_authored_camera is False
        )


@dataclass
class PreparedSong:
    name: str
    inputs: np.ndarray
    targets: np.ndarray
    cuts: np.ndarray
    times: np.ndarray
    body_scale: float
    source_manifest: str


@dataclass(frozen=True)
class AudioAlignedFrameMask:
    """Validated training-index frame mask for one logical song.

    Ranges use zero-based, half-open ``[start, end)`` frame coordinates and
    refer to the original prepared arrays.  The arrays are never sliced or
    rewritten as part of applying this contract.
    """

    song_id: str
    declared_frame_count: int
    ranges: tuple[tuple[int, int], ...]
    audio_aligned_end_frame: int | None
    audio_fingerprint: str | None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data",
        nargs="*",
        default=[],
        help="Dataset folders or dataset_manifest.json files. Defaults to DatasetExports.",
    )
    parser.add_argument("--val-song", default="", help="Song name held out for validation.")
    parser.add_argument(
        "--val-song-id",
        default="",
        help="Hold out one 16-character song ID instead of every matching song name.",
    )
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--window", type=int, default=240, help="Frames per training window.")
    parser.add_argument("--train-stride", type=int, default=60)
    parser.add_argument("--hidden-size", type=int, default=128)
    parser.add_argument("--layers", type=int, default=6)
    parser.add_argument("--dropout", type=float, default=0.15)
    parser.add_argument("--learning-rate", type=float, default=1e-3)
    parser.add_argument(
        "--patience",
        type=int,
        default=5,
        help="Stop after this many epochs without validation improvement.",
    )
    parser.add_argument("--seed", type=int, default=20260729)
    parser.add_argument("--num-workers", type=int, default=0)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=Path(__file__).resolve().parent / "cache",
    )
    parser.add_argument(
        "--skip-unavailable",
        action="store_true",
        help=(
            "Skip raw-only songs that have neither valid preparation inputs "
            "nor portable prepared tensors."
        ),
    )
    parser.add_argument(
        "--training-index",
        type=Path,
        help=(
            "Optional reports/training_index.json. Training uses only the "
            "selected purpose in the train split plus the requested validation song."
        ),
    )
    parser.add_argument(
        "--training-purpose",
        choices=FRAME_MASK_TRAINING_PURPOSES,
        default="jointTcn",
    )
    parser.add_argument(
        "--single-subject-only",
        action="store_true",
        help=(
            "Train only on songs marked singleSubjectTrainingEligible=true "
            "in the catalog supplied with --catalog."
        ),
    )
    parser.add_argument(
        "--catalog",
        type=Path,
        help=(
            "Catalog JSON used by --single-subject-only. It must match the "
            "current dataset manifests and training inputs."
        ),
    )
    parser.add_argument(
        "--run-dir",
        type=Path,
        default=Path(__file__).resolve().parent / "runs" / time.strftime("%Y%m%d_%H%M%S"),
    )
    parser.add_argument("--rebuild-cache", action="store_true")
    args = parser.parse_args()
    if args.single_subject_only and args.catalog is None:
        parser.error(
            "--single-subject-only requires an explicit --catalog path."
        )
    if args.single_subject_only and args.training_index is None:
        parser.error(
            "--single-subject-only requires an explicit --training-index path "
            "so evaluation groups cannot leak across the training split."
        )
    if args.catalog is not None and not args.single_subject_only:
        parser.error(
            "--catalog is only used together with --single-subject-only."
        )
    return args


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def find_manifest_paths(values: Iterable[str]) -> list[Path]:
    paths: list[Path] = []
    requested = [Path(value).resolve() for value in values]
    if not requested:
        requested = [(Path.cwd() / "DatasetExports").resolve()]

    for path in requested:
        if path.is_file() and path.name == "dataset_manifest.json":
            paths.append(path)
        elif (path / "dataset_manifest.json").is_file():
            paths.append(path / "dataset_manifest.json")
        elif path.is_dir():
            paths.extend(path.glob("TimelineCamera_*/dataset_manifest.json"))
            paths.extend(path.glob("GenerationInput_*/dataset_manifest.json"))

    unique = sorted(set(path.resolve() for path in paths))
    if not unique:
        raise FileNotFoundError("No dataset_manifest.json files were found.")
    return unique


def load_records(
    manifest_paths: Iterable[Path],
    skip_unavailable: bool = False,
) -> list[SongRecord]:
    records: list[SongRecord] = []
    for manifest_path in manifest_paths:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        if manifest["sampleRate"] != 60:
            raise ValueError(f"{manifest_path} is not sampled at 60 FPS.")
        dataset_dir = manifest_path.parent
        project_root = dataset_dir.parent.parent
        for song in manifest["songs"]:
            record = SongRecord(
                manifest_path=manifest_path,
                dataset_dir=dataset_dir,
                project_root=project_root,
                manifest=manifest,
                song=song,
            )
            has_audio = record.audio_path.is_file()
            has_prepared = record.portable_prepared_path.is_file()
            unavailable_reason = ""
            if record.is_generation_input:
                required_target_paths = (
                    record.audio_path,
                    record.song_dir / "joints_world.f32",
                    record.song_dir / "root.f32",
                    record.song_dir / "time.f64",
                )
                missing_target_paths = [
                    str(path)
                    for path in required_target_paths
                    if not path.is_file()
                ]
                if missing_target_paths:
                    unavailable_reason = (
                        "Generation input is missing motion/audio sources: "
                        + ", ".join(missing_target_paths)
                    )
                else:
                    try:
                        canonical_joint_indices(manifest["sharedBones"])
                    except (KeyError, ValueError) as error:
                        unavailable_reason = str(error)
            elif not has_audio and not has_prepared:
                unavailable_reason = (
                    "Neither source audio nor portable prepared data was found: "
                    f"{record.audio_path}"
                )
            elif has_audio and not has_prepared:
                try:
                    canonical_joint_indices(manifest["sharedBones"])
                except (KeyError, ValueError) as error:
                    unavailable_reason = str(error)
            if unavailable_reason:
                if skip_unavailable:
                    continue
                raise FileNotFoundError(unavailable_reason)
            records.append(record)
    if len(records) < 2:
        raise ValueError("At least two songs are required for a song-level validation split.")
    return records


def select_record(
    records: Iterable[SongRecord],
    song_name: str = "",
    song_id: str = "",
    dataset_id: str = "",
) -> SongRecord:
    candidates = list(records)
    if song_name:
        candidates = [record for record in candidates if record.name == song_name]
    if song_id:
        candidates = [record for record in candidates if record.id == song_id]
    if dataset_id:
        candidates = [
            record for record in candidates if record.dataset_id == dataset_id
        ]

    selector_parts = []
    if song_name:
        selector_parts.append(f"name={song_name!r}")
    if song_id:
        selector_parts.append(f"id={song_id!r}")
    if dataset_id:
        selector_parts.append(f"dataset={dataset_id!r}")
    selector = ", ".join(selector_parts) or "(no selector)"

    if len(candidates) == 1:
        return candidates[0]
    if not candidates:
        raise ValueError(f"No song matched {selector}.")

    preview = "\n".join(
        (
            f"  {record.id} | {record.name} | "
            f"{record.dataset_id}/{record.song['folderName']}"
        )
        for record in candidates[:20]
    )
    suffix = ""
    if len(candidates) > 20:
        suffix = f"\n  ... and {len(candidates) - 20} more"
    raise ValueError(
        f"Song selector {selector} matched {len(candidates)} records. "
        "Add --song-id or --dataset-id to select one:\n"
        f"{preview}{suffix}"
    )


def filter_records_by_training_index(
    records: Iterable[SongRecord],
    index_path: Path,
    purpose: str,
    validation_ids: set[str],
) -> list[SongRecord]:
    _, entries = load_training_index_entries(index_path)
    missing_evaluation_groups = sorted(
        song_id
        for song_id, entry in entries.items()
        if not str(entry.get("evaluationGroupId", "")).strip()
    )
    if missing_evaluation_groups:
        raise ValueError(
            "Training index entries are missing evaluationGroupId: "
            + ", ".join(missing_evaluation_groups[:20])
        )
    group_splits: dict[str, set[str]] = {}
    for entry in entries.values():
        evaluation_group = str(entry["evaluationGroupId"])
        group_splits.setdefault(evaluation_group, set()).add(
            str(entry.get("recommendedSplit", ""))
        )
    conflicting_groups = {
        group: sorted(splits)
        for group, splits in group_splits.items()
        if len(splits) > 1
    }
    if conflicting_groups:
        raise ValueError(
            "Training index has evaluation groups spanning splits: "
            + json.dumps(
                conflicting_groups,
                ensure_ascii=False,
                sort_keys=True,
            )
        )
    ineligible_validation = sorted(
        song_id
        for song_id in validation_ids
        if song_id in entries
        and entries[song_id].get("usage", {}).get(purpose) is not True
    )
    if ineligible_validation:
        raise ValueError(
            f"Validation songs are not eligible for training purpose "
            f"{purpose!r}: "
            + ", ".join(ineligible_validation)
        )
    validation_groups = {
        str(entries[song_id].get("evaluationGroupId", ""))
        for song_id in validation_ids
        if song_id in entries
        and entries[song_id].get("evaluationGroupId")
    }
    selected: list[SongRecord] = []
    missing: list[str] = []
    for record in records:
        entry = entries.get(record.id)
        if entry is None:
            missing.append(record.id)
            continue
        is_validation = record.id in validation_ids
        is_training = (
            bool(entry.get("usage", {}).get(purpose, False))
            and entry.get("recommendedSplit") == "train"
            and str(entry.get("evaluationGroupId", ""))
            not in validation_groups
        )
        if is_validation or is_training:
            selected.append(record)

    if missing:
        raise ValueError(
            f"Training index is missing {len(missing)} available song IDs. "
            "Regenerate it with tools/build_training_index.py."
        )
    selected_ids = {record.id for record in selected}
    absent_validation = validation_ids - selected_ids
    if absent_validation:
        raise ValueError(
            "Validation songs are missing from the training index: "
            + ", ".join(sorted(absent_validation))
        )
    leakage_ids = sorted(
        record.id
        for record in selected
        if record.id not in validation_ids
        and entries[record.id].get("recommendedSplit") == "train"
        and str(entries[record.id].get("evaluationGroupId", ""))
        in validation_groups
    )
    if leakage_ids:
        raise ValueError(
            "Training selection leaks a validation evaluation group through "
            "song IDs: "
            + ", ".join(leakage_ids)
        )
    if len(selected) < 2:
        raise ValueError(
            f"Training index purpose {purpose!r} selected fewer than two songs."
        )
    return selected


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def verify_mask_audio_fingerprint(
    record: SongRecord,
    mask: AudioAlignedFrameMask,
    file_hashes: dict[Path, str],
) -> None:
    """Bind an indexed mask to materialized audio bytes for this record."""

    audio_path = record.audio_path
    if not audio_path.is_file():
        raise ValueError(
            f"Masked training song {record.id!r} audio file is missing: "
            f"{audio_path}"
        )
    resolved_path = audio_path.resolve()
    with resolved_path.open("rb") as stream:
        if stream.read(len(LFS_POINTER_PREFIX)) == LFS_POINTER_PREFIX:
            raise ValueError(
                f"Masked training song {record.id!r} audio is still a Git "
                f"LFS pointer: {audio_path}"
            )
    actual_fingerprint = file_hashes.get(resolved_path)
    if actual_fingerprint is None:
        actual_fingerprint = sha256_file(resolved_path)
        file_hashes[resolved_path] = actual_fingerprint
    if actual_fingerprint != mask.audio_fingerprint:
        raise ValueError(
            f"Masked training song {record.id!r} audioFingerprint does not "
            f"match the materialized audio file: {audio_path}"
        )


def mask_content_aggregate_sha256(
    masks: dict[str, AudioAlignedFrameMask],
) -> str:
    """Hash canonical selected mask coordinates and verified audio identity."""

    payload = [
        {
            "songId": song_id,
            "frameCount": mask.declared_frame_count,
            "ranges": [list(frame_range) for frame_range in mask.ranges],
            "audioFingerprint": mask.audio_fingerprint,
        }
        for song_id, mask in sorted(masks.items())
    ]
    canonical = json.dumps(
        payload,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    return hashlib.sha256(canonical).hexdigest()


def load_training_index_entries(
    index_path: Path,
) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]:
    """Load an index without silently accepting malformed or duplicate rows."""

    payload = json.loads(index_path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("Training index root must be a JSON object.")
    raw_entries = payload.get("entries")
    if not isinstance(raw_entries, list):
        raise ValueError("Training index entries must be a JSON array.")

    entries: dict[str, dict[str, Any]] = {}
    for entry_index, entry in enumerate(raw_entries):
        if not isinstance(entry, dict):
            raise ValueError(
                f"Training index entry {entry_index} must be a JSON object."
            )
        song_id_value = entry.get("songId")
        if not isinstance(song_id_value, str) or not song_id_value.strip():
            raise ValueError(
                f"Training index entry {entry_index} has an invalid songId."
            )
        song_id = song_id_value.strip()
        if song_id in entries:
            raise ValueError(f"Training index contains duplicate songId {song_id!r}.")
        entries[song_id] = entry
    return payload, entries


def parse_audio_aligned_frame_mask(
    entry: dict[str, Any],
    *,
    song_id: str,
) -> AudioAlignedFrameMask | None:
    """Parse one strict canonical half-open mask, or return legacy ``None``."""

    if "audioAlignedFrameMaskRanges" not in entry:
        partial_fields = [
            field
            for field in (
                "audioAlignedEndFrame",
                "audioAlignedFrameCount",
            )
            if field in entry
        ]
        if partial_fields:
            raise ValueError(
                f"Training index song {song_id!r} has partial audio frame "
                "mask metadata without audioAlignedFrameMaskRanges: "
                + ", ".join(partial_fields)
            )
        return None

    frame_count_value = entry.get("frameCount")
    if type(frame_count_value) is not int or frame_count_value < 0:
        raise ValueError(
            f"Training index song {song_id!r} has an audio frame mask but "
            "frameCount is not a non-negative integer."
        )
    frame_count = frame_count_value
    parsed_ranges = runtime_data.normalize_audio_aligned_frame_mask_ranges(
        entry["audioAlignedFrameMaskRanges"],
        frame_count,
        label=f"Training index song {song_id!r} audioAlignedFrameMaskRanges",
    )
    if frame_count > 0 and not parsed_ranges:
        raise ValueError(
            f"Training index song {song_id!r} "
            "audioAlignedFrameMaskRanges must admit at least one frame."
        )

    aligned_end_value = entry.get("audioAlignedEndFrame")
    if aligned_end_value is not None:
        if type(aligned_end_value) is not int:
            raise ValueError(
                f"Training index song {song_id!r} audioAlignedEndFrame "
                "must be an integer when present."
            )
        expected_aligned_end = parsed_ranges[-1][1] if parsed_ranges else 0
        if aligned_end_value != expected_aligned_end:
            raise ValueError(
                f"Training index song {song_id!r} audioAlignedEndFrame does "
                "not match the final half-open mask boundary."
            )

    aligned_count_value = entry.get("audioAlignedFrameCount")
    parsed_frame_count = sum(end - start for start, end in parsed_ranges)
    if aligned_count_value is not None:
        if type(aligned_count_value) is not int:
            raise ValueError(
                f"Training index song {song_id!r} audioAlignedFrameCount "
                "must be an integer when present."
            )
        if aligned_count_value != parsed_frame_count:
            raise ValueError(
                f"Training index song {song_id!r} audioAlignedFrameCount does "
                "not match audioAlignedFrameMaskRanges."
            )

    audio_fingerprint_value = entry.get("audioFingerprint")
    if not isinstance(audio_fingerprint_value, str):
        raise ValueError(
            f"Training index song {song_id!r} with an audio frame mask must "
            "declare audioFingerprint as a SHA-256 value."
        )
    if re.fullmatch(
        r"[0-9a-f]{64}",
        audio_fingerprint_value,
    ) is None:
        raise ValueError(
            f"Training index song {song_id!r} audioFingerprint must be a "
            "SHA-256 value when present."
        )
    return AudioAlignedFrameMask(
        song_id=song_id,
        declared_frame_count=frame_count,
        ranges=parsed_ranges,
        audio_aligned_end_frame=aligned_end_value,
        audio_fingerprint=audio_fingerprint_value,
    )


def load_audio_aligned_frame_masks(
    index_path: Path,
    records: Iterable[SongRecord],
    purpose: str,
) -> tuple[dict[str, AudioAlignedFrameMask], dict[str, Any]]:
    """Load selected training masks and their immutable source provenance."""

    payload, entries = load_training_index_entries(index_path)
    if purpose not in FRAME_MASK_TRAINING_PURPOSES:
        raise ValueError(
            f"Training purpose {purpose!r} has no audio frame-mask consumer."
        )
    selected_records = list(records)
    selected_ids = [record.id for record in selected_records]
    duplicate_selected_ids = sorted(
        song_id for song_id, count in Counter(selected_ids).items() if count > 1
    )
    if duplicate_selected_ids:
        raise ValueError(
            "Selected training records contain duplicate song IDs: "
            + ", ".join(duplicate_selected_ids)
        )

    parsed_masks: dict[str, AudioAlignedFrameMask] = {}
    for song_id, entry in entries.items():
        if entry.get("usage", {}).get(purpose) is not True:
            continue
        mask = parse_audio_aligned_frame_mask(entry, song_id=song_id)
        if mask is not None:
            parsed_masks[song_id] = mask

    masks: dict[str, AudioAlignedFrameMask] = {}
    audio_hashes: dict[Path, str] = {}
    for record in selected_records:
        song_id = record.id
        entry = entries.get(song_id)
        if entry is None:
            raise ValueError(
                f"Training index is missing selected song ID {song_id}."
            )
        if entry.get("usage", {}).get(purpose) is not True:
            raise ValueError(
                f"Selected song {song_id} is not eligible for {purpose}."
            )
        mask = parsed_masks.get(song_id)
        if mask is not None:
            verify_mask_audio_fingerprint(record, mask, audio_hashes)
            masks[song_id] = mask

    quality_gate = payload.get("qualityGate")
    if not isinstance(quality_gate, dict):
        quality_gate = {}
    provenance = {
        "contractVersion": "audio-aligned-frame-mask-v1",
        "rangeSemantics": "zero_based_half_open",
        "windowSelectionSemantics": "range_anchored_no_boundary_crossing",
        "cutBoundaryCensoring": (
            "zero_partial_mask_range_start_labels_with_full_range_legacy_parity"
        ),
        "windowCountReconciliation": (
            "legacy=eligibleLegacy+excludedLegacy; "
            "used=reusedLegacy+addedRangeAnchored; "
            "eligibleLegacy=reusedLegacy+omittedEligibleLegacy"
        ),
        "trainingPurpose": purpose,
        "consumerEnabled": True,
        "audioFingerprintVerification": "materialized_file_sha256_exact",
        "verifiedAudioFingerprintCount": len(masks),
        "maskContentAggregateSha256": mask_content_aggregate_sha256(masks),
        "trainingIndex": str(index_path.resolve()),
        "trainingIndexSha256": sha256_file(index_path),
        "candidateIndexSchemaVersion": payload.get("candidateIndexSchemaVersion"),
        "candidatePolicyVersion": payload.get(
            "candidatePolicyVersion",
            quality_gate.get("candidatePolicyVersion"),
        ),
        "catalogSha256": payload.get(
            "catalogSha256",
            quality_gate.get("catalogSha256"),
        ),
        "qualityReportSha256": quality_gate.get("qualityReportSha256"),
        "audioMaskConsumptionPolicy": quality_gate.get("audioMaskConsumptionPolicy"),
        "selectedSongCount": len(selected_ids),
        "maskedSongCount": len(masks),
        "legacyUnmaskedSongCount": len(selected_ids) - len(masks),
    }
    return masks, provenance


def normalized_catalog_folder(value: object) -> str:
    if not isinstance(value, str):
        return ""
    return value.replace("\\", "/").strip("/")


def record_folder_name(record: SongRecord) -> str:
    value = record.song.get("folderName")
    if not isinstance(value, str):
        return ""
    normalized = value.replace("\\", "/")
    if (
        not normalized
        or normalized.startswith("/")
        or re.match(r"^[A-Za-z]:/", normalized)
        or any(part in {"", ".", ".."} for part in normalized.split("/"))
    ):
        return ""
    return normalized


def catalog_safety_exclusion(catalog_song: dict[str, Any]) -> str | None:
    join_status = catalog_song.get("subjectTargetingJoinStatus")
    if isinstance(join_status, str) and join_status.startswith("conflict_"):
        return "subject_label_match_conflict"
    if not isinstance(join_status, str) or not join_status.startswith("matched_"):
        return "subject_label_not_safely_matched"
    if catalog_song.get("subjectTargeting") != "single_subject":
        return "catalog_subject_mode_not_single"
    if catalog_song.get("cameraTrackMuted") is not False:
        return "catalog_cameraTrackMuted_not_false"
    evidence = catalog_song.get("subjectTargetingEvidence")
    if not isinstance(evidence, dict):
        return "catalog_subject_evidence_missing"

    def evidence_is_explicitly_true(field: str) -> bool:
        values = evidence.get(f"{field}Values")
        if isinstance(values, list):
            return bool(values) and all(value is True for value in values)
        return evidence.get(field) is True

    for field in (
        "selectedAnimatorTargetMatchKnown",
        "allTargetsMatchSelectedExporterAnimator",
        "exportedCharacterPathMatchKnown",
        "exportedCharacterPathMatchesSelectedAnimator",
    ):
        if not evidence_is_explicitly_true(field):
            return f"catalog_subject_evidence_{field}_not_true"
    for field in SAFE_CATALOG_BOOLEAN_FIELDS:
        if catalog_song.get(field) is not True:
            return f"catalog_{field}_not_true"
    return None


def catalog_identity_exclusion(
    record: SongRecord,
    catalog_dataset: dict[str, Any],
    catalog_song: dict[str, Any],
) -> str | None:
    """Reject stale or ambiguous revision metadata before trusting labels."""

    manifest_fingerprint = record.manifest.get("sourceFingerprint")
    if catalog_dataset.get("isRevision") is not record.is_revision:
        return "catalog_dataset_revision_mismatch"
    if (
        catalog_dataset.get("revisionOfDatasetId")
        != record.revision_of_dataset_id
    ):
        return "catalog_dataset_revision_parent_mismatch"
    if catalog_dataset.get("sourceFingerprint") != manifest_fingerprint:
        return "stale_catalog_source_fingerprint"
    if record.is_revision:
        suffix = record.dataset_id.rsplit("_rev_", 1)[-1]
        if (
            not isinstance(manifest_fingerprint, str)
            or not manifest_fingerprint.casefold().startswith(suffix)
        ):
            return "revision_source_fingerprint_invalid"

    expected_revision_song_id = None
    if record.is_revision:
        scene_path = str(record.manifest.get("scenePath", ""))
        folder_name = str(record.song.get("folderName", ""))
        base_value = f"{scene_path}\0{record.name}\0{folder_name}".encode(
            "utf-8"
        )
        expected_revision_song_id = hashlib.sha256(base_value).hexdigest()[:16]
    expected_song_fields = {
        "id": record.id,
        "datasetId": record.dataset_id,
        "sourceFingerprint": manifest_fingerprint,
        "isRevision": record.is_revision,
        "revisionOfDatasetId": record.revision_of_dataset_id,
        "revisionOfSongId": expected_revision_song_id,
    }
    for field, expected in expected_song_fields.items():
        if catalog_song.get(field) != expected:
            return f"catalog_song_{field}_mismatch"
    return None


def training_input_paths(
    record: SongRecord,
    rebuild_cache: bool,
) -> tuple[str, list[tuple[str, Path]]] | None:
    if record.portable_prepared_path.is_file() and not rebuild_cache:
        return "prepared", [("preparedFile", record.portable_prepared_path)]

    paths: list[tuple[str, Path]] = []
    for field in RAW_TRAINING_FILE_FIELDS:
        if field == "audioFile":
            try:
                path = record.audio_path
            except (KeyError, TypeError):
                return None
        else:
            file_name = record.song.get(field)
            if not isinstance(file_name, str) or not file_name:
                return None
            path = record.song_dir / file_name
        paths.append((field, path))
    return "raw", paths


def validate_training_input_hashes(
    record: SongRecord,
    catalog_song: dict[str, Any],
    rebuild_cache: bool,
    file_hashes: dict[Path, str],
) -> tuple[str | None, int, str]:
    selected_inputs = training_input_paths(record, rebuild_cache)
    if selected_inputs is None:
        return "training_input_path_unavailable", 0, "unknown"
    input_mode, inputs = selected_inputs
    catalog_files = catalog_song.get("files")
    if not isinstance(catalog_files, dict):
        return "catalog_input_files_missing", 0, input_mode

    data_root = record.dataset_dir.parent.resolve()
    verified_count = 0
    for field, path in inputs:
        resolved_path = path.resolve()
        if not resolved_path.is_file():
            return "training_input_file_missing", verified_count, input_mode
        try:
            relative_path = resolved_path.relative_to(data_root).as_posix()
        except ValueError:
            return "training_input_outside_data_root", verified_count, input_mode

        file_info = catalog_files.get(field)
        if not isinstance(file_info, dict):
            return "catalog_input_file_missing", verified_count, input_mode
        if file_info.get("exists") is not True:
            return "catalog_input_file_not_available", verified_count, input_mode
        catalog_path = normalized_catalog_folder(file_info.get("path"))
        if catalog_path != normalized_catalog_folder(relative_path):
            return "catalog_input_path_mismatch", verified_count, input_mode
        expected_bytes = file_info.get("bytes")
        if (
            not isinstance(expected_bytes, int)
            or isinstance(expected_bytes, bool)
            or resolved_path.stat().st_size != expected_bytes
        ):
            return "stale_catalog_input_size", verified_count, input_mode
        expected_hash = file_info.get("sha256")
        if not isinstance(expected_hash, str) or re.fullmatch(
            r"[0-9a-fA-F]{64}",
            expected_hash,
        ) is None:
            return "catalog_input_hash_missing", verified_count, input_mode
        actual_hash = file_hashes.get(resolved_path)
        if actual_hash is None:
            actual_hash = sha256_file(resolved_path)
            file_hashes[resolved_path] = actual_hash
        if actual_hash.casefold() != expected_hash.casefold():
            return "stale_catalog_input_hash", verified_count, input_mode
        verified_count += 1
    return None, verified_count, input_mode


def filter_single_subject_records(
    records: list[SongRecord],
    catalog_path: Path,
    rebuild_cache: bool = False,
) -> tuple[list[SongRecord], dict[str, Any]]:
    """Select only fresh, unambiguous single-subject training records."""

    catalog_path = catalog_path.resolve()
    try:
        catalog = json.loads(catalog_path.read_text(encoding="utf-8-sig"))
    except (OSError, UnicodeError, json.JSONDecodeError) as error:
        raise ValueError(
            f"Unable to read catalog '{catalog_path}': {error}"
        ) from error
    if not isinstance(catalog, dict):
        raise ValueError(f"Catalog root must be an object: {catalog_path}")
    datasets = catalog.get("datasets")
    songs = catalog.get("songs")
    if not isinstance(datasets, list) or not isinstance(songs, list):
        raise ValueError("Catalog must contain 'datasets' and 'songs' arrays.")

    datasets_by_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
    invalid_dataset_rows = 0
    for dataset in datasets:
        if not isinstance(dataset, dict) or not isinstance(
            dataset.get("id"),
            str,
        ):
            invalid_dataset_rows += 1
            continue
        datasets_by_id[dataset["id"]].append(dataset)

    songs_by_folder: dict[str, list[dict[str, Any]]] = defaultdict(list)
    invalid_song_rows = 0
    for song in songs:
        if not isinstance(song, dict):
            invalid_song_rows += 1
            continue
        folder = normalized_catalog_folder(song.get("folder"))
        if not folder:
            invalid_song_rows += 1
            continue
        songs_by_folder[folder].append(song)

    selected: list[SongRecord] = []
    exclusions: Counter[str] = Counter()
    manifest_hashes: dict[Path, str] = {}
    file_hashes: dict[Path, str] = {}
    input_mode_counts: Counter[str] = Counter()
    verified_input_files = 0

    dataset_paths: dict[str, set[Path]] = defaultdict(set)
    logical_song_counts: Counter[tuple[str, str]] = Counter()
    resolved_song_path_counts: Counter[Path] = Counter()
    for record in records:
        dataset_id = record.dataset_id
        folder_name = record_folder_name(record)
        dataset_paths[dataset_id].add(record.manifest_path.resolve())
        logical_song_counts[(dataset_id, folder_name)] += 1
        resolved_song_path_counts[record.song_dir.resolve()] += 1

    for record in records:
        dataset_id = record.dataset_id
        folder_name = record_folder_name(record)
        if not folder_name:
            exclusions["invalid_input_song_folder"] += 1
            continue
        resolved_dataset_path = record.dataset_dir.resolve()
        resolved_song_path = record.song_dir.resolve()
        try:
            resolved_song_path.relative_to(resolved_dataset_path)
        except ValueError:
            exclusions["input_song_path_outside_dataset"] += 1
            continue
        if len(dataset_paths[dataset_id]) != 1:
            exclusions["input_dataset_path_conflict"] += 1
            continue
        if logical_song_counts[(dataset_id, folder_name)] != 1:
            exclusions["input_logical_song_conflict"] += 1
            continue
        if resolved_song_path_counts[resolved_song_path] != 1:
            exclusions["input_resolved_song_path_conflict"] += 1
            continue

        catalog_datasets = datasets_by_id.get(dataset_id, [])
        if not catalog_datasets:
            exclusions["missing_catalog_dataset"] += 1
            continue
        if len(catalog_datasets) != 1:
            exclusions["catalog_dataset_match_conflict"] += 1
            continue
        catalog_dataset = catalog_datasets[0]
        expected_manifest_hash = catalog_dataset.get("manifestSha256")
        if not isinstance(expected_manifest_hash, str) or not expected_manifest_hash:
            exclusions["catalog_manifest_hash_missing"] += 1
            continue
        manifest_path = record.manifest_path.resolve()
        manifest_hash = manifest_hashes.get(manifest_path)
        if manifest_hash is None:
            manifest_hash = sha256_file(manifest_path)
            manifest_hashes[manifest_path] = manifest_hash
        if manifest_hash.casefold() != expected_manifest_hash.casefold():
            exclusions["stale_catalog_manifest"] += 1
            continue

        folder = normalized_catalog_folder(f"{dataset_id}/{folder_name}")
        catalog_songs = songs_by_folder.get(folder, [])
        if not catalog_songs:
            exclusions["missing_catalog_song"] += 1
            continue
        if len(catalog_songs) != 1:
            exclusions["catalog_song_match_conflict"] += 1
            continue
        catalog_song = catalog_songs[0]
        identity_exclusion = catalog_identity_exclusion(
            record,
            catalog_dataset,
            catalog_song,
        )
        if identity_exclusion is not None:
            exclusions[identity_exclusion] += 1
            continue
        safety_exclusion = catalog_safety_exclusion(catalog_song)
        if safety_exclusion is not None:
            exclusions[safety_exclusion] += 1
            continue
        input_exclusion, verified_count, input_mode = (
            validate_training_input_hashes(
                record,
                catalog_song,
                rebuild_cache,
                file_hashes,
            )
        )
        if input_exclusion is not None:
            exclusions[input_exclusion] += 1
            continue
        selected.append(record)
        verified_input_files += verified_count
        input_mode_counts[input_mode] += 1

    summary: dict[str, Any] = {
        "catalog": str(catalog_path),
        "catalogSha256": sha256_file(catalog_path),
        "candidateSongs": len(records),
        "selectedSongs": len(selected),
        "excludedSongs": len(records) - len(selected),
        "exclusions": dict(sorted(exclusions.items())),
        "invalidCatalogDatasetRows": invalid_dataset_rows,
        "invalidCatalogSongRows": invalid_song_rows,
        "verifiedInputFiles": verified_input_files,
        "inputModeCounts": dict(sorted(input_mode_counts.items())),
    }
    print(
        "Single-subject catalog selection: "
        f"{len(selected)}/{len(records)} songs included from {catalog_path}."
    )
    for reason, count in sorted(exclusions.items()):
        print(f"  excluded {count}: {reason}")
    if invalid_dataset_rows or invalid_song_rows:
        print(
            "  ignored invalid catalog rows: "
            f"datasets={invalid_dataset_rows}, songs={invalid_song_rows}"
        )
    return selected, summary


def normalized_leaf(path: str) -> str:
    leaf = path.replace("\\", "/").rsplit("/", 1)[-1].lower()
    leaf = leaf.replace("bip001", "")
    return re.sub(r"[^a-z0-9]", "", leaf)


def canonical_joint_indices(names: list[str]) -> list[int]:
    normalized = [normalized_leaf(name) for name in names]
    result: list[int] = []
    missing: list[str] = []
    for canonical_name, aliases in CANONICAL_JOINT_ALIASES.items():
        match = next(
            (index for index, value in enumerate(normalized) if value in aliases),
            None,
        )
        if match is None:
            missing.append(canonical_name)
        else:
            result.append(match)
    if missing:
        raise ValueError(f"Missing canonical joints: {', '.join(missing)}")
    return result


def quat_conjugate(q: np.ndarray) -> np.ndarray:
    result = q.copy()
    result[..., :3] *= -1.0
    return result


def quat_multiply(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    av, aw = a[..., :3], a[..., 3:4]
    bv, bw = b[..., :3], b[..., 3:4]
    xyz = aw * bv + bw * av + np.cross(av, bv)
    w = aw * bw - np.sum(av * bv, axis=-1, keepdims=True)
    return np.concatenate((xyz, w), axis=-1)


def quat_rotate(q: np.ndarray, v: np.ndarray) -> np.ndarray:
    qv = q[..., :3]
    qw = q[..., 3:4]
    return v + 2.0 * (
        qw * np.cross(qv, v) + np.cross(qv, np.cross(qv, v))
    )


def quat_to_rotation_6d(q: np.ndarray) -> np.ndarray:
    q = q / np.maximum(np.linalg.norm(q, axis=-1, keepdims=True), 1e-8)
    x, y, z, w = (q[..., index] for index in range(4))
    first_column = np.stack(
        (
            1.0 - 2.0 * (y * y + z * z),
            2.0 * (x * y + z * w),
            2.0 * (x * z - y * w),
        ),
        axis=-1,
    )
    second_column = np.stack(
        (
            2.0 * (x * y - z * w),
            1.0 - 2.0 * (x * x + z * z),
            2.0 * (y * z + x * w),
        ),
        axis=-1,
    )
    return np.concatenate((first_column, second_column), axis=-1)


def extract_audio_features(
    path: Path,
    timeline_times: np.ndarray,
    audio_start: float,
) -> np.ndarray:
    waveform, _ = librosa.load(path, sr=AUDIO_RATE, mono=True)
    mel = librosa.feature.melspectrogram(
        y=waveform,
        sr=AUDIO_RATE,
        n_fft=1024,
        hop_length=AUDIO_HOP,
        n_mels=MEL_BANDS,
        fmin=30,
        fmax=AUDIO_RATE // 2,
        center=True,
        power=2.0,
    )
    mel = np.log1p(mel).T
    rms = librosa.feature.rms(
        y=waveform,
        frame_length=1024,
        hop_length=AUDIO_HOP,
        center=True,
    ).T
    onset = librosa.onset.onset_strength(
        y=waveform,
        sr=AUDIO_RATE,
        hop_length=AUDIO_HOP,
        center=True,
    )[:, None]
    feature_frames = min(len(mel), len(rms), len(onset))
    audio_features = np.concatenate(
        (mel[:feature_frames], rms[:feature_frames], onset[:feature_frames]),
        axis=1,
    ).astype(np.float32)
    feature_times = np.arange(feature_frames, dtype=np.float64) / 60.0
    query_times = timeline_times - float(audio_start)
    aligned = np.zeros((len(timeline_times), audio_features.shape[1]), dtype=np.float32)
    valid = (query_times >= 0.0) & (query_times <= feature_times[-1])
    for feature_index in range(audio_features.shape[1]):
        aligned[valid, feature_index] = np.interp(
            query_times[valid],
            feature_times,
            audio_features[:, feature_index],
        )
    return aligned


def cache_path_for(record: SongRecord, cache_dir: Path) -> Path:
    source_paths = [
        record.manifest_path,
        record.audio_path,
        record.song_dir / "joints_world.f32",
        record.song_dir / "camera.f32",
        record.song_dir / "root.f32",
        record.song_dir / "shot_index.i32",
        record.song_dir / "time.f64",
    ]
    signature = {
        "pipelineVersion": PIPELINE_VERSION,
        "song": record.name,
        "sources": [
            (str(path), path.stat().st_size, path.stat().st_mtime_ns)
            for path in source_paths
        ],
    }
    digest = hashlib.sha1(
        json.dumps(signature, ensure_ascii=False).encode("utf-8")
    ).hexdigest()[:12]
    safe_name = re.sub(r"[^0-9A-Za-z가-힣_-]+", "_", record.name).strip("_")
    return cache_dir / f"{safe_name}_{digest}.npz"


def generation_input_cache_path_for(
    record: SongRecord,
    cache_dir: Path,
) -> Path:
    """Cache identity for a target that never includes authored camera data."""

    return runtime_data.generation_input_cache_path_for(
        record,
        cache_dir,
        pipeline_version=PIPELINE_VERSION,
    )


def prepare_song_inputs(
    record: SongRecord,
) -> tuple[np.ndarray, np.ndarray, float, np.ndarray, np.ndarray]:
    """Prepare motion/audio inputs shared by training and camera-free targets."""

    frames = int(record.song["frameCount"])
    joint_count = int(record.song["sharedBoneCount"])
    source_joints = np.fromfile(
        record.song_dir / "joints_world.f32", dtype="<f4"
    ).reshape(frames, joint_count, 3)
    root = np.fromfile(
        record.song_dir / "root.f32", dtype="<f4"
    ).reshape(frames, 14)
    times = np.fromfile(record.song_dir / "time.f64", dtype="<f8")
    if len(times) != frames:
        raise ValueError(
            f"Timeline time count does not match frameCount for {record.name}: "
            f"{len(times)} != {frames}"
        )

    indices = canonical_joint_indices(record.manifest["sharedBones"])
    joints = source_joints[:, indices]
    root_position = root[:, :3]
    root_rotation = root[:, 3:7]
    inverse_root_rotation = quat_conjugate(root_rotation)

    joints_local = quat_rotate(
        inverse_root_rotation[:, None, :],
        joints - root_position[:, None, :],
    )
    pelvis_index = CANONICAL_JOINTS.index("pelvis")
    head_index = CANONICAL_JOINTS.index("head")
    body_scale = float(
        np.median(
            np.linalg.norm(
                joints_local[:, head_index] - joints_local[:, pelvis_index],
                axis=-1,
            )
        )
    )
    if not math.isfinite(body_scale) or body_scale < 1e-4:
        raise ValueError(f"Invalid body scale for {record.name}: {body_scale}")

    joints_local /= body_scale
    joint_velocity = np.zeros_like(joints_local)
    joint_velocity[1:] = np.diff(joints_local, axis=0) * 60.0

    audio = extract_audio_features(
        record.audio_path,
        times,
        float(record.song["audioStart"]),
    )
    inputs = np.concatenate(
        (
            joints_local.reshape(frames, -1),
            joint_velocity.reshape(frames, -1),
            audio,
        ),
        axis=1,
    ).astype(np.float32)
    if not (np.isfinite(inputs).all() and np.isfinite(times).all()):
        raise ValueError(f"Non-finite input data found in {record.name}.")
    return inputs, times, body_scale, root_position, inverse_root_rotation


def prepare_song(record: SongRecord) -> PreparedSong:
    frames = int(record.song["frameCount"])
    (
        inputs,
        times,
        body_scale,
        root_position,
        inverse_root_rotation,
    ) = prepare_song_inputs(record)
    camera = np.fromfile(record.song_dir / "camera.f32", dtype="<f4").reshape(frames, 9)
    shot_index = np.fromfile(record.song_dir / "shot_index.i32", dtype="<i4")
    if len(shot_index) != frames:
        raise ValueError(
            f"Shot index count does not match frameCount for {record.name}: "
            f"{len(shot_index)} != {frames}"
        )

    camera_local_position = quat_rotate(
        inverse_root_rotation,
        camera[:, :3] - root_position,
    ) / body_scale
    camera_relative_rotation = quat_multiply(
        inverse_root_rotation,
        camera[:, 3:7],
    )
    camera_rotation_6d = quat_to_rotation_6d(camera_relative_rotation)
    targets = np.concatenate(
        (
            camera_local_position,
            camera_rotation_6d,
            camera[:, 7:9],
        ),
        axis=1,
    ).astype(np.float32)

    cuts = np.zeros(frames, dtype=np.float32)
    cuts[1:] = (shot_index[1:] != shot_index[:-1]).astype(np.float32)

    if not (
        np.isfinite(inputs).all()
        and np.isfinite(targets).all()
        and np.isfinite(times).all()
    ):
        raise ValueError(f"Non-finite training data found in {record.name}.")

    return PreparedSong(
        name=record.name,
        inputs=inputs,
        targets=targets,
        cuts=cuts,
        times=times,
        body_scale=body_scale,
        source_manifest=str(record.manifest_path),
    )


def prepare_generation_input_song(record: SongRecord) -> PreparedSong:
    """Prepare a new motion/audio target without touching camera artifacts."""

    if not record.is_generation_input:
        raise ValueError(
            "Camera-free preparation requires inputKind='generation_input' "
            "and hasAuthoredCamera=false."
        )
    inputs, times, body_scale, _, _ = prepare_song_inputs(record)
    frames = len(inputs)
    return PreparedSong(
        name=record.name,
        inputs=inputs,
        # These arrays satisfy the common PreparedSong shape contract only.
        # Hybrid generation never reads target camera labels or authored cuts.
        targets=np.zeros((frames, 11), dtype=np.float32),
        cuts=np.zeros(frames, dtype=np.float32),
        times=times,
        body_scale=body_scale,
        source_manifest=str(record.manifest_path),
    )


def load_or_prepare_generation_input_song(
    record: SongRecord,
    cache_dir: Path,
    rebuild: bool,
) -> PreparedSong:
    """Use the worker's atomic, validated generation-input cache contract."""

    prepared = runtime_data.load_or_prepare_generation_input_song(
        record,
        cache_dir,
        rebuild,
        pipeline_version=PIPELINE_VERSION,
        prepare_song_fn=prepare_generation_input_song,
    )
    return PreparedSong(
        name=prepared.name,
        inputs=prepared.inputs,
        targets=prepared.targets,
        cuts=prepared.cuts,
        times=prepared.times,
        body_scale=prepared.body_scale,
        source_manifest=prepared.source_manifest,
    )


def load_or_prepare_song(
    record: SongRecord,
    cache_dir: Path,
    rebuild: bool,
    allow_cached_raw: bool = True,
) -> PreparedSong:
    if record.portable_prepared_path.is_file() and not rebuild:
        with np.load(record.portable_prepared_path, allow_pickle=False) as data:
            return PreparedSong(
                name=str(data["name"].item()),
                inputs=data["inputs"],
                targets=data["targets"],
                cuts=data["cuts"],
                times=data["times"],
                body_scale=float(data["body_scale"].item()),
                source_manifest=str(data["source_manifest"].item()),
            )

    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_path = cache_path_for(record, cache_dir)
    if cache_path.is_file() and not rebuild and allow_cached_raw:
        with np.load(cache_path, allow_pickle=False) as data:
            return PreparedSong(
                name=str(data["name"].item()),
                inputs=data["inputs"],
                targets=data["targets"],
                cuts=data["cuts"],
                times=data["times"],
                body_scale=float(data["body_scale"].item()),
                source_manifest=str(data["source_manifest"].item()),
            )

    prepared = prepare_song(record)
    np.savez(
        cache_path,
        name=np.asarray(prepared.name),
        inputs=prepared.inputs,
        targets=prepared.targets,
        cuts=prepared.cuts,
        times=prepared.times,
        body_scale=np.asarray(prepared.body_scale, dtype=np.float32),
        source_manifest=np.asarray(prepared.source_manifest),
    )
    return prepared


def resolve_prepared_frame_ranges(
    prepared_records: Sequence[tuple[SongRecord, PreparedSong]],
    masks: dict[str, AudioAlignedFrameMask],
) -> dict[str, tuple[tuple[int, int], ...]]:
    """Validate masks against prepared arrays without changing those arrays."""

    resolved: dict[str, tuple[tuple[int, int], ...]] = {}
    for record, song in prepared_records:
        frame_count = len(song.inputs)
        mask = masks.get(record.id)
        if mask is None:
            resolved[record.id] = ((0, frame_count),) if frame_count else ()
            continue

        array_lengths = {
            "inputs": frame_count,
            "targets": len(song.targets),
            "cuts": len(song.cuts),
            "times": len(song.times),
        }
        if len(set(array_lengths.values())) != 1:
            raise ValueError(
                f"Prepared arrays for masked song {record.id!r} have "
                "different frame counts: " + json.dumps(array_lengths, sort_keys=True)
            )
        if frame_count != mask.declared_frame_count:
            raise ValueError(
                f"Training index frameCount for masked song {record.id!r} "
                f"is {mask.declared_frame_count}, but prepared arrays contain "
                f"{frame_count} frames."
            )
        resolved[record.id] = mask.ranges
    return resolved


def concatenate_prepared_frame_ranges(
    prepared_records: Sequence[tuple[SongRecord, PreparedSong]],
    ranges_by_song_id: dict[str, tuple[tuple[int, int], ...]],
    attribute: str,
    *,
    cut_range_start_censor_song_ids: set[str] | None = None,
) -> np.ndarray:
    """Concatenate only admitted source frames while preserving source arrays."""

    censor_song_ids = cut_range_start_censor_song_ids or set()
    if censor_song_ids and attribute != "cuts":
        raise ValueError(
            "Range-start cut censoring can only be applied to cut labels."
        )
    chunks: list[np.ndarray] = []
    for record, song in prepared_records:
        ranges = ranges_by_song_id[record.id]
        preserve_full_range = ranges == ((0, len(song.cuts)),)
        for start, end in ranges:
            chunk = getattr(song, attribute)[start:end]
            if record.id in censor_song_ids and not preserve_full_range:
                chunk = np.array(chunk, copy=True)
                chunk[0] = 0.0
            chunks.append(chunk)
    if not chunks:
        raise ValueError(f"Audio-aligned frame masks selected no {attribute} frames.")
    return np.concatenate(chunks)


class WindowDataset(Dataset):
    def __init__(
        self,
        songs: list[PreparedSong],
        window: int,
        stride: int,
        input_mean: np.ndarray,
        input_std: np.ndarray,
        target_mean: np.ndarray,
        target_std: np.ndarray,
        frame_ranges: Sequence[Sequence[tuple[int, int]]] | None = None,
        censor_cut_range_starts: Sequence[bool] | None = None,
    ) -> None:
        if frame_ranges is not None:
            if window <= 0:
                raise ValueError("Window size must be positive.")
            if stride <= 0:
                raise ValueError("Window stride must be positive.")
            if len(frame_ranges) != len(songs):
                raise ValueError("Frame ranges must contain one item per song.")
        if censor_cut_range_starts is not None:
            if len(censor_cut_range_starts) != len(songs):
                raise ValueError(
                    "Cut range-start censor flags must contain one item per song."
                )
            if frame_ranges is None and any(censor_cut_range_starts):
                raise ValueError(
                    "Cut range-start censoring requires explicit frame ranges."
                )
        self.songs = songs
        self.window = window
        self.input_mean = input_mean
        self.input_std = input_std
        self.target_mean = target_mean
        self.target_std = target_std
        self.windows: list[tuple[int, int]] = []
        self.song_window_counts: list[int] = []
        self.song_unmasked_window_counts: list[int] = []
        self.song_eligible_legacy_window_counts: list[int] = []
        self.song_excluded_legacy_window_counts: list[int] = []
        self.song_reused_legacy_window_counts: list[int] = []
        self.song_added_range_anchored_window_counts: list[int] = []
        self.song_omitted_eligible_legacy_window_counts: list[int] = []
        self.song_included_frame_counts: list[int] = []
        self.song_excluded_frame_counts: list[int] = []
        self.song_censored_cut_range_start_counts: list[int] = []
        self.song_positive_censored_cut_range_start_counts: list[int] = []
        self.song_censored_cut_range_starts: list[tuple[int, ...]] = []
        for song_index, song in enumerate(songs):
            frame_count = len(song.inputs)
            unmasked_starts = list(range(0, frame_count - window + 1, stride))
            if frame_ranges is None:
                ranges = ((0, frame_count),) if frame_count else ()
            else:
                raw_ranges = [
                    list(frame_range)
                    if isinstance(frame_range, tuple)
                    else frame_range
                    for frame_range in frame_ranges[song_index]
                ]
                ranges = runtime_data.normalize_audio_aligned_frame_mask_ranges(
                    raw_ranges,
                    frame_count,
                    label=f"WindowDataset song {song_index} frame ranges",
                )
            censor_range_starts = (
                frame_ranges is not None
                if censor_cut_range_starts is None
                else bool(censor_cut_range_starts[song_index])
            )
            censored_cut_range_starts = (
                tuple(start for start, _ in ranges)
                if censor_range_starts
                and ranges != ((0, frame_count),)
                else ()
            )
            positive_censored_count = sum(
                float(song.cuts[start]) >= 0.5
                for start in censored_cut_range_starts
            )

            selected_starts: list[int] = []
            for start, end in ranges:
                selected_starts.extend(
                    range(start, end - window + 1, stride)
                )
            self.windows.extend(
                (song_index, window_start)
                for window_start in selected_starts
            )
            eligible_legacy_starts = [
                start
                for start in unmasked_starts
                if runtime_data.audio_aligned_frame_mask_range_index(
                    start,
                    start + window,
                    ranges,
                )
                is not None
            ]
            selected_start_set = set(selected_starts)
            eligible_legacy_start_set = set(eligible_legacy_starts)
            reused_legacy_count = len(
                selected_start_set & eligible_legacy_start_set
            )
            included_frame_count = sum(end - start for start, end in ranges)
            self.song_window_counts.append(len(selected_starts))
            self.song_unmasked_window_counts.append(len(unmasked_starts))
            self.song_eligible_legacy_window_counts.append(
                len(eligible_legacy_starts)
            )
            self.song_excluded_legacy_window_counts.append(
                len(unmasked_starts) - len(eligible_legacy_starts)
            )
            self.song_reused_legacy_window_counts.append(reused_legacy_count)
            self.song_added_range_anchored_window_counts.append(
                len(selected_start_set - eligible_legacy_start_set)
            )
            self.song_omitted_eligible_legacy_window_counts.append(
                len(eligible_legacy_start_set - selected_start_set)
            )
            self.song_included_frame_counts.append(included_frame_count)
            self.song_excluded_frame_counts.append(frame_count - included_frame_count)
            self.song_censored_cut_range_start_counts.append(
                len(censored_cut_range_starts)
            )
            self.song_positive_censored_cut_range_start_counts.append(
                positive_censored_count
            )
            self.song_censored_cut_range_starts.append(
                censored_cut_range_starts
            )

    def __len__(self) -> int:
        return len(self.windows)

    def __getitem__(self, index: int):
        song_index, start = self.windows[index]
        song = self.songs[song_index]
        end = start + self.window
        inputs = (song.inputs[start:end] - self.input_mean) / self.input_std
        targets = (song.targets[start:end] - self.target_mean) / self.target_std
        cuts = song.cuts[start:end]
        censored_offsets = [
            frame - start
            for frame in self.song_censored_cut_range_starts[song_index]
            if start <= frame < end
        ]
        if censored_offsets:
            cuts = np.array(cuts, copy=True)
            cuts[censored_offsets] = 0.0
        return (
            torch.from_numpy(inputs),
            torch.from_numpy(targets),
            torch.from_numpy(cuts),
            torch.tensor(song.body_scale, dtype=torch.float32),
        )


def build_audio_frame_mask_report(
    provenance: dict[str, Any],
    masks: dict[str, AudioAlignedFrameMask],
    ranges_by_song_id: dict[str, tuple[tuple[int, int], ...]],
    training_records: Sequence[tuple[SongRecord, PreparedSong]],
    validation_records: Sequence[tuple[SongRecord, PreparedSong]],
    training_dataset: WindowDataset,
    validation_dataset: WindowDataset,
) -> dict[str, Any]:
    """Describe exact mask provenance and split-level consumption counts."""

    def split_report(
        prepared_records: Sequence[tuple[SongRecord, PreparedSong]],
        dataset: WindowDataset,
    ) -> dict[str, Any]:
        songs: list[dict[str, Any]] = []
        for song_index, (record, song) in enumerate(prepared_records):
            ranges = ranges_by_song_id[record.id]
            songs.append(
                {
                    "songId": record.id,
                    "songName": song.name,
                    "maskPresent": record.id in masks,
                    "ranges": [list(frame_range) for frame_range in ranges],
                    "sourceFrameCount": len(song.inputs),
                    "usedFrameCount": dataset.song_included_frame_counts[song_index],
                    "excludedFrameCount": dataset.song_excluded_frame_counts[
                        song_index
                    ],
                    "usedWindowCount": dataset.song_window_counts[song_index],
                    "legacyUnmaskedWindowCount": (
                        dataset.song_unmasked_window_counts[song_index]
                    ),
                    "eligibleLegacyWindowCount": (
                        dataset.song_eligible_legacy_window_counts[song_index]
                    ),
                    "excludedLegacyWindowCount": (
                        dataset.song_excluded_legacy_window_counts[song_index]
                    ),
                    "reusedLegacyWindowCount": (
                        dataset.song_reused_legacy_window_counts[song_index]
                    ),
                    "addedRangeAnchoredWindowCount": (
                        dataset.song_added_range_anchored_window_counts[song_index]
                    ),
                    "omittedEligibleLegacyWindowCount": (
                        dataset.song_omitted_eligible_legacy_window_counts[
                            song_index
                        ]
                    ),
                    "censoredCutRangeStartCount": (
                        dataset.song_censored_cut_range_start_counts[song_index]
                    ),
                    "positiveCensoredCutRangeStartCount": (
                        dataset.song_positive_censored_cut_range_start_counts[
                            song_index
                        ]
                    ),
                }
            )
        masked_songs = [row for row in songs if row["maskPresent"]]
        return {
            "songCount": len(songs),
            "maskedSongCount": sum(row["maskPresent"] for row in songs),
            "sourceFrameCount": sum(row["sourceFrameCount"] for row in songs),
            "usedFrameCount": sum(row["usedFrameCount"] for row in songs),
            "excludedFrameCount": sum(row["excludedFrameCount"] for row in songs),
            "usedWindowCount": sum(row["usedWindowCount"] for row in songs),
            "legacyUnmaskedWindowCount": sum(
                row["legacyUnmaskedWindowCount"] for row in songs
            ),
            "eligibleLegacyWindowCount": sum(
                row["eligibleLegacyWindowCount"] for row in songs
            ),
            "excludedLegacyWindowCount": sum(
                row["excludedLegacyWindowCount"] for row in songs
            ),
            "reusedLegacyWindowCount": sum(
                row["reusedLegacyWindowCount"] for row in songs
            ),
            "addedRangeAnchoredWindowCount": sum(
                row["addedRangeAnchoredWindowCount"] for row in songs
            ),
            "omittedEligibleLegacyWindowCount": sum(
                row["omittedEligibleLegacyWindowCount"] for row in songs
            ),
            "censoredCutRangeStartCount": sum(
                row["censoredCutRangeStartCount"] for row in songs
            ),
            "positiveCensoredCutRangeStartCount": sum(
                row["positiveCensoredCutRangeStartCount"] for row in songs
            ),
            "maskedSourceFrameCount": sum(
                row["sourceFrameCount"] for row in masked_songs
            ),
            "maskedUsedFrameCount": sum(
                row["usedFrameCount"] for row in masked_songs
            ),
            "maskedExcludedFrameCount": sum(
                row["excludedFrameCount"] for row in masked_songs
            ),
            "maskedUsedWindowCount": sum(
                row["usedWindowCount"] for row in masked_songs
            ),
            "maskedEligibleLegacyWindowCount": sum(
                row["eligibleLegacyWindowCount"] for row in masked_songs
            ),
            "maskedExcludedLegacyWindowCount": sum(
                row["excludedLegacyWindowCount"] for row in masked_songs
            ),
            "maskedReusedLegacyWindowCount": sum(
                row["reusedLegacyWindowCount"] for row in masked_songs
            ),
            "maskedAddedRangeAnchoredWindowCount": sum(
                row["addedRangeAnchoredWindowCount"] for row in masked_songs
            ),
            "maskedOmittedEligibleLegacyWindowCount": sum(
                row["omittedEligibleLegacyWindowCount"] for row in masked_songs
            ),
            "songs": songs,
        }

    training = split_report(training_records, training_dataset)
    validation = split_report(validation_records, validation_dataset)
    report = dict(provenance)
    report.update(
        {
            "applied": bool(masks),
            "maskedRecords": [
                {
                    "songId": song_id,
                    "declaredFrameCount": mask.declared_frame_count,
                    "ranges": [list(frame_range) for frame_range in mask.ranges],
                    "audioAlignedEndFrame": mask.audio_aligned_end_frame,
                    "audioFingerprint": mask.audio_fingerprint,
                    "audioFingerprintVerified": True,
                }
                for song_id, mask in sorted(masks.items())
            ],
            "training": training,
            "validation": validation,
            "totals": {
                key: training[key] + validation[key]
                for key in (
                    "songCount",
                    "maskedSongCount",
                    "sourceFrameCount",
                    "usedFrameCount",
                    "excludedFrameCount",
                    "usedWindowCount",
                    "legacyUnmaskedWindowCount",
                    "eligibleLegacyWindowCount",
                    "excludedLegacyWindowCount",
                    "reusedLegacyWindowCount",
                    "addedRangeAnchoredWindowCount",
                    "omittedEligibleLegacyWindowCount",
                    "censoredCutRangeStartCount",
                    "positiveCensoredCutRangeStartCount",
                    "maskedSourceFrameCount",
                    "maskedUsedFrameCount",
                    "maskedExcludedFrameCount",
                    "maskedUsedWindowCount",
                    "maskedEligibleLegacyWindowCount",
                    "maskedExcludedLegacyWindowCount",
                    "maskedReusedLegacyWindowCount",
                    "maskedAddedRangeAnchoredWindowCount",
                    "maskedOmittedEligibleLegacyWindowCount",
                )
            },
        }
    )
    return report


class ResidualTemporalBlock(nn.Module):
    def __init__(self, hidden_size: int, dilation: int, dropout: float) -> None:
        super().__init__()
        self.network = nn.Sequential(
            nn.Conv1d(
                hidden_size,
                hidden_size,
                kernel_size=3,
                padding=dilation,
                dilation=dilation,
            ),
            nn.GroupNorm(1, hidden_size),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Conv1d(
                hidden_size,
                hidden_size,
                kernel_size=3,
                padding=dilation,
                dilation=dilation,
            ),
            nn.GroupNorm(1, hidden_size),
            nn.GELU(),
            nn.Dropout(dropout),
        )

    def forward(self, values: torch.Tensor) -> torch.Tensor:
        return values + self.network(values)


class CameraDirector(nn.Module):
    def __init__(
        self,
        input_size: int,
        output_size: int,
        hidden_size: int,
        layers: int,
        dropout: float,
    ) -> None:
        super().__init__()
        self.input_projection = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.LayerNorm(hidden_size),
            nn.GELU(),
        )
        self.temporal = nn.Sequential(
            *(
                ResidualTemporalBlock(
                    hidden_size,
                    dilation=2**layer,
                    dropout=dropout,
                )
                for layer in range(layers)
            )
        )
        self.head = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Linear(hidden_size, hidden_size),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, output_size + 1),
        )

    def forward(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        encoded = self.input_projection(inputs)
        temporal = self.temporal(encoded.transpose(1, 2)).transpose(1, 2)
        output = self.head(temporal)
        return output[..., :-1], output[..., -1]


def grouped_loss(
    predicted: torch.Tensor,
    target: torch.Tensor,
    cut_logit: torch.Tensor,
    cut_target: torch.Tensor,
    cut_loss_fn: nn.Module,
) -> tuple[torch.Tensor, dict[str, float]]:
    position = torch.mean((predicted[..., :3] - target[..., :3]) ** 2)
    rotation = torch.mean((predicted[..., 3:9] - target[..., 3:9]) ** 2)
    lens = torch.mean((predicted[..., 9:11] - target[..., 9:11]) ** 2)
    cut = cut_loss_fn(cut_logit, cut_target)
    total = position + rotation + 0.5 * lens + 0.1 * cut
    return total, {
        "position": float(position.detach()),
        "rotation": float(rotation.detach()),
        "lens": float(lens.detach()),
        "cut": float(cut.detach()),
    }


def rotation_6d_to_matrix(values: np.ndarray) -> np.ndarray:
    first = values[..., :3]
    second = values[..., 3:6]
    first = first / np.maximum(np.linalg.norm(first, axis=-1, keepdims=True), 1e-8)
    second = second - np.sum(first * second, axis=-1, keepdims=True) * first
    second = second / np.maximum(np.linalg.norm(second, axis=-1, keepdims=True), 1e-8)
    third = np.cross(first, second)
    return np.stack((first, second, third), axis=-1)


def rotation_6d_to_quat(values: np.ndarray) -> np.ndarray:
    """Convert the compact camera rotation representation to XYZW quaternions."""

    matrix = rotation_6d_to_matrix(values)
    result = np.empty(matrix.shape[:-2] + (4,), dtype=np.float32)
    trace = np.trace(matrix, axis1=-2, axis2=-1)
    positive = trace > 0.0
    root = np.sqrt(np.maximum(trace[positive] + 1.0, 1e-8)) * 2.0
    result[positive, 3] = 0.25 * root
    result[positive, 0] = (matrix[positive, 2, 1] - matrix[positive, 1, 2]) / root
    result[positive, 1] = (matrix[positive, 0, 2] - matrix[positive, 2, 0]) / root
    result[positive, 2] = (matrix[positive, 1, 0] - matrix[positive, 0, 1]) / root
    remaining = ~positive
    for axis in range(3):
        mask = remaining & (matrix[..., axis, axis] >= matrix[..., (axis + 1) % 3, (axis + 1) % 3])
        mask &= matrix[..., axis, axis] >= matrix[..., (axis + 2) % 3, (axis + 2) % 3]
        if not np.any(mask):
            continue
        j = (axis + 1) % 3
        k = (axis + 2) % 3
        root = np.sqrt(np.maximum(1.0 + matrix[mask, axis, axis] - matrix[mask, j, j] - matrix[mask, k, k], 1e-8)) * 2.0
        result[mask, axis] = 0.25 * root
        result[mask, 3] = (matrix[mask, k, j] - matrix[mask, j, k]) / root
        result[mask, j] = (matrix[mask, j, axis] + matrix[mask, axis, j]) / root
        result[mask, k] = (matrix[mask, k, axis] + matrix[mask, axis, k]) / root
        remaining &= ~mask
    result /= np.maximum(np.linalg.norm(result, axis=-1, keepdims=True), 1e-8)
    return result


@torch.no_grad()
def evaluate(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    target_mean: np.ndarray,
    target_std: np.ndarray,
    cut_loss_fn: nn.Module,
) -> tuple[dict, dict[str, np.ndarray]]:
    model.eval()
    losses: list[float] = []
    predictions: list[np.ndarray] = []
    targets: list[np.ndarray] = []
    cut_probabilities: list[np.ndarray] = []
    cut_targets: list[np.ndarray] = []
    body_scales: list[np.ndarray] = []

    for inputs, target, cuts, scales in loader:
        inputs = inputs.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)
        cuts = cuts.to(device, non_blocking=True)
        predicted, cut_logit = model(inputs)
        loss, _ = grouped_loss(predicted, target, cut_logit, cuts, cut_loss_fn)
        losses.append(float(loss))
        predictions.append(predicted.cpu().numpy().reshape(-1, predicted.shape[-1]))
        targets.append(target.cpu().numpy().reshape(-1, target.shape[-1]))
        cut_probabilities.append(torch.sigmoid(cut_logit).cpu().numpy().reshape(-1))
        cut_targets.append(cuts.cpu().numpy().reshape(-1))
        body_scales.append(
            np.repeat(scales.numpy(), predicted.shape[1]).astype(np.float32)
        )

    predicted_standard = np.concatenate(predictions)
    target_standard = np.concatenate(targets)
    predicted = predicted_standard * target_std + target_mean
    target = target_standard * target_std + target_mean
    cut_probability = np.concatenate(cut_probabilities)
    cut_target = np.concatenate(cut_targets)
    scales = np.concatenate(body_scales)

    position_error = np.linalg.norm(predicted[:, :3] - target[:, :3], axis=-1)
    predicted_rotation = rotation_6d_to_matrix(predicted[:, 3:9])
    target_rotation = rotation_6d_to_matrix(target[:, 3:9])
    relative = np.swapaxes(predicted_rotation, -1, -2) @ target_rotation
    cosine = np.clip((np.trace(relative, axis1=1, axis2=2) - 1.0) * 0.5, -1.0, 1.0)
    rotation_error = np.degrees(np.arccos(cosine))

    predicted_cut = cut_probability >= 0.5
    actual_cut = cut_target >= 0.5
    true_positive = int(np.sum(predicted_cut & actual_cut))
    false_positive = int(np.sum(predicted_cut & ~actual_cut))
    false_negative = int(np.sum(~predicted_cut & actual_cut))
    precision = true_positive / max(1, true_positive + false_positive)
    recall = true_positive / max(1, true_positive + false_negative)
    f1 = 2.0 * precision * recall / max(1e-8, precision + recall)

    positive_count = int(np.sum(actual_cut))
    if positive_count > 0:
        order = np.argsort(-cut_probability)
        ordered_actual = actual_cut[order].astype(np.float64)
        cumulative_true_positive = np.cumsum(ordered_actual)
        rank = np.arange(1, len(order) + 1, dtype=np.float64)
        ranked_precision = cumulative_true_positive / rank
        ranked_recall = cumulative_true_positive / positive_count
        ranked_f1 = (
            2.0
            * ranked_precision
            * ranked_recall
            / np.maximum(ranked_precision + ranked_recall, 1e-8)
        )
        best_rank = int(np.argmax(ranked_f1))
        average_precision = float(
            np.sum(ranked_precision * ordered_actual) / positive_count
        )
        best_threshold = float(cut_probability[order[best_rank]])
        best_precision = float(ranked_precision[best_rank])
        best_recall = float(ranked_recall[best_rank])
        best_f1 = float(ranked_f1[best_rank])
    else:
        average_precision = 0.0
        best_threshold = 0.5
        best_precision = 0.0
        best_recall = 0.0
        best_f1 = 0.0

    metrics = {
        "loss": float(np.mean(losses)),
        "positionMeanErrorBodyScale": float(np.mean(position_error)),
        "positionMeanErrorMeters": float(np.mean(position_error * scales)),
        "rotationMeanErrorDegrees": float(np.mean(rotation_error)),
        "fovMeanAbsoluteErrorDegrees": float(
            np.mean(np.abs(predicted[:, 9] - target[:, 9]))
        ),
        "dutchMeanAbsoluteErrorDegrees": float(
            np.mean(np.abs(predicted[:, 10] - target[:, 10]))
        ),
        "cutPrecisionAt0_5": precision,
        "cutRecallAt0_5": recall,
        "cutF1At0_5": f1,
        "cutAveragePrecision": average_precision,
        "cutBestThreshold": best_threshold,
        "cutBestPrecision": best_precision,
        "cutBestRecall": best_recall,
        "cutBestF1": best_f1,
        "evaluatedFrames": int(len(predicted)),
    }
    arrays = {
        "predicted_camera_local": predicted.astype(np.float32),
        "target_camera_local": target.astype(np.float32),
        "predicted_cut_probability": cut_probability.astype(np.float32),
        "target_cut": cut_target.astype(np.float32),
    }
    return metrics, arrays


def main() -> None:
    configure_utf8_stdio()
    args = parse_args()
    seed_everything(args.seed)
    args.run_dir = args.run_dir.resolve()
    args.cache_dir = args.cache_dir.resolve()
    args.run_dir.mkdir(parents=True, exist_ok=True)

    manifest_paths = find_manifest_paths(args.data)
    records = load_records(
        manifest_paths,
        skip_unavailable=(args.skip_unavailable or args.single_subject_only),
    )
    single_subject_selection = None
    if args.single_subject_only:
        records, single_subject_selection = filter_single_subject_records(
            records,
            args.catalog,
            rebuild_cache=args.rebuild_cache,
        )
        if len(records) < 2:
            raise ValueError(
                "Single-subject catalog selection left fewer than two songs; "
                "a song-level validation split is not possible."
            )
    print(f"Found {len(records)} songs in {len(manifest_paths)} datasets.")

    if args.val_song_id:
        validation_record = select_record(
            records,
            song_name=args.val_song,
            song_id=args.val_song_id,
        )
        validation_name = validation_record.name
        validation_ids = {validation_record.id}
    else:
        validation_name = args.val_song or records[-1].name
        validation_ids = {
            record.id for record in records if record.name == validation_name
        }
        if not validation_ids:
            raise ValueError(f"Validation song not found: {validation_name}")

    audio_frame_masks: dict[str, AudioAlignedFrameMask] = {}
    audio_frame_mask_provenance: dict[str, Any] = {
        "contractVersion": "audio-aligned-frame-mask-v1",
        "rangeSemantics": "zero_based_half_open",
        "windowSelectionSemantics": "range_anchored_no_boundary_crossing",
        "cutBoundaryCensoring": (
            "zero_partial_mask_range_start_labels_with_full_range_legacy_parity"
        ),
        "windowCountReconciliation": (
            "legacy=eligibleLegacy+excludedLegacy; "
            "used=reusedLegacy+addedRangeAnchored; "
            "eligibleLegacy=reusedLegacy+omittedEligibleLegacy"
        ),
        "trainingPurpose": args.training_purpose,
        "consumerEnabled": args.training_purpose in FRAME_MASK_TRAINING_PURPOSES,
        "audioFingerprintVerification": None,
        "verifiedAudioFingerprintCount": 0,
        "maskContentAggregateSha256": None,
        "trainingIndex": None,
        "trainingIndexSha256": None,
        "candidateIndexSchemaVersion": None,
        "candidatePolicyVersion": None,
        "catalogSha256": None,
        "qualityReportSha256": None,
        "audioMaskConsumptionPolicy": None,
        "selectedSongCount": len(records),
        "maskedSongCount": 0,
        "legacyUnmaskedSongCount": len(records),
    }
    if args.training_index:
        training_index_path = args.training_index.resolve()
        records = filter_records_by_training_index(
            records,
            training_index_path,
            args.training_purpose,
            validation_ids,
        )
        audio_frame_masks, audio_frame_mask_provenance = load_audio_aligned_frame_masks(
            training_index_path,
            records,
            args.training_purpose,
        )
        print(
            f"Training index selected {len(records)} songs for "
            f"{args.training_purpose} including validation."
        )

    prepared_records: list[tuple[SongRecord, PreparedSong]] = []
    for record in records:
        print(f"Preparing {record.name} ...", flush=True)
        song = load_or_prepare_song(
            record,
            args.cache_dir,
            args.rebuild_cache,
            allow_cached_raw=not args.single_subject_only,
        )
        prepared_records.append((record, song))
        print(
            f"  {len(song.inputs):,} frames, input={song.inputs.shape[1]}, "
            f"bodyScale={song.body_scale:.4f}m"
        )

    validation_prepared_records = [
        (record, song)
        for record, song in prepared_records
        if record.id in validation_ids
    ]
    training_prepared_records = [
        (record, song)
        for record, song in prepared_records
        if record.id not in validation_ids
    ]
    validation_songs = [song for _, song in validation_prepared_records]
    training_songs = [song for _, song in training_prepared_records]
    if not training_songs:
        raise ValueError("The validation split leaves no training songs.")

    ranges_by_song_id = resolve_prepared_frame_ranges(
        prepared_records,
        audio_frame_masks,
    )
    if audio_frame_masks:
        train_inputs = concatenate_prepared_frame_ranges(
            training_prepared_records,
            ranges_by_song_id,
            "inputs",
        )
        train_targets = concatenate_prepared_frame_ranges(
            training_prepared_records,
            ranges_by_song_id,
            "targets",
        )
    else:
        train_inputs = np.concatenate([song.inputs for song in training_songs])
        train_targets = np.concatenate([song.targets for song in training_songs])
    input_mean = train_inputs.mean(axis=0, dtype=np.float64).astype(np.float32)
    input_std = train_inputs.std(axis=0, dtype=np.float64).astype(np.float32)
    target_mean = train_targets.mean(axis=0, dtype=np.float64).astype(np.float32)
    target_std = train_targets.std(axis=0, dtype=np.float64).astype(np.float32)
    input_std = np.maximum(input_std, 1e-5)
    target_std = np.maximum(target_std, 1e-5)

    training_frame_ranges = (
        [ranges_by_song_id[record.id] for record, _ in training_prepared_records]
        if audio_frame_masks
        else None
    )
    validation_frame_ranges = (
        [ranges_by_song_id[record.id] for record, _ in validation_prepared_records]
        if audio_frame_masks
        else None
    )
    training_dataset = WindowDataset(
        training_songs,
        args.window,
        args.train_stride,
        input_mean,
        input_std,
        target_mean,
        target_std,
        frame_ranges=training_frame_ranges,
        censor_cut_range_starts=(
            [
                record.id in audio_frame_masks
                for record, _ in training_prepared_records
            ]
            if audio_frame_masks
            else None
        ),
    )
    validation_dataset = WindowDataset(
        validation_songs,
        args.window,
        args.window,
        input_mean,
        input_std,
        target_mean,
        target_std,
        frame_ranges=validation_frame_ranges,
        censor_cut_range_starts=(
            [
                record.id in audio_frame_masks
                for record, _ in validation_prepared_records
            ]
            if audio_frame_masks
            else None
        ),
    )
    if not training_dataset or not validation_dataset:
        raise ValueError("The selected window size produced an empty dataset.")

    device = torch.device(args.device)
    pin_memory = device.type == "cuda"
    training_loader = DataLoader(
        training_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=pin_memory,
    )
    validation_loader = DataLoader(
        validation_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=pin_memory,
    )

    model = CameraDirector(
        input_size=training_songs[0].inputs.shape[1],
        output_size=training_songs[0].targets.shape[1],
        hidden_size=args.hidden_size,
        layers=args.layers,
        dropout=args.dropout,
    ).to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.learning_rate,
        weight_decay=1e-4,
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=args.epochs,
        eta_min=args.learning_rate * 0.05,
    )
    if audio_frame_masks:
        training_cuts = concatenate_prepared_frame_ranges(
            training_prepared_records,
            ranges_by_song_id,
            "cuts",
            cut_range_start_censor_song_ids=set(audio_frame_masks),
        )
    else:
        training_cuts = np.concatenate([song.cuts for song in training_songs])
    positive_count = max(1.0, float(training_cuts.sum()))
    negative_count = float(len(training_cuts) - positive_count)
    positive_weight = min(100.0, negative_count / positive_count)
    cut_loss_fn = nn.BCEWithLogitsLoss(
        pos_weight=torch.tensor(positive_weight, device=device)
    )
    scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")
    audio_frame_mask_report = build_audio_frame_mask_report(
        audio_frame_mask_provenance,
        audio_frame_masks,
        ranges_by_song_id,
        training_prepared_records,
        validation_prepared_records,
        training_dataset,
        validation_dataset,
    )
    audio_frame_mask_metrics = dict(audio_frame_mask_report)
    audio_frame_mask_metrics.pop("maskedRecords", None)
    for split_name in ("training", "validation"):
        split_metrics = dict(audio_frame_mask_report[split_name])
        split_metrics.pop("songs", None)
        audio_frame_mask_metrics[split_name] = split_metrics
    training_provenance = {
        "trainingPurpose": args.training_purpose,
        "trainingIndexSha256": audio_frame_mask_report[
            "trainingIndexSha256"
        ],
        "maskContractVersion": audio_frame_mask_report["contractVersion"],
        "maskRangeSemantics": audio_frame_mask_report["rangeSemantics"],
        "maskCutBoundaryCensoring": audio_frame_mask_report[
            "cutBoundaryCensoring"
        ],
        "maskContentAggregateSha256": audio_frame_mask_report[
            "maskContentAggregateSha256"
        ],
        "maskedSongCount": audio_frame_mask_report["maskedSongCount"],
        "verifiedAudioFingerprintCount": audio_frame_mask_report[
            "verifiedAudioFingerprintCount"
        ],
        "catalogSha256": audio_frame_mask_report["catalogSha256"],
        "qualityReportSha256": audio_frame_mask_report[
            "qualityReportSha256"
        ],
    }

    config = {
        key: str(value) if isinstance(value, Path) else value
        for key, value in vars(args).items()
    }
    config.update(
        {
            "run_dir": str(args.run_dir),
            "cache_dir": str(args.cache_dir),
            "catalog": str(args.catalog.resolve()) if args.catalog else None,
            "singleSubjectSelection": single_subject_selection,
            "manifests": [str(path) for path in manifest_paths],
            "trainingSongs": [song.name for song in training_songs],
            "validationSongs": [song.name for song in validation_songs],
            "trainingSongIds": [
                record.id
                for record, _ in prepared_records
                if record.id not in validation_ids
            ],
            "validationSongIds": sorted(validation_ids),
            "canonicalJoints": list(CANONICAL_JOINTS),
            "inputSize": int(training_songs[0].inputs.shape[1]),
            "outputSize": int(training_songs[0].targets.shape[1]),
            "trainingWindows": len(training_dataset),
            "validationWindows": len(validation_dataset),
            "audioAlignedFrameMask": audio_frame_mask_report,
            "trainingProvenance": training_provenance,
            "cutPositiveWeight": positive_weight,
            "deviceResolved": str(device),
            "architecture": "dilated_tcn",
        }
    )
    (args.run_dir / "config.json").write_text(
        json.dumps(config, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    np.savez(
        args.run_dir / "normalization.npz",
        input_mean=input_mean,
        input_std=input_std,
        target_mean=target_mean,
        target_std=target_std,
    )

    print(
        f"Training on {device}: {len(training_dataset):,} windows, "
        f"validating on '{validation_name}' ({len(validation_dataset):,} windows)."
    )
    history: list[dict] = []
    best_validation_loss = math.inf
    best_epoch = 0
    epochs_without_improvement = 0

    for epoch in range(1, args.epochs + 1):
        model.train()
        epoch_losses: list[float] = []
        for inputs, targets, cuts, _ in training_loader:
            inputs = inputs.to(device, non_blocking=True)
            targets = targets.to(device, non_blocking=True)
            cuts = cuts.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            with torch.autocast(
                device_type=device.type,
                dtype=torch.float16,
                enabled=device.type == "cuda",
            ):
                predicted, cut_logit = model(inputs)
                loss, _ = grouped_loss(
                    predicted,
                    targets,
                    cut_logit,
                    cuts,
                    cut_loss_fn,
                )
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
            epoch_losses.append(float(loss.detach()))

        validation_metrics, _ = evaluate(
            model,
            validation_loader,
            device,
            target_mean,
            target_std,
            cut_loss_fn,
        )
        scheduler.step()
        epoch_result = {
            "epoch": epoch,
            "trainLoss": float(np.mean(epoch_losses)),
            "learningRate": float(optimizer.param_groups[0]["lr"]),
            **validation_metrics,
        }
        history.append(epoch_result)
        print(
            f"Epoch {epoch:02d}/{args.epochs}: "
            f"train={epoch_result['trainLoss']:.4f}, "
            f"val={epoch_result['loss']:.4f}, "
            f"pos={epoch_result['positionMeanErrorMeters']:.3f}m, "
            f"rot={epoch_result['rotationMeanErrorDegrees']:.1f}deg, "
            f"fov={epoch_result['fovMeanAbsoluteErrorDegrees']:.2f}deg, "
            f"cutF1={epoch_result['cutF1At0_5']:.3f}",
            flush=True,
        )

        if validation_metrics["loss"] < best_validation_loss:
            best_validation_loss = validation_metrics["loss"]
            best_epoch = epoch
            epochs_without_improvement = 0
            torch.save(
                {
                    "modelState": model.state_dict(),
                    "modelConfig": {
                        "architecture": "dilated_tcn",
                        "input_size": training_songs[0].inputs.shape[1],
                        "output_size": training_songs[0].targets.shape[1],
                        "hidden_size": args.hidden_size,
                        "layers": args.layers,
                        "dropout": args.dropout,
                    },
                    "canonicalJoints": CANONICAL_JOINTS,
                    "inputMean": torch.from_numpy(input_mean),
                    "inputStd": torch.from_numpy(input_std),
                    "targetMean": torch.from_numpy(target_mean),
                    "targetStd": torch.from_numpy(target_std),
                    "trainingProvenance": training_provenance,
                    "epoch": epoch,
                    "validationLoss": best_validation_loss,
                },
                args.run_dir / "best_model.pt",
            )
        else:
            epochs_without_improvement += 1
            if epochs_without_improvement >= args.patience:
                print(
                    f"Early stopping at epoch {epoch}; "
                    f"best epoch was {best_epoch}.",
                    flush=True,
                )
                break

    checkpoint = torch.load(
        args.run_dir / "best_model.pt",
        map_location=device,
        weights_only=False,
    )
    model.load_state_dict(checkpoint["modelState"])
    final_metrics, prediction_arrays = evaluate(
        model,
        validation_loader,
        device,
        target_mean,
        target_std,
        cut_loss_fn,
    )
    final_metrics.update(
        {
            "bestEpoch": best_epoch,
            "bestValidationLoss": best_validation_loss,
            "validationSong": validation_name,
            "audioAlignedFrameMask": audio_frame_mask_metrics,
        }
    )
    np.savez(args.run_dir / "validation_predictions.npz", **prediction_arrays)
    (args.run_dir / "history.json").write_text(
        json.dumps(history, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    (args.run_dir / "metrics.json").write_text(
        json.dumps(final_metrics, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(final_metrics, ensure_ascii=False, indent=2))
    print(f"Run saved to: {args.run_dir}")


if __name__ == "__main__":
    main()
