"""Validated, causal pooled and frame-set local-feature sidecars.

No model/decoder imports: these utilities are also used by CPU head training.
Sidecars store encoder-native vectors, never a randomly initialized projection.
"""

from __future__ import annotations

from dataclasses import dataclass
import math
from pathlib import Path
from typing import Any

import torch

LOCAL_FORMAT = "qwen-local-pooled-v1"
MODALITIES = {"visual": ("visual",), "audio": ("audio",), "av": ("visual", "audio")}


@dataclass(frozen=True)
class FramePersonFeatures:
    """Ragged person sets indexed by aligned 5 Hz rows without frame duplication."""

    features: torch.Tensor
    frame_offsets: torch.LongTensor
    row_to_frame: torch.LongTensor
    chosen_frame_times: torch.Tensor
    previous_observation: torch.LongTensor | None = None

    def __post_init__(self) -> None:
        if (
            not isinstance(self.features, torch.Tensor)
            or self.features.ndim != 2
            or not self.features.is_floating_point()
            or not torch.isfinite(self.features).all()
            or self.features.shape[1] <= 0
        ):
            raise ValueError("Person features must be finite floating [N,D]")
        if (
            not isinstance(self.frame_offsets, torch.Tensor)
            or self.frame_offsets.dtype != torch.long
            or self.frame_offsets.ndim != 1
            or not len(self.frame_offsets)
            or int(self.frame_offsets[0]) != 0
            or int(self.frame_offsets[-1]) != len(self.features)
            or bool((self.frame_offsets[1:] < self.frame_offsets[:-1]).any())
        ):
            raise ValueError("Invalid person frame_offsets")
        if (
            not isinstance(self.row_to_frame, torch.Tensor)
            or self.row_to_frame.dtype != torch.long
            or self.row_to_frame.ndim != 1
        ):
            raise ValueError("row_to_frame must be int64 [T]")
        num_frames = len(self.frame_offsets) - 1
        if bool(((self.row_to_frame < -1) | (self.row_to_frame >= num_frames)).any()):
            raise ValueError("row_to_frame contains an invalid frame index")
        if (
            not isinstance(self.chosen_frame_times, torch.Tensor)
            or self.chosen_frame_times.dtype != torch.float64
            or self.chosen_frame_times.shape != self.row_to_frame.shape
        ):
            raise ValueError("chosen_frame_times must be float64 [T]")
        missing = self.row_to_frame < 0
        if not torch.equal(torch.isnan(self.chosen_frame_times), missing):
            raise ValueError("Only rows without an available frame may have NaN frame time")
        if not torch.isfinite(self.chosen_frame_times[~missing]).all():
            raise ValueError("Chosen person frame times must be finite")
        if self.previous_observation is not None:
            previous = self.previous_observation
            if (
                not isinstance(previous, torch.Tensor)
                or previous.dtype != torch.long
                or previous.shape != (len(self.features),)
            ):
                raise ValueError("previous_observation must be int64 [N]")
            indices = torch.arange(len(previous), dtype=torch.long)
            if bool(((previous < -1) | (previous >= indices)).any()):
                raise ValueError(
                    "Each previous observation must be -1 or point strictly backward"
                )

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

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

    @property
    def max_people(self) -> int:
        counts = self.frame_offsets[1:] - self.frame_offsets[:-1]
        return int(counts.max()) if len(counts) else 0

    def select_rows(self, indices: torch.Tensor) -> "FramePersonFeatures":
        indices = torch.as_tensor(indices)
        if indices.dtype != torch.long or indices.ndim != 1:
            raise ValueError("Person row selection must be int64 [N]")
        if len(indices) and (int(indices.min()) < 0 or int(indices.max()) >= len(self)):
            raise IndexError("Person row selection is out of range")
        return FramePersonFeatures(
            self.features,
            self.frame_offsets,
            self.row_to_frame[indices],
            self.chosen_frame_times[indices],
            self.previous_observation,
        )

    def _padded_observation_indices(
        self,
        row_indices: torch.Tensor | None = None,
    ) -> tuple[torch.LongTensor, torch.BoolTensor]:
        """Return current-frame observation indices and a real-person mask."""

        if row_indices is None:
            frames = self.row_to_frame
        else:
            row_indices = torch.as_tensor(row_indices)
            if row_indices.dtype != torch.long or row_indices.ndim != 1:
                raise ValueError("Batch row indices must be int64 [B]")
            if len(row_indices) and (
                int(row_indices.min()) < 0 or int(row_indices.max()) >= len(self)
            ):
                raise IndexError("Batch person row index is out of range")
            frames = self.row_to_frame[row_indices]
        valid_frames = frames >= 0
        safe_frames = frames.clamp_min(0)
        starts = torch.zeros_like(frames)
        counts = torch.zeros_like(frames)
        if bool(valid_frames.any()):
            starts[valid_frames] = self.frame_offsets[safe_frames[valid_frames]]
            counts[valid_frames] = (
                self.frame_offsets[safe_frames[valid_frames] + 1]
                - starts[valid_frames]
            )
        width = int(counts.max()) if len(counts) else 0
        mask = torch.arange(width)[None, :] < counts[:, None]
        if width == 0:
            return torch.zeros((len(frames), 0), dtype=torch.long), mask
        indices = starts[:, None] + torch.arange(width)[None, :]
        return indices.masked_fill(~mask, 0), mask

    def padded(
        self,
        row_indices: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.BoolTensor]:
        """Materialize only one batch as [B,P,D] plus a True-for-real mask."""

        safe_indices, mask = self._padded_observation_indices(row_indices)
        if safe_indices.shape[1] == 0:
            return self.features.new_zeros((len(safe_indices), 0, self.dimension)), mask
        values = self.features[safe_indices]
        return values.masked_fill(~mask[..., None], 0), mask

    def padded_history(
        self,
        row_indices: torch.Tensor | None = None,
        *,
        max_observations: int,
    ) -> tuple[torch.Tensor, torch.BoolTensor]:
        """Return causal per-current-person histories as ``[B,P,L,D]``.

        Histories are chronological and left aligned. Track links are built
        from provided IDs with an explicit causal gap/reset policy when the
        sidecar is loaded; no identity value is exposed to the model.
        """

        if max_observations <= 0:
            raise ValueError("max_observations must be positive")
        if self.previous_observation is None:
            raise ValueError("Track histories were not constructed for this person set")
        current, current_mask = self._padded_observation_indices(row_indices)
        batch, people = current.shape
        if people == 0:
            return (
                self.features.new_zeros(
                    (batch, 0, max_observations, self.dimension)
                ),
                torch.zeros(
                    (batch, 0, max_observations), dtype=torch.bool
                ),
            )

        reverse_indices = torch.zeros(
            (batch, people, max_observations), dtype=torch.long
        )
        reverse_mask = torch.zeros(
            (batch, people, max_observations), dtype=torch.bool
        )
        cursor = current.clone()
        cursor_valid = current_mask.clone()
        for step in range(max_observations):
            reverse_indices[:, :, step] = cursor.masked_fill(~cursor_valid, 0)
            reverse_mask[:, :, step] = cursor_valid
            safe = cursor.masked_fill(~cursor_valid, 0)
            cursor = self.previous_observation[safe]
            cursor_valid &= cursor >= 0

        lengths = reverse_mask.sum(dim=-1)
        output_positions = torch.arange(max_observations)[None, None, :]
        source_positions = lengths[:, :, None] - 1 - output_positions
        history_mask = output_positions < lengths[:, :, None]
        source_positions = source_positions.clamp_min(0)
        chronological = reverse_indices.gather(dim=-1, index=source_positions)
        chronological = chronological.masked_fill(~history_mask, 0)
        histories = self.features[chronological]
        histories = histories.masked_fill(~history_mask[..., None], 0)
        return histories, history_mask

    @classmethod
    def concatenate(
        cls,
        parts: list["FramePersonFeatures"],
    ) -> "FramePersonFeatures":
        if not parts:
            raise ValueError("Cannot concatenate no person feature sets")
        dimension = parts[0].dimension
        if any(part.dimension != dimension for part in parts):
            raise ValueError("Person feature dimensions differ across sessions")
        feature_parts = []
        offset_parts = [torch.zeros(1, dtype=torch.long)]
        row_parts = []
        time_parts = []
        feature_base = 0
        frame_base = 0
        previous_parts = []
        has_track_links = [part.previous_observation is not None for part in parts]
        if any(has_track_links) and not all(has_track_links):
            raise ValueError("Track history links are missing for some sessions")
        for part in parts:
            feature_parts.append(part.features)
            offset_parts.append(part.frame_offsets[1:] + feature_base)
            rows = part.row_to_frame.clone()
            valid = rows >= 0
            rows[valid] += frame_base
            row_parts.append(rows)
            time_parts.append(part.chosen_frame_times)
            if part.previous_observation is not None:
                previous = part.previous_observation.clone()
                valid_previous = previous >= 0
                previous[valid_previous] += feature_base
                previous_parts.append(previous)
            feature_base += len(part.features)
            frame_base += len(part.frame_offsets) - 1
        return cls(
            torch.cat(feature_parts, dim=0),
            torch.cat(offset_parts, dim=0),
            torch.cat(row_parts, dim=0),
            torch.cat(time_parts, dim=0),
            torch.cat(previous_parts, dim=0) if previous_parts else None,
        )


