"""Frozen EgoSpeak MiniROADA plus a trainable Omni residual branch.

This module deliberately does not modify the vendored EgoSpeak checkout.  The
``MiniROADABase`` parameter names match the released implementation so its raw
state dict can be loaded strictly.  Stage-one training freezes that complete
base and learns only a zero-initialized residual over the ten anticipation
logits.
"""

from __future__ import annotations

import hashlib
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import Dataset

from easycom_head import NUM_CLASSES, NUM_HORIZONS, STEP_SECONDS

OmniLayerFusion = Literal["single", "last", "mean"]


def load_torch_payload(path: Path, *, weights_only: bool = False) -> Any:
    """Load a torch payload while retaining compatibility with older PyTorch."""

    try:
        return torch.load(path, map_location="cpu", weights_only=weights_only)
    except TypeError:  # pragma: no cover - only exercised by older PyTorch.
        return torch.load(path, map_location="cpu")


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def normalize_session_name(value: str | int) -> str:
    text = str(value).strip()
    suffix = text.removeprefix("Session_")
    if not suffix.isdigit():
        raise ValueError(f"Invalid EasyCom session: {value!r}")
    return f"Session_{int(suffix)}"


class MiniROADABase(nn.Module):
    """Clean-room parameter-compatible implementation of EgoSpeak MiniROADA."""

    def __init__(
        self,
        *,
        rgb_dim: int = 2048,
        audio_dim: int = 5120,
        embedding_dim: int = 2048,
        hidden_dim: int = 1024,
        num_layers: int = 1,
        num_horizons: int = NUM_HORIZONS,
        num_classes: int = NUM_CLASSES,
        dropout: float = 0.2,
    ) -> None:
        super().__init__()
        self.rgb_dim = int(rgb_dim)
        self.audio_dim = int(audio_dim)
        self.embedding_dim = int(embedding_dim)
        self.hidden_dim = int(hidden_dim)
        self.num_layers = int(num_layers)
        self.num_horizons = int(num_horizons)
        self.num_classes = int(num_classes)

        self.layer1 = nn.Sequential(
            nn.Linear(self.rgb_dim + self.audio_dim, self.embedding_dim),
            nn.LayerNorm(self.embedding_dim),
            nn.ReLU(),
            nn.Dropout(p=dropout),
        )
        self.relu = nn.ReLU()
        self.gru = nn.GRU(
            self.embedding_dim,
            self.hidden_dim,
            self.num_layers,
            batch_first=True,
        )
        self.f_classification = nn.Sequential(
            nn.Linear(self.hidden_dim, self.num_classes)
        )
        self.anticipation_layer = nn.Sequential(
            nn.Linear(
                self.hidden_dim,
                self.num_horizons * self.hidden_dim,
            )
        )

    def encode(
        self,
        rgb: torch.Tensor,
        audio: torch.Tensor,
        recurrent_state: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if rgb.ndim != 3 or rgb.shape[-1] != self.rgb_dim:
            raise ValueError(
                f"Expected RGB [B,S,{self.rgb_dim}], got {tuple(rgb.shape)}"
            )
        if (
            audio.ndim != 3
            or audio.shape[:-1] != rgb.shape[:-1]
            or audio.shape[-1] != self.audio_dim
        ):
            raise ValueError(
                f"Expected audio [B,S,{self.audio_dim}] aligned with RGB, got "
                f"{tuple(audio.shape)}"
            )
        embedded = self.layer1(torch.cat((rgb, audio), dim=-1))
        hidden, next_state = self.gru(embedded, recurrent_state)
        return self.relu(hidden), next_state

    def anticipation_logits(self, hidden: torch.Tensor) -> torch.Tensor:
        shape = hidden.shape[:-1]
        anticipated = self.anticipation_layer(self.relu(hidden)).view(
            *shape,
            self.num_horizons,
            self.hidden_dim,
        )
        return self.f_classification(self.relu(anticipated))

    def current_logits(self, hidden: torch.Tensor) -> torch.Tensor:
        return self.f_classification(self.relu(hidden))

    def forward(
        self, rgb: torch.Tensor, audio: torch.Tensor
    ) -> dict[str, torch.Tensor]:
        hidden, _ = self.encode(rgb, audio)
        return {
            "logits": self.current_logits(hidden),
            "anticipation_logits": self.anticipation_logits(hidden),
        }


def _unwrap_state_dict(payload: Any) -> dict[str, torch.Tensor]:
    if not isinstance(payload, dict):
        raise TypeError("EgoSpeak checkpoint must contain a state dict")
    for key in ("state_dict", "model_state_dict"):
        nested = payload.get(key)
        if isinstance(nested, dict):
            payload = nested
            break
    if not payload or not all(
        isinstance(value, torch.Tensor) for value in payload.values()
    ):
        raise ValueError("EgoSpeak checkpoint does not contain tensor parameters")
    state = dict(payload)
    if all(key.startswith("module.") for key in state):
        state = {key.removeprefix("module."): value for key, value in state.items()}
    return state


def load_frozen_miniroada(checkpoint: Path, **model_kwargs: Any) -> MiniROADABase:
    checkpoint = checkpoint.expanduser().resolve()
    if not checkpoint.is_file():
        raise FileNotFoundError(checkpoint)
    base = MiniROADABase(**model_kwargs)
    state = _unwrap_state_dict(load_torch_payload(checkpoint, weights_only=True))
    base.load_state_dict(state, strict=True)
    base.requires_grad_(False)
    base.eval()
    return base


class EgoSpeakOmniResidual(nn.Module):
    """Feature fusion whose zero residual exactly recovers frozen MiniROADA."""

    def __init__(
        self,
        base: MiniROADABase,
        *,
        omni_dim: int,
        fusion_dim: int = 256,
        tail_embedding_dim: int = 32,
        num_tail_bins: int = NUM_HORIZONS,
        dropout: float = 0.2,
    ) -> None:
        super().__init__()
        if fusion_dim <= 0 or tail_embedding_dim <= 0 or num_tail_bins <= 0:
            raise ValueError("Fusion and tail dimensions must be positive")
        self.base = base
        self.omni_dim = int(omni_dim)
        self.fusion_dim = int(fusion_dim)
        self.tail_embedding_dim = int(tail_embedding_dim)
        self.num_tail_bins = int(num_tail_bins)

        self.ego_projection = nn.Sequential(
            nn.LayerNorm(base.hidden_dim),
            nn.Linear(base.hidden_dim, fusion_dim),
            nn.GELU(),
        )
        self.omni_projection = nn.Sequential(
            nn.LayerNorm(omni_dim),
            nn.Linear(omni_dim, fusion_dim),
            nn.GELU(),
        )
        self.tail_embedding = nn.Embedding(num_tail_bins, tail_embedding_dim)
        self.fusion = nn.Sequential(
            nn.LayerNorm(3 * fusion_dim + tail_embedding_dim),
            nn.Linear(3 * fusion_dim + tail_embedding_dim, fusion_dim),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        self.residual_output = nn.Linear(
            fusion_dim,
            base.num_horizons * base.num_classes,
        )
        nn.init.zeros_(self.residual_output.weight)
        nn.init.zeros_(self.residual_output.bias)

        self.base.requires_grad_(False)
        self.base.eval()

    def train(self, mode: bool = True) -> EgoSpeakOmniResidual:
        super().train(mode)
        # The released base remains deterministic: in particular, its dropout
        # must never reactivate during stage-one residual training.
        self.base.eval()
        return self

    def trainable_parameters(self) -> list[nn.Parameter]:
        return [parameter for parameter in self.parameters() if parameter.requires_grad]

    def _residual_logits(
        self,
        ego_hidden: torch.Tensor,
        omni: torch.Tensor,
        tail_bin: torch.Tensor,
    ) -> torch.Tensor:
        if omni.shape != (*ego_hidden.shape[:-1], self.omni_dim):
            raise ValueError(
                f"Omni/hidden shape mismatch: {tuple(omni.shape)} versus "
                f"{tuple(ego_hidden.shape)}"
            )
        if tail_bin.shape != ego_hidden.shape[:-1]:
            raise ValueError(
                f"Tail-bin shape mismatch: {tuple(tail_bin.shape)} versus "
                f"{tuple(ego_hidden.shape[:-1])}"
            )
        if tail_bin.dtype != torch.long:
            raise ValueError("tail_bin must be int64")
        if bool(((tail_bin < 0) | (tail_bin >= self.num_tail_bins)).any()):
            raise ValueError("tail_bin lies outside the configured block phase")

        ego = self.ego_projection(ego_hidden)
        global_state = self.omni_projection(omni)
        fused = torch.cat(
            (
                ego,
                global_state,
                ego * global_state,
                self.tail_embedding(tail_bin),
            ),
            dim=-1,
        )
        residual = self.residual_output(self.fusion(fused))
        return residual.view(
            *ego_hidden.shape[:-1],
            self.base.num_horizons,
            self.base.num_classes,
        )

    def forward_last(
        self,
        rgb: torch.Tensor,
        audio: torch.Tensor,
        omni: torch.Tensor,
        tail_bin: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        with torch.no_grad():
            hidden, _ = self.base.encode(rgb, audio)
            ego_hidden = hidden[:, -1, :]
            base_logits = self.base.anticipation_logits(ego_hidden)
        residual = self._residual_logits(ego_hidden.detach(), omni, tail_bin)
        return base_logits + residual, base_logits

    def forward_sequence_chunk(
        self,
        rgb: torch.Tensor,
        audio: torch.Tensor,
        omni: torch.Tensor,
        tail_bin: torch.Tensor,
        recurrent_state: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        with torch.no_grad():
            hidden, next_state = self.base.encode(rgb, audio, recurrent_state)
            base_logits = self.base.anticipation_logits(hidden)
        residual = self._residual_logits(hidden.detach(), omni, tail_bin)
        return base_logits + residual, base_logits, next_state


def collapse_omni_states(
    states: torch.Tensor,
    mode: OmniLayerFusion,
) -> torch.Tensor:
    if not isinstance(states, torch.Tensor) or states.ndim not in (2, 3):
        raise ValueError("Omni states must be [T,D] or [T,L,D]")
    if not torch.is_floating_point(states) or not torch.isfinite(states).all():
        raise ValueError("Omni states must be finite floating-point values")
    if states.ndim == 2:
        if mode not in ("single", "last", "mean"):
            raise ValueError(f"Unsupported Omni layer fusion: {mode}")
        return states.float().contiguous()
    if mode == "single":
        raise ValueError("Stacked Omni states require --omni-layer-fusion last or mean")
    if mode == "last":
        return states[:, -1, :].float().contiguous()
    if mode == "mean":
        return states.float().mean(dim=1).contiguous()
    raise ValueError(f"Unsupported Omni layer fusion: {mode}")


@dataclass
class EasyComSessionArrays:
    name: str
    rgb: np.ndarray
    audio: np.ndarray
    targets: np.ndarray
    omni: torch.Tensor
    tail_bins: torch.LongTensor
    omni_metadata: dict[str, Any]

    @property
    def num_steps(self) -> int:
        return len(self.targets)


def _require_feature_array(path: Path, expected_dim: int) -> np.ndarray:
    if not path.is_file():
        raise FileNotFoundError(path)
    array = np.load(path, mmap_mode="r", allow_pickle=False)
    if array.ndim != 2 or array.shape[1] != expected_dim:
        raise ValueError(f"Expected [T,{expected_dim}], got {array.shape}: {path}")
    if not np.issubdtype(array.dtype, np.floating):
        raise ValueError(f"Expected floating-point features: {path}")
    return array


def _require_target_array(path: Path, expected_dim: int) -> np.ndarray:
    if not path.is_file():
        raise FileNotFoundError(path)
    array = np.load(path, mmap_mode="r", allow_pickle=False)
    if array.ndim != 2 or array.shape[1] != expected_dim:
        raise ValueError(f"Expected [T,{expected_dim}], got {array.shape}: {path}")
    if not (
        np.issubdtype(array.dtype, np.number) or np.issubdtype(array.dtype, np.bool_)
    ):
        raise ValueError(f"Expected numeric activity targets: {path}")
    return array


def load_easycom_session_arrays(
    egospeak_feature_root: Path,
    omni_feature_root: Path,
    session: str | int,
    *,
    expected_omni_mode: str | None = "open-stream-replay",
    omni_layer_fusion: OmniLayerFusion = "single",
    rgb_dim: int = 2048,
    audio_dim: int = 5120,
    num_classes: int = NUM_CLASSES,
) -> EasyComSessionArrays:
    name = normalize_session_name(session)
    ego_root = egospeak_feature_root.expanduser().resolve()
    omni_root = omni_feature_root.expanduser().resolve()
    rgb = _require_feature_array(
        ego_root / "rgb_kinetics_resnet50" / f"{name}.npy",
        rgb_dim,
    )
    audio = _require_feature_array(
        ego_root / "audio_wav2vec2" / f"{name}.npy",
        audio_dim,
    )
    targets = _require_target_array(
        ego_root / "target_perframe" / f"{name}.npy",
        num_classes,
    )
    lengths = {len(rgb), len(audio), len(targets)}
    if len(lengths) != 1:
        raise ValueError(
            f"{name}: unequal EgoSpeak RGB/audio/target lengths "
            f"{len(rgb)}/{len(audio)}/{len(targets)}"
        )
    unique_targets = np.unique(np.asarray(targets))
    if not np.isin(unique_targets, (0, 1)).all():
        raise ValueError(
            f"{name}: target_perframe contains values outside binary activity: "
            f"{unique_targets.tolist()}"
        )

    omni_path = omni_root / f"{name}.pt"
    if not omni_path.is_file():
        raise FileNotFoundError(omni_path)
    payload = load_torch_payload(omni_path)
    if not isinstance(payload, dict):
        raise TypeError(f"Invalid Omni payload: {omni_path}")
    mode = payload.get("mode")
    if expected_omni_mode is not None and mode != expected_omni_mode:
        raise ValueError(
            f"{name}: Omni mode is {mode!r}, expected {expected_omni_mode!r}"
        )
    omni = collapse_omni_states(payload.get("states"), omni_layer_fusion)
    if len(omni) != len(targets):
        raise ValueError(
            f"{name}: Omni/target length mismatch {len(omni)} vs {len(targets)}"
        )
    if payload.get("label_indices") is None:
        raise ValueError(f"{name}: Omni payload lacks label_indices")
    label_indices = torch.as_tensor(payload["label_indices"], dtype=torch.long)
    expected_indices = torch.arange(len(targets), dtype=torch.long)
    if label_indices.ndim != 1 or not torch.equal(label_indices, expected_indices):
        raise ValueError(f"{name}: Omni label_indices are not dense zero-origin rows")

    config = payload.get("config")
    if not isinstance(config, dict):
        raise TypeError(f"{name}: Omni payload lacks extraction config")
    step_seconds = float(config.get("step_seconds", float("nan")))
    native_block_seconds = float(config.get("native_block_seconds", float("nan")))
    if not np.isclose(step_seconds, STEP_SECONDS, rtol=0.0, atol=1e-9):
        raise ValueError(
            f"{name}: Omni step_seconds={step_seconds}, expected {STEP_SECONDS}"
        )
    block_steps_float = native_block_seconds / step_seconds
    block_steps = round(block_steps_float)
    if block_steps <= 0 or not np.isclose(
        block_steps_float, block_steps, rtol=0.0, atol=1e-9
    ):
        raise ValueError(
            f"{name}: native block {native_block_seconds}s is not an integer "
            f"multiple of {step_seconds}s"
        )
    if block_steps != NUM_HORIZONS:
        raise ValueError(
            f"{name}: expected ten 0.2-s tail bins in a 2-s block, got {block_steps}"
        )

    times = torch.as_tensor(payload.get("times"), dtype=torch.float64)
    start = float(payload.get("start", 0.0))
    expected_times = (
        start + (torch.arange(len(targets), dtype=torch.float64) + 1) * step_seconds
    )
    if times.ndim != 1 or not torch.allclose(
        times, expected_times, rtol=0.0, atol=1e-6
    ):
        raise ValueError(f"{name}: Omni times do not match the dense 5 Hz grid")
    if bool(payload.get("partial", False)):
        raise ValueError(f"{name}: partial Omni smoke feature cannot train the model")

    tail_bins = torch.remainder(expected_indices, block_steps)
    metadata = {
        "path": str(omni_path),
        "mode": mode,
        "model_id": payload.get("model_id"),
        "config": config,
        "video_sampling_grid": payload.get("video_sampling_grid"),
        "media_interval_policy": payload.get("media_interval_policy"),
        "open_stream_policy": payload.get("open_stream_policy"),
        "state_shape_before_layer_fusion": list(payload["states"].shape),
        "state_shape": list(omni.shape),
    }
    return EasyComSessionArrays(
        name=name,
        rgb=rgb,
        audio=audio,
        targets=targets,
        omni=omni,
        tail_bins=tail_bins,
        omni_metadata=metadata,
    )


class EgoSpeakOmniWindowDataset(Dataset[tuple[torch.Tensor, ...]]):
    """MiniROADA windows plus the aligned current Omni state."""

    def __init__(
        self,
        sessions: Sequence[EasyComSessionArrays],
        *,
        window_size: int = 128,
        stride: int = 4,
        num_horizons: int = NUM_HORIZONS,
        seed: int = 0,
        max_examples: int | None = None,
    ) -> None:
        if not sessions:
            raise ValueError("At least one training session is required")
        if window_size <= 0 or stride <= 0 or num_horizons <= 0:
            raise ValueError("Window, stride, and horizon counts must be positive")
        if max_examples is not None and max_examples <= 0:
            raise ValueError("max_examples must be positive")
        self.sessions = list(sessions)
        self.window_size = int(window_size)
        self.stride = int(stride)
        self.num_horizons = int(num_horizons)
        self.seed = int(seed)
        self.max_examples = max_examples
        self.indices: list[tuple[int, int]] = []
        self.resample(0)

    def resample(self, epoch: int) -> None:
        rng = np.random.default_rng(self.seed + int(epoch))
        indices: list[tuple[int, int]] = []
        for session_index, session in enumerate(self.sessions):
            last_source = session.num_steps - self.num_horizons - 1
            if last_source < 0:
                raise ValueError(
                    f"{session.name}: only {session.num_steps} rows for "
                    f"{self.num_horizons} horizons"
                )
            offset = int(rng.integers(self.stride))
            indices.extend(
                (session_index, source_index)
                for source_index in range(offset, last_source + 1, self.stride)
            )
        if self.max_examples is not None:
            rng.shuffle(indices)
            indices = indices[: self.max_examples]
        if not indices:
            raise ValueError("Training resampling produced no examples")
        self.indices = indices

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

    def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
        session_index, source_index = self.indices[index]
        session = self.sessions[session_index]
        first = max(0, source_index - self.window_size + 1)
        length = source_index - first + 1
        rgb = torch.zeros(
            self.window_size,
            session.rgb.shape[1],
            dtype=torch.float32,
        )
        audio = torch.zeros(
            self.window_size,
            session.audio.shape[1],
            dtype=torch.float32,
        )
        rgb[-length:] = torch.from_numpy(
            np.asarray(session.rgb[first : source_index + 1], dtype=np.float32).copy()
        )
        audio[-length:] = torch.from_numpy(
            np.asarray(session.audio[first : source_index + 1], dtype=np.float32).copy()
        )
        future = torch.from_numpy(
            np.asarray(
                session.targets[
                    source_index + 1 : source_index + self.num_horizons + 1
                ],
                dtype=np.float32,
            ).copy()
        )
        return (
            rgb,
            audio,
            session.omni[source_index].float(),
            session.tail_bins[source_index],
            future,
            torch.tensor(session_index, dtype=torch.long),
            torch.tensor(source_index, dtype=torch.long),
        )


def miniroad_normalized_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
) -> torch.Tensor:
    """Released MiniROAD multi-positive anticipation objective."""

    if logits.shape != targets.shape or logits.shape[-1] != NUM_CLASSES:
        raise ValueError(
            f"Expected equal [...,{NUM_CLASSES}] logits/targets, got "
            f"{tuple(logits.shape)} and {tuple(targets.shape)}"
        )
    normalized_targets = F.normalize(targets, dim=-1)
    return torch.sum(
        -normalized_targets * F.log_softmax(logits, dim=-1),
        dim=-1,
    ).mean()


def align_source_predictions(
    raw_prediction: np.ndarray,
    target_steps: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Align MiniROADA [source,horizon,class] forecasts to target rows."""

    raw_prediction = np.asarray(raw_prediction, dtype=np.float32)
    expected = (target_steps - NUM_HORIZONS, NUM_HORIZONS, NUM_CLASSES)
    if raw_prediction.shape != expected:
        raise ValueError(
            f"Expected raw prediction {expected}, got {raw_prediction.shape}"
        )
    if not np.isfinite(raw_prediction).all():
        raise ValueError("Raw prediction contains NaN/Inf")
    predictions = np.zeros(
        (target_steps, NUM_CLASSES, NUM_HORIZONS),
        dtype=np.float32,
    )
    valid = np.zeros((target_steps, NUM_HORIZONS), dtype=bool)
    source_rows = np.arange(len(raw_prediction), dtype=np.int64)
    for horizon in range(NUM_HORIZONS):
        target_rows = source_rows + horizon + 1
        predictions[target_rows, :, horizon] = raw_prediction[:, horizon, :]
        valid[target_rows, horizon] = True
    return predictions, valid