"""Shared head, labels, and metrics for the EasyCom/Ego4D experiments."""

from __future__ import annotations

import csv
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, Sequence
from uuid import UUID

import numpy as np
import torch
from torch import nn

if TYPE_CHECKING:
    from local_turn_features import FramePersonFeatures


OFFICIAL_TRAIN_SESSIONS = (
    "Session_7",
    "Session_11",
    "Session_10",
    "Session_6",
    "Session_9",
    "Session_4",
    "Session_8",
    "Session_5",
    "Session_12",
)
OFFICIAL_TEST_SESSIONS = ("Session_3", "Session_2", "Session_1")
CLASS_NAMES = (
    "background",
    "camera_wearer_speaking",
    "other_people_speaking",
)
FOUR_STATE_NAMES = (
    "silence",
    "wearer_only",
    "others_only",
    "overlap",
)
STEP_SECONDS = 0.2
NUM_HORIZONS = 10
NUM_CLASSES = 3
NUM_FOUR_STATES = len(FOUR_STATE_NAMES)
LOSS_NAMES = (
    "lstr-mce",
    "miniroad-normalized",
    "independent-bce",
    "four-state-ce",
)
TRANSITION_OBJECTIVE_NAMES = (
    "none",
    "r1-supervised",
    "r2-expected-reward",
)
DATASET_NAMES = ("easycom", "ego4d")
LAYER_FUSION_MODES = (
    "single",
    "fixed-mean",
    "scalar-mixture",
    "concat-mlp",
)
CURRENT_FRAME_PERSON_SET_FUSION_MODES = (
    "frame-set-attention",
    "oracle-frame-set-attention",
    "audio-person-frame-set",
    "gated-person-state-adapter",
)
TRACK_PERSON_FUSION_MODES = ("track-gru-residual",)
PERSON_SET_FUSION_MODES = (
    *CURRENT_FRAME_PERSON_SET_FUSION_MODES,
    *TRACK_PERSON_FUSION_MODES,
)
LOCAL_FUSION_MODES = ("pooled", *PERSON_SET_FUSION_MODES)
LOCAL_FUSION_SCOPES = {
    "pooled": "pooled-baseline",
    "frame-set-attention": "f1-current-frame-person-set",
    "oracle-frame-set-attention": "causal-oracle-visible-person-speaking-history",
    "audio-person-frame-set": "f5-screening-current-frame-no-person-history",
    "gated-person-state-adapter": "f6-screening-post-thinker-not-internal",
    "track-gru-residual": "causal-anonymous-track-gru-pooled-baseline-residual",
}
DatasetName = Literal["easycom", "ego4d"]