def reference_alignment(payload: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]:
    """Require the same full, zero-origin 5 Hz timeline used by head training."""
    states = payload.get("states")
    if not isinstance(states, torch.Tensor) or states.ndim not in (2, 3) or not len(states):
        raise ValueError("Reference must contain nonempty [T,D] or [T,L,D] states")
    if float(payload.get("start", 0)) != 0:
        raise ValueError("Local features require a zero-origin reference (start=0)")
    indices = torch.as_tensor(payload.get("label_indices", []))
    if indices.dtype != torch.long or not torch.equal(indices, torch.arange(len(states))):
        raise ValueError("Reference label_indices must be dense int64 [0,...,T-1]")
    times = torch.as_tensor(payload.get("times", []), dtype=torch.float64)
    expected = (torch.arange(len(states), dtype=torch.float64) + 1) / 5
    # Existing feature files may have serialized times as float32.
    if times.shape != expected.shape or not torch.allclose(times, expected, atol=5e-4, rtol=0):
        raise ValueError("Reference times must be [0.2,0.4,...] seconds at 5 Hz")
    return times, indices


def _load_validated_payload(
    path: Path,
    reference: dict[str, Any],
    *,
    sample_id: str,
    dataset: str,
    modalities: str,
    expected_contract: dict[str, Any] | None,
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], dict[str, Any], dict[str, Any]]:
    if modalities not in MODALITIES:
        raise ValueError(f"Unknown local modalities: {modalities}")
    times, indices = reference_alignment(reference)
    payload = torch.load(path, map_location="cpu", weights_only=False)
    if not isinstance(payload, dict) or payload.get("format") != LOCAL_FORMAT:
        raise ValueError(f"Invalid local feature format: {path}")
    if payload.get("sample_id") != sample_id or payload.get("dataset") != dataset:
        raise ValueError(f"Local sample/dataset mismatch: {path}")
    if payload.get("partial", False):
        raise ValueError(f"Partial smoke-test features cannot be used for training: {path}")
    local_indices = torch.as_tensor(payload.get("label_indices", []))
    local_times = torch.as_tensor(payload.get("times", []), dtype=torch.float64)
    if local_indices.dtype != torch.long or not torch.equal(indices, local_indices):
        raise ValueError(f"Local label_indices mismatch: {path}")
    if local_times.shape != times.shape or not torch.allclose(local_times, times, atol=1e-7, rtol=0):
        raise ValueError(f"Local times mismatch: {path}")
    config = payload.get("config", {})
    if config.get("causal_policy") != "half-open-trailing-window" or config.get("uses_speech_labels") is not False:
        raise ValueError(f"Local features lack the causal/no-speech-label contract: {path}")
    contract = {"format": LOCAL_FORMAT, "modalities": modalities, "config": config}
    if expected_contract is not None and contract != expected_contract:
        raise ValueError(f"Local encoder/pooling contract differs across sessions/checkpoint: {path}")
    return times, indices, payload, config, contract


