"""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",
)
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",
)
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 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,
    ) -> 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,
    *,
    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) -> torch.Tensor:
    """Validate the four legal EasyCom/EgoSpeak three-label activity rows."""

    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).to(rounded.dtype)
    if not torch.equal(rounded[..., 0], expected_background):
        raise ValueError(
            "Activity targets must be one of [1,0,0], [0,1,0], "
            "[0,0,1], or [0,1,1]"
        )
    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 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 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))
    _validated_binary_activity_targets(targets)
    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 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)}")