"""Shared head, labels, and metrics for the EasyCom EgoSpeak experiment."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, Sequence

import numpy as np
import torch
from torch import nn


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",
)
STEP_SECONDS = 0.2
NUM_HORIZONS = 10
NUM_CLASSES = 3


class TurnPredictionHead(nn.Module):
    """Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits."""

    def __init__(
        self,
        hidden_size: int = 2048,
        num_horizons: int = NUM_HORIZONS,
        num_classes: int = NUM_CLASSES,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.num_horizons = num_horizons
        self.num_classes = num_classes
        self.dropout = dropout
        self.net = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, num_horizons * num_classes),
        )

    def forward(self, states: torch.Tensor) -> torch.Tensor:
        logits = self.net(states)
        return logits.view(
            *states.shape[:-1], self.num_horizons, self.num_classes
        )


def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> 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_CLASSES:
        raise ValueError(
            f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}"
        )


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


@dataclass
class SessionData:
    name: str
    states: torch.Tensor
    targets: torch.Tensor
    label_indices: torch.LongTensor
    feature_metadata: dict[str, Any]

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

    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",
) -> SessionData:
    """Load and validate one saved Omni feature payload and target array."""

    name = normalize_session_name(session)
    feature_path = Path(feature_root) / f"{name}.pt"
    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 != 2:
        raise ValueError(f"states must be [T,D] in {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))
    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",
        )
    }
    return SessionData(name, states, targets, label_indices, metadata)


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)
    return config


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