class TurnPredictionHead(nn.Module):
    """Map one Omni state to configurable ``[horizon, output]`` logits."""

    def __init__(
        self,
        hidden_size: int = 2048,
        num_horizons: int = NUM_HORIZONS,
        num_classes: int = NUM_CLASSES,
        dropout: float = 0.1,
        num_input_layers: int = 1,
        learnable_layer_mixture: bool = False,
        layer_fusion: str | None = None,
        transition_auxiliary: bool = False,
        local_visual_dim: int = 0,
        local_audio_dim: int = 0,
        local_projection_dim: int = 256,
        local_fusion: str = "pooled",
        person_attention_dim: int = 256,
        person_attention_heads: int = 4,
        person_history_frames: int = 16,
        person_track_max_gap_seconds: float = 1.0,
    ) -> None:
        super().__init__()
        if num_input_layers <= 0:
            raise ValueError("num_input_layers must be positive")
        if layer_fusion is None:
            layer_fusion = (
                "scalar-mixture" if learnable_layer_mixture else "single"
            )
        if layer_fusion not in LAYER_FUSION_MODES:
            raise ValueError(
                f"Unknown layer_fusion={layer_fusion!r}; expected one of "
                f"{LAYER_FUSION_MODES}"
            )
        if learnable_layer_mixture and layer_fusion != "scalar-mixture":
            raise ValueError(
                "learnable_layer_mixture=True is only compatible with "
                "layer_fusion='scalar-mixture'"
            )
        if layer_fusion != "single" and num_input_layers < 2:
            raise ValueError(
                f"layer_fusion={layer_fusion!r} requires at least two input layers"
            )
        if layer_fusion == "single" and num_input_layers != 1:
            raise ValueError(
                "Multiple input layers require a non-single layer_fusion"
            )
        self.hidden_size = hidden_size
        self.num_horizons = num_horizons
        self.num_classes = num_classes
        self.dropout = dropout
        self.num_input_layers = num_input_layers
        self.layer_fusion = layer_fusion
        self.learnable_layer_mixture = layer_fusion == "scalar-mixture"
        self.transition_auxiliary = transition_auxiliary
        if min(local_visual_dim, local_audio_dim, local_projection_dim) < 0:
            raise ValueError("Local dimensions must be non-negative")
        self.local_visual_dim = local_visual_dim
        self.local_audio_dim = local_audio_dim
        self.local_projection_dim = local_projection_dim
        if local_fusion not in LOCAL_FUSION_MODES:
            raise ValueError(
                f"Unknown local_fusion={local_fusion!r}; expected one of "
                f"{LOCAL_FUSION_MODES}"
            )
        if person_attention_dim <= 0 or person_attention_heads <= 0:
            raise ValueError("Person attention dimension/heads must be positive")
        if person_attention_dim % person_attention_heads:
            raise ValueError("Person attention dimension must be divisible by heads")
        if person_history_frames <= 0:
            raise ValueError("person_history_frames must be positive")
        if person_track_max_gap_seconds <= 0:
            raise ValueError("person_track_max_gap_seconds must be positive")
        if local_fusion in PERSON_SET_FUSION_MODES and (
            local_visual_dim <= 0 or local_audio_dim <= 0
        ):
            raise ValueError(
                "Person-set fusion requires both local visual people and audio"
            )
        self.local_fusion = local_fusion
        self.person_attention_dim = person_attention_dim
        self.person_attention_heads = person_attention_heads
        self.person_history_frames = person_history_frames
        self.person_track_max_gap_seconds = person_track_max_gap_seconds
        if self.learnable_layer_mixture:
            # Zero initialization exactly reproduces a uniform layer mean at
            # step zero while allowing training to select useful query layers.
            self.layer_logits = nn.Parameter(torch.zeros(num_input_layers))
        else:
            self.register_parameter("layer_logits", None)
        if layer_fusion == "concat-mlp":
            fusion_input_size = num_input_layers * hidden_size
            self.layer_fusion_net: nn.Module | None = nn.Sequential(
                nn.LayerNorm(fusion_input_size),
                nn.Linear(fusion_input_size, hidden_size),
                nn.GELU(),
            )
        else:
            self.layer_fusion_net = None
        self.local_projections = nn.ModuleDict()
        fused_size = hidden_size
        projected_modalities = (
            (("audio", local_audio_dim),)
            if local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES
            else (("visual", local_visual_dim), ("audio", local_audio_dim))
        )
        for name, dimension in projected_modalities:
            if dimension:
                self.local_projections[name] = (
                    nn.Sequential(
                        nn.LayerNorm(dimension),
                        nn.Linear(dimension, local_projection_dim),
                        nn.GELU(),
                    )
                    if local_projection_dim
                    else nn.Identity()
                )
                fused_size += (local_projection_dim or dimension) + 1
        self.person_input_projection: nn.Module | None = None
        self.person_state_query: nn.Module | None = None
        self.person_audio_query: nn.Module | None = None
        self.person_attention: nn.MultiheadAttention | None = None
        self.null_person: nn.Parameter | None = None
        self.audio_person_attention: nn.MultiheadAttention | None = None
        self.audio_person_gate: nn.Parameter | None = None
        self.person_state_adapter: nn.Module | None = None
        self.person_state_adapter_gate: nn.Parameter | None = None
        self.person_track_gru: nn.GRU | None = None
        self.person_track_delta: nn.Module | None = None
        self.person_track_residual_gate: nn.Parameter | None = None
        if local_fusion in PERSON_SET_FUSION_MODES:
            self.person_input_projection = nn.Sequential(
                nn.LayerNorm(local_visual_dim),
                nn.Linear(local_visual_dim, person_attention_dim),
                nn.GELU(),
            )
            self.person_state_query = nn.Sequential(
                nn.LayerNorm(hidden_size),
                nn.Linear(hidden_size, person_attention_dim),
            )
            self.person_audio_query = nn.Sequential(
                nn.LayerNorm(local_audio_dim),
                nn.Linear(local_audio_dim, person_attention_dim),
            )
            self.person_attention = nn.MultiheadAttention(
                person_attention_dim,
                person_attention_heads,
                dropout=dropout,
                batch_first=True,
            )
            self.null_person = nn.Parameter(
                torch.zeros(1, 1, person_attention_dim)
            )
            nn.init.normal_(self.null_person, std=0.02)
            if local_fusion == "audio-person-frame-set":
                # Screening version of the F5 reverse read: an audio-only
                # query first reads the current person memory, then a
                # zero-initialized residual gate adds that evidence to the
                # regular G+A floor query.  This remains a current-frame
                # module; it is not the full history-conditioned F5 block.
                self.audio_person_attention = nn.MultiheadAttention(
                    person_attention_dim,
                    person_attention_heads,
                    dropout=dropout,
                    batch_first=True,
                )
                self.audio_person_gate = nn.Parameter(torch.zeros(()))
            if local_fusion == "gated-person-state-adapter":
                # Screening version of F6 for precomputed frozen states.  It
                # injects selected person evidence after the Thinker readout,
                # not into a Thinker transformer block.  The scalar gate makes
                # the visual residual exactly zero at initialization.
                self.person_state_adapter = nn.Sequential(
                    nn.LayerNorm(person_attention_dim),
                    nn.Linear(person_attention_dim, person_attention_dim),
                    nn.GELU(),
                    nn.Linear(person_attention_dim, hidden_size),
                )
                self.person_state_adapter_gate = nn.Parameter(torch.zeros(()))
            elif local_fusion == "track-gru-residual":
                # Apply one shared causal GRU independently to every anonymous
                # face track. The regular pooled G+A+V head remains the exact
                # base path; this branch can only add a zero-gated correction.
                self.person_track_gru = nn.GRU(
                    person_attention_dim,
                    person_attention_dim,
                    num_layers=1,
                    batch_first=True,
                )
                delta_input = hidden_size // 2 + person_attention_dim + 1
                self.person_track_delta = nn.Sequential(
                    nn.LayerNorm(delta_input),
                    nn.Linear(delta_input, hidden_size // 2),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(hidden_size // 2, num_horizons * num_classes),
                )
                self.person_track_residual_gate = nn.Parameter(torch.zeros(()))
            else:
                fused_size += person_attention_dim + 1
        self.net = nn.Sequential(
            nn.LayerNorm(fused_size),
            nn.Linear(fused_size, hidden_size // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, num_horizons * num_classes),
        )
        self.transition_head: nn.Module | None
        if transition_auxiliary:
            # Share the activity MLP trunk and add only a small binary readout
            # for each future horizon.
            self.transition_head = nn.Linear(hidden_size // 2, num_horizons)
        else:
            self.transition_head = None

    def _fuse_layers(self, states: torch.Tensor) -> torch.Tensor:
        if states.shape[-1] != self.hidden_size:
            raise ValueError(
                f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}"
            )
        if self.layer_fusion != "single":
            if states.ndim < 3 or states.shape[-2] != self.num_input_layers:
                raise ValueError(
                    "Multi-layer input must end in "
                    f"[{self.num_input_layers},{self.hidden_size}], got "
                    f"{tuple(states.shape)}"
                )
        if self.layer_fusion == "fixed-mean":
            states = states.mean(dim=-2)
        elif self.layer_fusion == "scalar-mixture":
            assert self.layer_logits is not None
            weights = self.layer_logits.softmax(dim=0).to(states.dtype)
            view_shape = (1,) * (states.ndim - 2) + (
                self.num_input_layers,
                1,
            )
            states = (states * weights.view(view_shape)).sum(dim=-2)
        elif self.layer_fusion == "concat-mlp":
            assert self.layer_fusion_net is not None
            states = states.flatten(start_dim=-2)
            states = self.layer_fusion_net(states)
        return states

    @property
    def local_input_size(self) -> int:
        if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
            return self.local_audio_dim + 1
        return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)

    def _frame_set_context(
        self,
        states: torch.Tensor,
        local: torch.Tensor,
        person_features: torch.Tensor | None,
        person_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        if states.ndim != 2:
            raise ValueError("Person-set fusion expects batched [B,D] states")
        if local.shape != (len(states), self.local_audio_dim + 1):
            raise ValueError(
                "Person-set fusion requires aligned [audio,valid] local features"
            )
        if (
            person_features is None
            or person_features.ndim != 3
            or person_features.shape[0] != len(states)
            or person_features.shape[2] != self.local_visual_dim
        ):
            raise ValueError(
                "Person-set fusion requires person features [B,P,visual_dim]"
            )
        if (
            person_mask is None
            or person_mask.dtype != torch.bool
            or person_mask.shape != person_features.shape[:2]
        ):
            raise ValueError("Person-set fusion requires a boolean person mask [B,P]")
        if not torch.isfinite(person_features).all():
            raise ValueError("Person features contain NaN/Inf")
        assert self.person_input_projection is not None
        assert self.person_state_query is not None
        assert self.person_audio_query is not None
        assert self.person_attention is not None
        assert self.null_person is not None

        audio_values = local[:, : self.local_audio_dim]
        audio_valid = local[:, self.local_audio_dim : self.local_audio_dim + 1]
        audio_projected = self.local_projections["audio"](audio_values) * audio_valid
        audio_query = self.person_audio_query(audio_values) * audio_valid
        query = self.person_state_query(states) + audio_query
        people = self.person_input_projection(person_features)
        null = self.null_person.expand(len(states), -1, -1)
        memory = torch.cat((null, people), dim=1)
        null_mask = torch.ones(
            (len(states), 1),
            dtype=torch.bool,
            device=person_mask.device,
        )
        memory_valid = torch.cat((null_mask, person_mask), dim=1)
        if self.audio_person_attention is not None:
            assert self.audio_person_gate is not None
            audio_context, _ = self.audio_person_attention(
                audio_query[:, None, :],
                memory,
                memory,
                key_padding_mask=~memory_valid,
                need_weights=False,
            )
            query = query + (
                torch.tanh(self.audio_person_gate)
                * audio_context[:, 0, :]
                * audio_valid
            )
        context, _ = self.person_attention(
            query[:, None, :],
            memory,
            memory,
            key_padding_mask=~memory_valid,
            need_weights=False,
        )
        person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype)
        return audio_projected, audio_valid, context[:, 0, :], person_valid

    def _track_gru_context(
        self,
        states: torch.Tensor,
        local: torch.Tensor,
        person_histories: torch.Tensor | None,
        history_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Encode each causal track with a shared GRU, then query the set."""

        if states.ndim != 2:
            raise ValueError("Track GRU fusion expects batched [B,D] states")
        if local.shape != (len(states), self.local_input_size):
            raise ValueError("Track GRU fusion requires pooled [visual,audio,valid] input")
        if (
            person_histories is None
            or person_histories.ndim != 4
            or person_histories.shape[0] != len(states)
            or person_histories.shape[2] != self.person_history_frames
            or person_histories.shape[3] != self.local_visual_dim
        ):
            raise ValueError(
                "Track GRU fusion requires person histories "
                "[B,P,person_history_frames,visual_dim]"
            )
        if (
            history_mask is None
            or history_mask.dtype != torch.bool
            or history_mask.shape != person_histories.shape[:3]
        ):
            raise ValueError("Track GRU fusion requires a boolean history mask [B,P,L]")
        if not torch.isfinite(person_histories).all():
            raise ValueError("Person track histories contain NaN/Inf")
        if bool((history_mask[:, :, 1:] & ~history_mask[:, :, :-1]).any()):
            raise ValueError("Person track histories must be left aligned")
        assert self.person_input_projection is not None
        assert self.person_state_query is not None
        assert self.person_audio_query is not None
        assert self.person_attention is not None
        assert self.null_person is not None
        assert self.person_track_gru is not None

        batch, people, length, _ = person_histories.shape
        person_valid = history_mask.any(dim=-1)
        if people:
            projected = self.person_input_projection(person_histories)
            flat = projected.reshape(batch * people, length, -1)
            lengths = history_mask.sum(dim=-1).reshape(-1)
            safe_lengths = lengths.clamp_min(1).to(device="cpu")
            packed = nn.utils.rnn.pack_padded_sequence(
                flat,
                safe_lengths,
                batch_first=True,
                enforce_sorted=False,
            )
            _, hidden = self.person_track_gru(packed)
            track_states = hidden[-1].reshape(batch, people, -1)
            track_states = track_states * person_valid[..., None]
        else:
            track_states = states.new_zeros(
                (batch, 0, self.person_attention_dim)
            )

        visual_width = self.local_visual_dim + 1
        audio_values = local[:, visual_width : visual_width + self.local_audio_dim]
        audio_valid = local[:, visual_width + self.local_audio_dim :]
        if audio_valid.shape != (batch, 1):
            raise AssertionError("Unexpected pooled AV local feature layout")
        query = self.person_state_query(states)
        query = query + self.person_audio_query(audio_values) * audio_valid
        memory = torch.cat(
            (self.null_person.expand(batch, -1, -1), track_states), dim=1
        )
        memory_valid = torch.cat(
            (
                torch.ones((batch, 1), dtype=torch.bool, device=states.device),
                person_valid,
            ),
            dim=1,
        )
        context, _ = self.person_attention(
            query[:, None, :],
            memory,
            memory,
            key_padding_mask=~memory_valid,
            need_weights=False,
        )
        any_person = person_valid.any(dim=1, keepdim=True).to(states.dtype)
        return context[:, 0, :] * any_person, any_person

    def _shared_features(
        self,
        states: torch.Tensor,
        local: torch.Tensor | None = None,
        person_features: torch.Tensor | None = None,
        person_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Return the activity MLP representation shared by both readouts."""

        states = self._fuse_layers(states)
        if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
            if local is None:
                raise ValueError("Person-set fusion requires local audio features")
            audio, audio_valid, person_context, person_valid = self._frame_set_context(
                states,
                local,
                person_features,
                person_mask,
            )
            if self.local_fusion == "gated-person-state-adapter":
                assert self.person_state_adapter is not None
                assert self.person_state_adapter_gate is not None
                states = states + (
                    torch.tanh(self.person_state_adapter_gate)
                    * self.person_state_adapter(person_context)
                    * person_valid
                )
                states = torch.cat((states, audio, audio_valid), dim=-1)
            else:
                states = torch.cat(
                    (states, audio, audio_valid, person_context, person_valid),
                    dim=-1,
                )
        elif self.local_input_size:
            if local is None or local.shape != (*states.shape[:-1], self.local_input_size):
                raise ValueError(
                    f"Head requires aligned local features of width {self.local_input_size}"
                )
            pieces = [states]
            offset = 0
            for name, dimension in (
                ("visual", self.local_visual_dim),
                ("audio", self.local_audio_dim),
            ):
                if dimension:
                    values = local[..., offset : offset + dimension]
                    valid = local[..., offset + dimension : offset + dimension + 1]
                    # Mask AFTER projection too: a learned bias is not an observation.
                    pieces.extend((self.local_projections[name](values) * valid, valid))
                    offset += dimension + 1
            states = torch.cat(pieces, dim=-1)
        elif local is not None and local.shape[-1] != 0:
            raise ValueError("Local features supplied to a baseline head")
        if self.local_fusion not in PERSON_SET_FUSION_MODES and (
            person_features is not None or person_mask is not None
        ):
            raise ValueError("Person sets supplied to a non-frame-set head")
        return self.net[:-1](states)

    def _activity_logits(
        self,
        states: torch.Tensor,
        local: torch.Tensor | None,
        person_features: torch.Tensor | None,
        person_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        features = self._shared_features(
            states,
            local,
            person_features,
            person_mask,
        )
        flat_logits = self.net[-1](features)
        if self.local_fusion == "track-gru-residual":
            if local is None:
                raise ValueError("Track GRU residual requires pooled local AV")
            assert self.person_track_delta is not None
            assert self.person_track_residual_gate is not None
            fused_states = self._fuse_layers(states)
            context, person_valid = self._track_gru_context(
                fused_states,
                local,
                person_features,
                person_mask,
            )
            delta = self.person_track_delta(
                torch.cat((features, context, person_valid), dim=-1)
            )
            flat_logits = flat_logits + (
                torch.tanh(self.person_track_residual_gate)
                * delta
                * person_valid
            )
        return features, flat_logits

    def forward(
        self,
        states: torch.Tensor,
        local: torch.Tensor | None = None,
        person_features: torch.Tensor | None = None,
        person_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        features, logits = self._activity_logits(
            states,
            local,
            person_features,
            person_mask,
        )
        return logits.view(
            *features.shape[:-1], self.num_horizons, self.num_classes
        )

    def forward_with_transition(
        self,
        states: torch.Tensor,
        local: torch.Tensor | None = None,
        person_features: torch.Tensor | None = None,
        person_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return activity logits and binary transition logits."""

        if self.transition_head is None:
            raise RuntimeError(
                "This head was created without transition_auxiliary=True"
            )
        features, flat_activity_logits = self._activity_logits(
            states,
            local,
            person_features,
            person_mask,
        )
        activity_logits = flat_activity_logits.view(
            *features.shape[:-1],
            self.num_horizons,
            self.num_classes,
        )
        transition_logits = self.transition_head(features)
        return activity_logits, transition_logits

    def layer_mixture_weights(self) -> torch.Tensor | None:
        """Return normalized oldest-to-newest layer weights, if enabled."""

        if self.layer_logits is None:
            return None
        return self.layer_logits.softmax(dim=0)

    def person_interaction_gates(self) -> dict[str, float]:
        """Return bounded residual strengths for screening fusion modes."""

        result: dict[str, float] = {}
        if self.audio_person_gate is not None:
            result["audio_person"] = float(
                torch.tanh(self.audio_person_gate.detach()).cpu()
            )
        if self.person_state_adapter_gate is not None:
            result["person_state_adapter"] = float(
                torch.tanh(self.person_state_adapter_gate.detach()).cpu()
            )
        if self.person_track_residual_gate is not None:
            result["person_track_residual"] = float(
                torch.tanh(self.person_track_residual_gate.detach()).cpu()
            )
        return result


def _validate_loss_inputs(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    num_output_classes: int = NUM_CLASSES,
) -> None:
    if logits.shape != targets.shape:
        raise ValueError(
            f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
            f"{tuple(targets.shape)}"
        )
    if logits.shape[-1] != num_output_classes:
        raise ValueError(
            f"Expected {num_output_classes} classes, got {logits.shape[-1]}"
        )


def _validated_binary_activity_targets(
    targets: torch.Tensor,
    *,
    allow_noncanonical_binary: bool = False,
) -> torch.Tensor:
    """Validate three-label activity rows.

    EasyCom uses four explicit activity states.  Ego4D's released target files
    are consumed by EgoSpeak's MCE objective as binary/multi-positive weights,
    so the compatibility path preserves any binary row instead of silently
    rewriting it to an EasyCom state.
    """

    if targets.ndim < 1 or targets.shape[-1] != NUM_CLASSES:
        raise ValueError(
            f"Expected activity targets [...,{NUM_CLASSES}], got "
            f"{tuple(targets.shape)}"
        )
    if not torch.isfinite(targets).all():
        raise ValueError("Activity targets contain NaN or Inf")
    rounded = targets.round()
    if not torch.allclose(targets, rounded, atol=1e-6, rtol=0.0):
        raise ValueError("Activity targets must contain only binary 0/1 values")
    expected_background = rounded[..., 1:].sum(dim=-1) == 0
    valid = rounded[..., 0].bool() == expected_background
    if allow_noncanonical_binary:
        return rounded
    if not bool(valid.all()):
        invalid = torch.unique(rounded[~valid].reshape(-1, NUM_CLASSES), dim=0)
        invalid_rows = ", ".join(
            str([int(value) for value in row.tolist()])
            for row in invalid[:8]
        )
        if len(invalid) > 8:
            invalid_rows += ", ..."
        raise ValueError(
            "Activity targets must be one of [1,0,0], [0,1,0], "
            "[0,0,1], or [0,1,1]; observed invalid "
            f"row(s): {invalid_rows}"
        )
    return rounded


def activity_targets_to_four_states(targets: torch.Tensor) -> torch.LongTensor:
    """Map three-label activity targets to silence/wearer/others/overlap."""

    targets = _validated_binary_activity_targets(targets)
    return (targets[..., 1] + 2 * targets[..., 2]).long()


def transition_targets_from_activity(
    activity_sequence: torch.Tensor,
) -> torch.Tensor:
    """Mark every activity-state change between adjacent 5 Hz target rows.

    The input ends in [time, 3] and the output ends in [time - 1].
    A silence/onset, offset, speaker handover, or overlap-state change is one
    positive transition.
    """

    activity_sequence = _validated_binary_activity_targets(activity_sequence)
    if activity_sequence.ndim < 2 or activity_sequence.shape[-2] < 2:
        raise ValueError(
            "Activity sequence must contain at least two target rows, got "
            f"{tuple(activity_sequence.shape)}"
        )
    changed = activity_sequence[..., 1:, :] != activity_sequence[..., :-1, :]
    return changed.any(dim=-1).to(activity_sequence.dtype)


def weighted_transition_bce_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    pos_weight: float = 3.0,
) -> torch.Tensor:
    """Binary transition loss with an explicit positive-class weight."""

    if logits.shape != targets.shape:
        raise ValueError(
            f"Transition logits/targets mismatch: {tuple(logits.shape)} vs "
            f"{tuple(targets.shape)}"
        )
    if pos_weight <= 0:
        raise ValueError("Transition pos_weight must be positive")
    targets = targets.to(device=logits.device, dtype=logits.dtype)
    if not torch.isfinite(targets).all() or not torch.all(
        (targets == 0) | (targets == 1)
    ):
        raise ValueError("Transition targets must contain only finite 0/1 values")
    return nn.functional.binary_cross_entropy_with_logits(
        logits,
        targets,
        pos_weight=logits.new_tensor(pos_weight),
    )


def expected_transition_reward(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    pos_weight: float = 3.0,
) -> torch.Tensor:
    """Exact expected weighted correctness of a Bernoulli transition policy.

    Correct transition/non-transition actions receive reward one and incorrect
    actions zero. The expectation is evaluated analytically rather than by
    sampling, and normalized by the sum of example weights to remain in [0, 1].
    """

    if logits.shape != targets.shape:
        raise ValueError(
            f"Transition logits/targets mismatch: {tuple(logits.shape)} vs "
            f"{tuple(targets.shape)}"
        )
    if pos_weight <= 0:
        raise ValueError("Transition pos_weight must be positive")
    targets = targets.to(device=logits.device, dtype=logits.dtype)
    if not torch.isfinite(targets).all() or not torch.all(
        (targets == 0) | (targets == 1)
    ):
        raise ValueError("Transition targets must contain only finite 0/1 values")
    probability = logits.sigmoid()
    probability_correct = (
        targets * probability + (1.0 - targets) * (1.0 - probability)
    )
    weights = 1.0 + (pos_weight - 1.0) * targets
    return (weights * probability_correct).sum() / weights.sum()


def transition_aware_objective(
    activity_loss: torch.Tensor,
    transition_logits: torch.Tensor,
    transition_targets: torch.Tensor,
    *,
    objective: str,
    transition_weight: float = 0.2,
    pos_weight: float = 3.0,
    reward_weight: float = 0.1,
) -> dict[str, torch.Tensor]:
    """Compose the R1 or R2 transition-aware training objective."""

    if objective not in TRANSITION_OBJECTIVE_NAMES[1:]:
        raise ValueError(
            f"Transition objective must be R1 or R2, got {objective!r}"
        )
    if transition_weight < 0 or reward_weight < 0:
        raise ValueError("Transition and reward weights must be non-negative")
    transition_loss = weighted_transition_bce_loss(
        transition_logits,
        transition_targets,
        pos_weight=pos_weight,
    )
    expected_reward = expected_transition_reward(
        transition_logits,
        transition_targets,
        pos_weight=pos_weight,
    )
    applied_reward = (
        expected_reward
        if objective == "r2-expected-reward"
        else expected_reward.detach().new_zeros(())
    )
    total = (
        activity_loss
        + transition_weight * transition_loss
        - reward_weight * applied_reward
    )
    return {
        "total": total,
        "activity": activity_loss,
        "transition": transition_loss,
        "expected_reward": expected_reward,
        "applied_reward": applied_reward,
    }


def output_classes_for_loss(loss_name: str) -> int:
    """Return the number of head logits required by one training objective."""

    if loss_name not in LOSS_NAMES:
        raise ValueError(f"Unknown loss {loss_name!r}; expected one of {LOSS_NAMES}")
    return NUM_FOUR_STATES if loss_name == "four-state-ce" else NUM_CLASSES


def output_space_for_loss(loss_name: str) -> str:
    """Return a stable checkpoint description of the head output space."""

    output_classes_for_loss(loss_name)
    if loss_name == "four-state-ce":
        return "four-state-softmax"
    if loss_name == "independent-bce":
        return "three-label-sigmoid"
    return "three-class-softmax"


def head_output_names_for_loss(loss_name: str) -> tuple[str, ...]:
    output_classes_for_loss(loss_name)
    return FOUR_STATE_NAMES if loss_name == "four-state-ce" else CLASS_NAMES


def egospeak_multihorizon_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    reduction: Literal["mean", "sum", "none"] = "mean",
) -> torch.Tensor:
    """Reproduce the LSTR ``MCE`` loss used by the supplied shell/config.

    EasyCom targets may contain both active speaker classes as ``[0, 1, 1]``.
    LSTR preserves those weights directly and sums their negative log-softmax
    terms. It does not normalize the target row.
    """

    _validate_loss_inputs(logits, targets)
    targets = targets.to(device=logits.device, dtype=logits.dtype)
    loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
    if reduction == "mean":
        return loss.mean()
    if reduction == "sum":
        return loss.sum()
    if reduction == "none":
        return loss
    raise ValueError(f"Unknown reduction: {reduction}")


def miniroad_multihorizon_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    reduction: Literal["mean", "sum", "none"] = "mean",
) -> torch.Tensor:
    """Reproduce MiniROAD's L2-normalized multi-positive anticipation loss."""

    _validate_loss_inputs(logits, targets)
    targets = nn.functional.normalize(
        targets.to(device=logits.device, dtype=logits.dtype),
        p=2,
        dim=-1,
    )
    loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
    if reduction == "mean":
        return loss.mean()
    if reduction == "sum":
        return loss.sum()
    if reduction == "none":
        return loss
    raise ValueError(f"Unknown reduction: {reduction}")


def independent_bce_multihorizon_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    reduction: Literal["mean", "sum", "none"] = "mean",
) -> torch.Tensor:
    """Independent sigmoid/BCE loss for the three marginal activity labels."""

    _validate_loss_inputs(logits, targets)
    targets = targets.to(
        device=logits.device,
        dtype=logits.dtype,
    )
    element_loss = nn.functional.binary_cross_entropy_with_logits(
        logits,
        targets,
        reduction="none",
    )
    if reduction == "mean":
        return element_loss.mean()
    if reduction == "sum":
        return element_loss.sum()
    if reduction == "none":
        return element_loss
    raise ValueError(f"Unknown reduction: {reduction}")


def four_state_multihorizon_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    reduction: Literal["mean", "sum", "none"] = "mean",
) -> torch.Tensor:
    """Four-state CE for silence, wearer-only, others-only, and overlap."""

    if logits.ndim < 1 or logits.shape[-1] != NUM_FOUR_STATES:
        raise ValueError(
            f"Expected four-state logits [...,{NUM_FOUR_STATES}], got "
            f"{tuple(logits.shape)}"
        )
    if logits.shape[:-1] != targets.shape[:-1] or targets.shape[-1] != NUM_CLASSES:
        raise ValueError(
            "Four-state logits/activity target shape mismatch: "
            f"{tuple(logits.shape)} vs {tuple(targets.shape)}"
        )
    state_targets = (targets[..., 1] + 2 * targets[..., 2]).long().to(
        logits.device
    )
    loss = nn.functional.cross_entropy(
        logits.reshape(-1, NUM_FOUR_STATES),
        state_targets.reshape(-1),
        reduction="none",
    ).reshape(logits.shape[:-1])
    if reduction == "mean":
        return loss.mean()
    if reduction == "sum":
        return loss.sum()
    if reduction == "none":
        return loss
    raise ValueError(f"Unknown reduction: {reduction}")


def marginal_probabilities_from_logits(
    logits: torch.Tensor,
    loss_name: str,
) -> torch.Tensor:
    """Convert any supported head output to three EgoSpeak marginals.

    Four-state marginals deliberately are not renormalized: wearer and others
    may both be near one when the overlap state is likely.
    """

    expected_classes = output_classes_for_loss(loss_name)
    if logits.ndim < 1 or logits.shape[-1] != expected_classes:
        raise ValueError(
            f"Loss {loss_name!r} expects {expected_classes} logits, got "
            f"{tuple(logits.shape)}"
        )
    if loss_name == "independent-bce":
        return logits.sigmoid()
    probabilities = logits.softmax(dim=-1)
    if loss_name != "four-state-ce":
        return probabilities
    return torch.stack(
        (
            probabilities[..., 0],
            probabilities[..., 1] + probabilities[..., 3],
            probabilities[..., 2] + probabilities[..., 3],
        ),
        dim=-1,
    )


def multihorizon_cross_entropy(
    logits: torch.Tensor,
    targets: torch.Tensor,
) -> torch.Tensor:
    """Backward-compatible alias for the LSTR/EgoSpeak loss."""

    return egospeak_multihorizon_loss(logits, targets)


def normalize_session_name(value: str | int) -> str:
    text = str(value)
    if text.startswith("Session_"):
        suffix = text.removeprefix("Session_")
    else:
        suffix = text
    if not suffix.isdigit():
        raise ValueError(f"Invalid EasyCom session: {value}")
    return f"Session_{int(suffix)}"


def normalize_sample_name(
    value: str | int,
    dataset: DatasetName = "easycom",
) -> str:
    """Return the canonical EasyCom session name or Ego4D clip UUID."""

    if dataset == "easycom":
        return normalize_session_name(value)
    if dataset != "ego4d":
        raise ValueError(f"Unknown dataset: {dataset}")

    text = str(value).strip()
    try:
        return str(UUID(text))
    except (ValueError, AttributeError) as error:
        raise ValueError(f"Invalid Ego4D clip UUID: {value}") from error


def resolve_feature_path(
    feature_root: str | Path,
    sample: str | int,
    *,
    dataset: DatasetName = "easycom",
) -> Path:
    """Read flat features or Ego4D train/test subdirectories without link farms.

    Manifest membership, not directory discovery, determines the data split.
    Multiple copies are rejected rather than silently choosing stale features.
    """
    name = normalize_sample_name(sample, dataset)
    root = Path(feature_root)
    candidates = [root / f"{name}.pt"]
    if dataset == "ego4d":
        candidates += [root / split / f"{name}.pt" for split in ("train", "test")]
    found = [path for path in candidates if path.is_file()]
    if len(found) > 1:
        raise ValueError(f"Ambiguous feature copies for {name}: {found}; select one feature root")
    if not found:
        raise FileNotFoundError(f"No features for {name}; checked: " + ", ".join(map(str, candidates)))
    return found[0]


def load_ego4d_manifest_split(
    manifest: str | Path,
    split: Literal["train", "test", "all"],
) -> list[str]:
    """Load ordered UIDs from ``ego4d_egospeak.tsv``; all preserves row order."""

    if split not in {"train", "test", "all"}:
        raise ValueError(f"Invalid Ego4D split: {split}")
    path = Path(manifest)
    if not path.is_file():
        raise FileNotFoundError(path)

    selected: list[str] = []
    seen: set[str] = set()
    with path.open(newline="") as file:
        reader = csv.DictReader(file, delimiter="\t")
        columns = set(reader.fieldnames or [])
        required = {"split", "uid"}
        if not required.issubset(columns):
            raise ValueError(
                f"Manifest must contain {sorted(required)}, got "
                f"{reader.fieldnames}: {path}"
            )
        for row_number, row in enumerate(reader, start=2):
            row_split = row["split"].strip()
            if row_split not in {"train", "test"}:
                raise ValueError(
                    f"Invalid split {row_split!r} at {path}:{row_number}"
                )
            uid = normalize_sample_name(row["uid"], "ego4d")
            if uid in seen:
                raise ValueError(
                    f"Duplicate Ego4D UID {uid} at {path}:{row_number}"
                )
            seen.add(uid)
            if split == "all" or row_split == split:
                selected.append(uid)

    if not selected:
        raise ValueError(f"No {split!r} samples found in {path}")
    return selected


@dataclass
class SessionData:
    name: str
    states: torch.Tensor
    targets: torch.Tensor
    label_indices: torch.LongTensor
    feature_metadata: dict[str, Any]
    local_features: torch.Tensor | None = None
    local_person_features: "FramePersonFeatures | None" = None

    @property
    def hidden_size(self) -> int:
        return int(self.states.shape[-1])

    @property
    def num_input_layers(self) -> int:
        return 1 if self.states.ndim == 2 else int(self.states.shape[-2])

    def future_examples(
        self,
        num_horizons: int = NUM_HORIZONS,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
        """Return states, future labels, and target indices for valid rows."""

        offsets = torch.arange(1, num_horizons + 1, dtype=torch.long)
        future_indices = self.label_indices[:, None] + offsets[None, :]
        valid = (future_indices >= 0).all(dim=1) & (
            future_indices < len(self.targets)
        ).all(dim=1)
        rows = future_indices[valid]
        return self.states[valid], self.targets[rows], rows


def load_session_data(
    feature_root: str | Path,
    target_root: str | Path,
    session: str | int,
    *,
    expected_mode: str | None = "causal-prefix",
    dataset: DatasetName = "easycom",
    local_feature_root: str | Path | None = None,
    local_modalities: str = "av",
    local_contract: dict[str, Any] | None = None,
    local_fusion: str = "pooled",
    person_track_max_gap_seconds: float = 1.0,
) -> SessionData:
    """Load and validate one saved Omni feature payload and target array."""

    name = normalize_sample_name(session, dataset)
    feature_path = resolve_feature_path(feature_root, name, dataset=dataset)
    target_path = Path(target_root) / f"{name}.npy"
    if not feature_path.is_file():
        raise FileNotFoundError(feature_path)
    if not target_path.is_file():
        raise FileNotFoundError(target_path)

    payload = torch.load(feature_path, map_location="cpu", weights_only=False)
    if not isinstance(payload, dict) or "states" not in payload:
        raise ValueError(f"Invalid feature payload: {feature_path}")
    states = payload["states"]
    if not isinstance(states, torch.Tensor) or states.ndim not in {2, 3}:
        raise ValueError(f"states must be [T,D] or [T,L,D] in {feature_path}")
    if states.ndim == 3 and states.shape[1] < 2:
        raise ValueError(
            f"Stacked states must contain at least two layers: {feature_path}"
        )
    if not torch.is_floating_point(states):
        raise ValueError(f"states must be floating point in {feature_path}")
    if not torch.isfinite(states).all():
        raise ValueError(f"states contain NaN/Inf in {feature_path}")

    mode = payload.get("mode")
    if expected_mode is not None and mode != expected_mode:
        raise ValueError(
            f"{feature_path} mode is {mode!r}, expected {expected_mode!r}"
        )
    label_indices = payload.get("label_indices")
    if label_indices is None:
        raise ValueError(f"Missing label_indices in {feature_path}")
    label_indices = torch.as_tensor(label_indices, dtype=torch.long)
    if label_indices.ndim != 1 or len(label_indices) != len(states):
        raise ValueError(f"Invalid label_indices in {feature_path}")
    if len(label_indices) and not torch.equal(
        label_indices,
        torch.arange(len(label_indices), dtype=torch.long),
    ):
        raise ValueError(
            f"{feature_path} does not use the verified zero-based dense alignment"
        )

    target_array = np.load(target_path, allow_pickle=False)
    if target_array.ndim != 2 or target_array.shape[1] != NUM_CLASSES:
        raise ValueError(
            f"Expected target [T,{NUM_CLASSES}], got {target_array.shape}: "
            f"{target_path}"
        )
    targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
    try:
        _validated_binary_activity_targets(
            targets,
            allow_noncanonical_binary=dataset == "ego4d",
        )
    except ValueError as error:
        raise ValueError(f"{error}: {target_path}") from error
    if len(label_indices) and int(label_indices[-1]) >= len(targets):
        raise ValueError(
            f"Last label index {int(label_indices[-1])} exceeds target length "
            f"{len(targets)} in {name}"
        )
    if len(states) != len(targets):
        raise ValueError(
            f"{name}: feature/target length mismatch: {len(states)} vs "
            f"{len(targets)}"
        )

    metadata = {
        key: payload.get(key)
        for key in (
            "model_id",
            "mode",
            "config",
            "start",
            "end",
            "num_steps",
            "video_sampling_grid",
        )
    }
    expected_background = targets[:, 1:].sum(dim=-1) == 0
    canonical_rows = targets[:, 0].bool() == expected_background
    unique_rows, row_counts = torch.unique(
        targets.to(dtype=torch.int64),
        dim=0,
        return_counts=True,
    )
    metadata["num_noncanonical_target_rows"] = int(
        (~canonical_rows).sum().item()
    )
    metadata["target_row_pattern_counts"] = {
        "".join(str(int(value)) for value in row.tolist()): int(count)
        for row, count in zip(unique_rows, row_counts)
    }
    if local_fusion not in LOCAL_FUSION_MODES:
        raise ValueError(f"Unknown local fusion: {local_fusion}")
    local = None
    local_people = None
    if local_feature_root is not None:
        local_path = resolve_feature_path(
            local_feature_root,
            name,
            dataset=dataset,
        )
        if local_fusion in TRACK_PERSON_FUSION_MODES:
            if local_modalities != "av":
                raise ValueError(
                    "Track person fusion requires --local-modalities av"
                )
            from local_turn_features import load_track_features

            local, local_people, contract = load_track_features(
                local_path,
                payload,
                sample_id=name,
                dataset=dataset,
                track_max_gap_seconds=person_track_max_gap_seconds,
                expected_contract=local_contract,
            )
        elif local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
            if local_modalities != "av":
                raise ValueError(
                    "Person-set fusion requires --local-modalities av"
                )
            from local_turn_features import load_frame_set_features

            local, local_people, contract = load_frame_set_features(
                local_path,
                payload,
                sample_id=name,
                dataset=dataset,
                expected_contract=local_contract,
                oracle_analysis=local_fusion == "oracle-frame-set-attention",
            )
        else:
            from local_turn_features import load_local_features

            local, contract = load_local_features(
                local_path,
                payload,
                sample_id=name,
                dataset=dataset,
                modalities=local_modalities,
                expected_contract=local_contract,
            )
        metadata["local_contract"] = contract
    elif local_contract is not None:
        raise ValueError("--local-feature-root is required for this checkpoint")
    elif local_fusion != "pooled":
        raise ValueError("Person-set fusion requires --local-feature-root")
    return SessionData(
        name,
        states,
        targets,
        label_indices,
        metadata,
        local,
        local_people,
    )


def average_precision_metrics(
    scores: np.ndarray,
    targets: np.ndarray,
) -> dict[str, Any]:
    """Compute EgoSpeak frame AP for all three classes, including background."""

    from sklearn.metrics import average_precision_score

    scores = np.asarray(scores, dtype=np.float64)
    targets = np.asarray(targets, dtype=np.float64)
    if scores.shape != targets.shape or scores.ndim != 2:
        raise ValueError(
            f"scores/targets must have equal [N,C] shape, got "
            f"{scores.shape} and {targets.shape}"
        )
    if scores.shape[1] != NUM_CLASSES:
        raise ValueError(f"Expected {NUM_CLASSES} classes, got {scores.shape[1]}")

    aps = average_precision_score(targets, scores, average=None)
    return {
        "num_steps": int(len(scores)),
        "class_ap": {
            name: float(value) for name, value in zip(CLASS_NAMES, aps)
        },
        "map": float(np.mean(aps)),
    }


def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
    """Return a validated head config, including legacy-checkpoint defaults."""

    config = dict(payload.get("head_config", {}))
    config.setdefault("hidden_size", 2048)
    config.setdefault("num_horizons", NUM_HORIZONS)
    config.setdefault("num_classes", NUM_CLASSES)
    config.setdefault("dropout", 0.1)
    config.setdefault("num_input_layers", 1)
    config.setdefault("learnable_layer_mixture", False)
    config.setdefault("transition_auxiliary", False)
    config.setdefault("local_visual_dim", 0)
    config.setdefault("local_audio_dim", 0)
    config.setdefault("local_projection_dim", 256)
    config.setdefault("local_fusion", "pooled")
    config.setdefault("person_attention_dim", 256)
    config.setdefault("person_attention_heads", 4)
    config.setdefault("person_history_frames", 16)
    config.setdefault("person_track_max_gap_seconds", 1.0)
    config.setdefault(
        "layer_fusion",
        (
            "scalar-mixture"
            if config["learnable_layer_mixture"]
            else "single"
        ),
    )
    return config


def checkpoint_loss_name(payload: dict[str, Any]) -> str:
    """Read the objective from a checkpoint, defaulting legacy files to MCE."""

    training_config = payload.get("training_config", {})
    if not isinstance(training_config, dict):
        raise ValueError("Checkpoint training_config must be a mapping")
    loss_name = str(training_config.get("loss", "lstr-mce"))
    expected_classes = output_classes_for_loss(loss_name)
    actual_classes = int(checkpoint_head_config(payload)["num_classes"])
    if actual_classes != expected_classes:
        raise ValueError(
            f"Checkpoint loss {loss_name!r} requires {expected_classes} head "
            f"outputs, but head_config has {actual_classes}"
        )
    recorded_space = payload.get("output_space")
    expected_space = output_space_for_loss(loss_name)
    if recorded_space is not None and recorded_space != expected_space:
        raise ValueError(
            f"Checkpoint output_space is {recorded_space!r}, expected "
            f"{expected_space!r} for loss {loss_name!r}"
        )
    return loss_name


def validate_disjoint_sessions(
    first: Sequence[str],
    second: Sequence[str],
) -> None:
    overlap = set(first) & set(second)
    if overlap:
        raise ValueError(f"Session sets overlap: {sorted(overlap)}")