def load_local_features(
    path: Path,
    reference: dict[str, Any],
    *,
    sample_id: str,
    dataset: str,
    modalities: str,
    expected_contract: dict[str, Any] | None = None,
) -> tuple[torch.Tensor, dict[str, Any]]:
    """Pack each requested modality as [raw vector, validity bit]. Fail closed."""
    times, _, payload, config, contract = _load_validated_payload(
        path,
        reference,
        sample_id=sample_id,
        dataset=dataset,
        modalities=modalities,
        expected_contract=expected_contract,
    )
    chunks = []
    for name in MODALITIES[modalities]:
        values = payload.get(name)
        valid = payload.get(f"{name}_valid")
        dimension = config.get(f"{name}_dim")
        if not isinstance(dimension, int) or dimension <= 0:
            raise ValueError(f"Missing {name}_dim: {path}")
        if not isinstance(values, torch.Tensor) or values.shape != (len(times), dimension):
            raise ValueError(f"Invalid {name} shape: {path}")
        if not values.is_floating_point() or not torch.isfinite(values).all():
            raise ValueError(f"Non-finite/non-floating {name} features: {path}")
        if not isinstance(valid, torch.Tensor) or valid.dtype != torch.bool or valid.shape != (len(times),):
            raise ValueError(f"Invalid {name}_valid mask: {path}")
        if torch.count_nonzero(values[~valid]):
            raise ValueError(f"Invalid {name} rows must be stored as zeros: {path}")
        chunks.extend((values.float(), valid[:, None].float()))
    return torch.cat(chunks, dim=-1), contract


