"""Validated, causal local-feature sidecars for the pooled B1 baseline.

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

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


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_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."""
    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}")
    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 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