"""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 Any, Literal, Sequence
from uuid import UUID

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
DATASET_NAMES = ("easycom", "ego4d")
LAYER_FUSION_MODES = (
    "single",
    "fixed-mean",
    "scalar-mixture",
    "concat-mlp",
)
DatasetName = Literal["easycom", "ego4d"]


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,
        num_input_layers: int = 1,
        learnable_layer_mixture: bool = False,
        layer_fusion: str | None = None,
    ) -> 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"
        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.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:
        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)
        logits = self.net(states)
        return logits.view(
            *states.shape[:-1], self.num_horizons, self.num_classes
        )

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


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 load_ego4d_manifest_split(
    manifest: str | Path,
    split: Literal["train", "test"],
) -> list[str]:
    """Load one ordered UID split from ``ego4d_egospeak.tsv``."""

    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 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]

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

    name = normalize_sample_name(session, dataset)
    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 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))
    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)
    config.setdefault("num_input_layers", 1)
    config.setdefault("learnable_layer_mixture", False)
    config.setdefault(
        "layer_fusion",
        (
            "scalar-mixture"
            if config["learnable_layer_mixture"]
            else "single"
        ),
    )
    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)}")