def load_frame_set_features(
    path: Path,
    reference: dict[str, Any],
    *,
    sample_id: str,
    dataset: str,
    expected_contract: dict[str, Any] | None = None,
    track_max_gap_seconds: float | None = None,
) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
    """Load direct audio and the latest causal per-frame person set per 5 Hz row."""

    times, _, payload, config, contract = _load_validated_payload(
        path,
        reference,
        sample_id=sample_id,
        dataset=dataset,
        modalities="av",
        expected_contract=expected_contract,
    )
    audio = payload.get("audio")
    audio_valid = payload.get("audio_valid")
    audio_dim = config.get("audio_dim")
    if not isinstance(audio_dim, int) or audio_dim <= 0:
        raise ValueError(f"Missing audio_dim: {path}")
    if not isinstance(audio, torch.Tensor) or audio.shape != (len(times), audio_dim):
        raise ValueError(f"Invalid audio shape: {path}")
    if not audio.is_floating_point() or not torch.isfinite(audio).all():
        raise ValueError(f"Non-finite/non-floating audio features: {path}")
    if (
        not isinstance(audio_valid, torch.Tensor)
        or audio_valid.dtype != torch.bool
        or audio_valid.shape != (len(times),)
    ):
        raise ValueError(f"Invalid audio_valid mask: {path}")
    if torch.count_nonzero(audio[~audio_valid]):
        raise ValueError(f"Invalid audio rows must be stored as zeros: {path}")

    visual_dim = config.get("visual_dim")
    persons = payload.get("persons")
    if not isinstance(visual_dim, int) or visual_dim <= 0:
        raise ValueError(f"Missing visual_dim: {path}")
    if not isinstance(persons, dict):
        raise ValueError(f"Missing ragged person features: {path}")
    features = persons.get("features")
    frame_times = persons.get("frame_times")
    frame_offsets = persons.get("frame_offsets")
    records = persons.get("records")
    if (
        not isinstance(features, torch.Tensor)
        or features.ndim != 2
        or features.shape[1] != visual_dim
        or not features.is_floating_point()
        or not torch.isfinite(features).all()
    ):
        raise ValueError(f"Invalid person feature shape/content: {path}")
    if (
        not isinstance(frame_times, torch.Tensor)
        or frame_times.ndim != 1
        or not frame_times.is_floating_point()
    ):
        raise ValueError(f"Invalid person frame_times: {path}")
    frame_times = frame_times.to(dtype=torch.float64)
    if (
        not torch.isfinite(frame_times).all()
        or bool((frame_times < 0).any())
        or bool((frame_times[1:] <= frame_times[:-1]).any())
    ):
        raise ValueError(f"Person frame_times must be finite, nonnegative and increasing: {path}")
    if (
        not isinstance(frame_offsets, torch.Tensor)
        or frame_offsets.dtype != torch.long
        or frame_offsets.shape != (len(frame_times) + 1,)
        or not len(frame_offsets)
        or int(frame_offsets[0]) != 0
        or int(frame_offsets[-1]) != len(features)
        or bool((frame_offsets[1:] < frame_offsets[:-1]).any())
    ):
        raise ValueError(f"Invalid person frame_offsets: {path}")
    if not isinstance(records, list) or len(records) != len(features):
        raise ValueError(f"Person records/features length mismatch: {path}")
    if len(frame_times) and float(frame_times[-1]) > float(times[-1]) + 1e-8:
        raise ValueError(f"Person frame timeline extends beyond the reference: {path}")

    # Match the sidecar's half-open policy: a frame stamped exactly at the
    # state cutoff is not yet part of that state's observed prefix.
    row_to_frame = torch.searchsorted(frame_times, times, right=False) - 1
    chosen = torch.full(times.shape, float("nan"), dtype=torch.float64)
    has_frame = row_to_frame >= 0
    chosen[has_frame] = frame_times[row_to_frame[has_frame]]
    if bool((chosen[has_frame] >= times[has_frame]).any()):
        raise ValueError(f"Person frame mapping selected future visual evidence: {path}")
    previous_observation = None
    if track_max_gap_seconds is not None:
        if not math.isfinite(track_max_gap_seconds) or track_max_gap_seconds <= 0:
            raise ValueError("track_max_gap_seconds must be finite and positive")
        previous_observation = _build_previous_track_links(
            records,
            frame_times,
            frame_offsets,
            max_gap_seconds=track_max_gap_seconds,
            path=path,
        )
    frame_sets = FramePersonFeatures(
        features,
        frame_offsets,
        row_to_frame.to(dtype=torch.long),
        chosen,
        previous_observation,
    )
    direct_audio = torch.cat(
        (audio.float(), audio_valid[:, None].float()),
        dim=-1,
    )
    return direct_audio, frame_sets, contract


