"""Extract single-letter logits and compute top-k softmax scores.

The engine reads one next-token position, keeps logits for the 26 uppercase
ASCII letters, and converts the top-k letter logits into a probability
distribution. It does not sample or decode text.
"""

from __future__ import annotations

import math
import string
from dataclasses import dataclass

LETTERS: tuple[str, ...] = tuple(string.ascii_uppercase)
LETTER_INDEX: dict[str, int] = {c: i for i, c in enumerate(LETTERS)}

# Rank at which we stop tracking for the reported top-k. k <= 26 by construction.
MAX_TOP_K = len(LETTERS)


class LetterLogitError(ValueError):
    """Raised when the model's next-token distribution cannot be read as letters."""


@dataclass(frozen=True)
class LetterReadout:
    """The full 26-letter logit readout for one prompt, plus a top-k view.

    ``logits`` keeps every letter when a question has more options than ``top_k``.
    ``top`` is the reported top-k softmax distribution.
    ``letter_mass`` is retained as a nullable compatibility field. The optimized
    26-row output head intentionally cannot report it because it never computes
    a full-vocabulary distribution.
    """

    logits: dict[str, float]
    top: "SoftmaxDistribution"
    letter_mass: float | None = None


@dataclass(frozen=True)
class SoftmaxDistribution:
    """A probability distribution over a chosen slice of letters."""

    ranked: tuple[tuple[str, float], ...]
    probabilities: dict[str, float]
    top_k: int
    temperature: float

    @property
    def top_letter(self) -> str:
        return self.ranked[0][0]

    @property
    def top_probability(self) -> float:
        return self.probabilities[self.top_letter]


def _softmax(logits: list[float], temperature: float) -> list[float]:
    if temperature <= 0:
        raise ValueError("temperature must be > 0")
    scaled = [x / temperature for x in logits]
    peak = max(scaled)
    exps = [math.exp(x - peak) for x in scaled]
    total = sum(exps)
    if total <= 0 or not math.isfinite(total):
        raise LetterLogitError("softmax denominator is not finite")
    return [e / total for e in exps]


def distribution_from_letter_logits(
    letter_logits: dict[str, float],
    *,
    top_k: int = 5,
    temperature: float = 1.0,
) -> SoftmaxDistribution:
    """Softmax the top-k letter logits into a probability distribution.

    `top_k` is clamped to 2..26. The softmax is taken *only* over the k letters we
    keep, which is what TypeSafe-compatible consumers expect: the reported
    probabilities sum to 1 across the surfaced options.
    """
    unknown = set(letter_logits) - set(LETTER_INDEX)
    if unknown:
        raise LetterLogitError(f"non-letter logit keys: {sorted(unknown)}")
    if not letter_logits:
        raise LetterLogitError("no letter logits supplied")

    k = max(1, min(int(top_k), len(letter_logits), MAX_TOP_K))
    ranked = sorted(letter_logits.items(), key=lambda kv: kv[1], reverse=True)[:k]
    probs = _softmax([v for _, v in ranked], temperature)

    distribution = {letter: p for (letter, _), p in zip(ranked, probs)}
    # Keep the distribution in descending-probability order for stable output.
    ordered = dict(sorted(distribution.items(), key=lambda kv: kv[1], reverse=True))
    return SoftmaxDistribution(
        ranked=tuple((letter, logit) for letter, logit in ranked),
        probabilities=ordered,
        top_k=k,
        temperature=temperature,
    )


def gather_letter_logits(next_token_logits, letter_token_ids):
    """Return the ``[batch, 26]`` A-Z slice of a vocabulary-logit tensor.

    ``letter_token_ids`` may be the device-resident tensor owned by the engine.
    Keeping it on-device avoids allocating a Python list and an indexing tensor on
    every request.
    """
    if next_token_logits.dim() != 2:
        raise LetterLogitError("expected next-token logits of shape [batch, vocab]")

    import torch

    ids = (
        letter_token_ids
        if isinstance(letter_token_ids, torch.Tensor)
        else torch.tensor(list(letter_token_ids.values()), device=next_token_logits.device)
    )
    return next_token_logits.index_select(1, ids)


def batch_letter_logits(
    gathered_logits, letters: tuple[str, ...] | dict[str, int] = LETTERS
) -> list[dict[str, float]]:
    """Convert a previously gathered ``[batch, 26]`` A-Z tensor to dictionaries."""
    # Preserve the original public helper signature for downstream callers while
    # allowing the engine to pass the already-gathered slice.
    if isinstance(letters, dict):
        gathered_logits = gather_letter_logits(gathered_logits, letters)
        letters = tuple(letters)
    if gathered_logits.dim() != 2 or gathered_logits.shape[1] != len(letters):
        raise LetterLogitError("expected gathered letter logits of shape [batch, 26]")
    return [
        {letter: float(value) for letter, value in zip(letters, row)}
        for row in gathered_logits
    ]