"""Tiny PyTorch candidate scorer used only by the optional ML experiment."""

import math
from dataclasses import dataclass

import torch
from torch import nn

from .features import ACTION_FEATURES, FEATURE_DIMENSION

DISPOSITION_COUNT = 3
MAX_HIDDEN_DIMENSION = 128


@dataclass(frozen=True, slots=True)
class FeatureBatch:
    """Dense instruction vectors with padded, explicitly masked actions."""

    vectors: torch.Tensor
    actions: torch.Tensor
    action_mask: torch.Tensor


def _finite_row(values: list[float], width: int, label: str) -> None:
    if len(values) != width:
        raise ValueError(f"A {label} has the wrong feature width.")
    if any(type(value) not in (int, float) or not math.isfinite(value) for value in values):
        raise ValueError(f"A {label} must contain only finite numeric features.")


def collate_features(vectors: list[list[float]], actions: list[list[list[float]]]) -> FeatureBatch:
    """Validate and pad variable action rows while retaining a boolean validity mask."""
    if not vectors or len(vectors) != len(actions):
        raise ValueError("Feature batches must be nonempty and aligned.")
    for vector in vectors:
        _finite_row(vector, FEATURE_DIMENSION, "vector")
    maximum = max(1, *(len(items) for items in actions))
    padded: list[list[list[float]]] = []
    masks: list[list[bool]] = []
    for items in actions:
        for item in items:
            _finite_row(item, ACTION_FEATURES, "action")
        missing = maximum - len(items)
        padded.append(items + [[0.0] * ACTION_FEATURES for _ in range(missing)])
        masks.append([True] * len(items) + [False] * missing)
    return FeatureBatch(
        torch.tensor(vectors, dtype=torch.float32),
        torch.tensor(padded, dtype=torch.float32),
        torch.tensor(masks, dtype=torch.bool),
    )


class CandidateScorer(nn.Module):
    """Encode an instruction and jointly score disposition and valid actions."""

    def __init__(self, feature_dimension: int, hidden_dimension: int) -> None:
        """Create a bounded two-head scorer from explicit dimensions."""
        super().__init__()
        if (
            type(feature_dimension) is not int
            or feature_dimension != FEATURE_DIMENSION
            or type(hidden_dimension) is not int
            or not 1 <= hidden_dimension <= MAX_HIDDEN_DIMENSION
        ):
            raise ValueError("Model dimensions exceed the bounded experiment contract.")
        self.encoder = nn.Linear(feature_dimension, hidden_dimension)
        self.disposition = nn.Linear(hidden_dimension, DISPOSITION_COUNT)
        self.action = nn.Linear(ACTION_FEATURES, hidden_dimension)

    def validate_architecture(self, feature_dimension: int, hidden_dimension: int) -> None:
        """Reject a scorer that does not match a validated artifact configuration."""
        if (
            self.encoder.in_features != feature_dimension
            or self.encoder.out_features != hidden_dimension
            or self.disposition.in_features != hidden_dimension
            or self.action.out_features != hidden_dimension
            or self.action.in_features != ACTION_FEATURES
        ):
            raise ValueError("The model architecture does not match its configuration.")

    def forward(
        self, vectors: torch.Tensor, actions: torch.Tensor, action_mask: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return three-way disposition logits and masked per-action logits."""
        hidden = torch.tanh(self.encoder(vectors))
        disposition_logits = self.disposition(hidden)
        action_vectors = torch.tanh(self.action(actions))
        action_logits = torch.einsum("bh,bah->ba", hidden, action_vectors)
        return disposition_logits, action_logits.masked_fill(~action_mask, -1.0e9)