def _build_previous_track_links(
    records: list[Any],
    frame_times: torch.Tensor,
    frame_offsets: torch.LongTensor,
    *,
    max_gap_seconds: float,
    path: Path,
) -> torch.LongTensor:
    """Link observations causally by provided track ID without exposing IDs.

    A source-clip boundary always resets a track. Missing/unknown IDs are
    isolated observations instead of being incorrectly merged into one track.
    Duplicate known IDs within one frame fail closed.
    """

    previous = torch.full((len(records),), -1, dtype=torch.long)
    last_by_track: dict[tuple[int, str], tuple[int, float]] = {}
    unknown_values = {"", "-1", "none", "null", "unknown"}
    for frame_index, time_value in enumerate(frame_times.tolist()):
        start = int(frame_offsets[frame_index])
        end = int(frame_offsets[frame_index + 1])
        seen: set[tuple[int, str]] = set()
        for observation_index in range(start, end):
            record = records[observation_index]
            if not isinstance(record, dict):
                raise ValueError(f"Invalid person record at {observation_index}: {path}")
            raw_id = record.get("id")
            text_id = "" if raw_id is None else str(raw_id).strip()
            source_clip_index = record.get("source_clip_index", 0)
            if isinstance(source_clip_index, bool) or not isinstance(
                source_clip_index, int
            ):
                raise ValueError(
                    f"Invalid source_clip_index at person record "
                    f"{observation_index}: {path}"
                )
            if text_id.lower() in unknown_values:
                # Never join multiple anonymous detections under a shared -1.
                continue
            key = (source_clip_index, text_id)
            if key in seen:
                raise ValueError(
                    f"Duplicate person/track ID {text_id!r} in frame "
                    f"{frame_index}: {path}"
                )
            seen.add(key)
            old = last_by_track.get(key)
            if old is not None:
                old_index, old_time = old
                gap = float(time_value) - old_time
                if gap <= 0:
                    raise ValueError(f"Non-causal person track order: {path}")
                if gap <= max_gap_seconds + 1e-8:
                    previous[observation_index] = old_index
            last_by_track[key] = (observation_index, float(time_value))
    return previous


