#!/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 dataclasses import dataclass
from pathlib import Path
from typing import Iterable

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

# 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}".encode("utf-8")
        return hashlib.sha256(value).hexdigest()[:16]

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

    @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


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=("jointTcn", "trajectoryModel", "cutModel"),
        default="jointTcn",
    )
    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")
    return parser.parse_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]:
    payload = json.loads(index_path.read_text(encoding="utf-8"))
    entries = {
        str(entry["songId"]): entry
        for entry in payload.get("entries", [])
    }
    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,
            )
        )
    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 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,
) -> 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:
        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


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,
    ) -> None:
        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]] = []
        for song_index, song in enumerate(songs):
            self.windows.extend(
                (song_index, start)
                for start in range(0, len(song.inputs) - window + 1, stride)
            )

    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
        return (
            torch.from_numpy(inputs),
            torch.from_numpy(targets),
            torch.from_numpy(song.cuts[start:end]),
            torch.tensor(song.body_scale, dtype=torch.float32),
        )


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,
    )
    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}")

    if args.training_index:
        records = filter_records_by_training_index(
            records,
            args.training_index.resolve(),
            args.training_purpose,
            validation_ids,
        )
        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)
        prepared_records.append((record, song))
        print(
            f"  {len(song.inputs):,} frames, input={song.inputs.shape[1]}, "
            f"bodyScale={song.body_scale:.4f}m"
        )

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

    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_dataset = WindowDataset(
        training_songs,
        args.window,
        args.train_stride,
        input_mean,
        input_std,
        target_mean,
        target_std,
    )
    validation_dataset = WindowDataset(
        validation_songs,
        args.window,
        args.window,
        input_mean,
        input_std,
        target_mean,
        target_std,
    )
    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,
    )
    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")

    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),
            "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),
            "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),
                    "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,
        }
    )
    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()