def load_track_features(
    path: Path,
    reference: dict[str, Any],
    *,
    sample_id: str,
    dataset: str,
    track_max_gap_seconds: float,
    expected_contract: dict[str, Any] | None = None,
) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
    """Load the exact pooled AV baseline plus causal anonymous track links."""

    pooled, contract = load_local_features(
        path,
        reference,
        sample_id=sample_id,
        dataset=dataset,
        modalities="av",
        expected_contract=expected_contract,
    )
    _, people, second_contract = load_frame_set_features(
        path,
        reference,
        sample_id=sample_id,
        dataset=dataset,
        expected_contract=contract,
        track_max_gap_seconds=track_max_gap_seconds,
    )
    if second_contract != contract:
        raise AssertionError("Track and pooled local contracts unexpectedly differ")
    return pooled, people, contract


def local_head_dimensions(contract: dict[str, Any] | None) -> dict[str, int]:
    if contract is None:
        return {"local_visual_dim": 0, "local_audio_dim": 0}
    selected = MODALITIES[contract["modalities"]]
    return {
        f"local_{name}_dim": int(contract["config"][f"{name}_dim"]) if name in selected else 0
        for name in ("visual", "audio")
    }


def causal_frame_ticks(end: float, fps: float) -> range:
    """Global fixed lattice [0,end), excluding the frame exactly at end."""
    if end <= 0 or fps <= 0:
        raise ValueError("end and fps must be positive")
    return range(max(0, math.ceil(end * fps - 1e-8)))


def pooled_recent_visual(
    frame_times: torch.Tensor,
    frame_vectors: torch.Tensor,
    frame_valid: torch.Tensor,
    end: float,
    window_seconds: float,
) -> tuple[torch.Tensor, bool]:
    """Equal frame weighting after equal person weighting within each frame."""
    keep = ((frame_times >= max(0., end - window_seconds) - 1e-8)
            & (frame_times < end - 1e-8) & frame_valid)
    if not bool(keep.any()):
        return frame_vectors.new_zeros(frame_vectors.shape[-1]), False
    return frame_vectors[keep].float().mean(0), True


def easycom_boxes(rows: list[dict[str, Any]], kind: str) -> dict[int, list[dict[str, Any]]]:
    """Map EasyCom face (1-based) / head (0-based) records to decoded indices."""
    if kind not in ("face", "head"):
        raise ValueError(f"Unknown box kind: {kind}")
    result = {}
    for row in rows:
        frame = int(row["Frame_Number"]) - (1 if kind == "face" else 0)
        people = row.get("Participants", [])
        if isinstance(people, dict):
            people = [people]
        if frame < 0 or frame in result or not isinstance(people, list):
            raise ValueError("Invalid/duplicate EasyCom box frame")
        result[frame] = [
            {"id": str(p.get("Participant_ID", -1)),
             "xyxy": [float(p[k]) for k in ("x1", "y1", "x2", "y2")]}
            for p in people
        ]
    return result


def crop_box(
    xyxy: list[float], width: int, height: int, *,
    offset_x: int = 0, offset_y: int = 0, expansion: float = 1.,
) -> tuple[int, int, int, int] | None:
    """Transform to camera coordinates, expand, clip; reject degenerate boxes."""
    if len(xyxy) != 4 or not all(math.isfinite(v) for v in xyxy):
        return None
    x1, y1, x2, y2 = xyxy
    if x2 <= x1 or y2 <= y1 or expansion <= 0:
        return None
    cx, cy = (x1 + x2) / 2 - offset_x, (y1 + y2) / 2 - offset_y
    half_w, half_h = (x2 - x1) * expansion / 2, (y2 - y1) * expansion / 2
    box = (max(0, math.floor(cx - half_w)), max(0, math.floor(cy - half_h)),
           min(width, math.ceil(cx + half_w)), min(height, math.ceil(cy + half_h)))
    w, h = box[2] - box[0], box[3] - box[1]
    # Qwen's smart_resize rejects aspect ratios above 200 (e.g. a sliver of
    # an offscreen box). Such a crop is not a usable person observation.
    return box if min(w, h) >= 2 and max(w, h) / min(w, h) <= 200 else None