"""Exact memory-bounded vocabulary projection, sampling, and backward."""

from __future__ import annotations

import math
from collections.abc import Sequence

import torch
from torch.nn import functional as F


def _stable_max(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Return ``max`` indices that stay in-range on MPS NaN/-inf rows."""

    best, index = values.max(dim=-1)
    width = values.shape[-1]
    if width <= 0:
        return best, index
    return best, index.clamp(0, width - 1)


class _ChunkedVocabStatistics(torch.autograd.Function):
    """One projection pass for CE and exact Gumbel-max categorical sampling."""

    @staticmethod
    def forward(
        ctx,
        hidden: torch.Tensor,
        weight: torch.Tensor,
        labels: torch.LongTensor,
        softcap: float,
        temperature: float,
        chunk_size: int,
        top_k: int | None = 40,
        min_p: float | None = 0.05,
    ) -> tuple[
        torch.Tensor,
        torch.Tensor,
        torch.LongTensor,
        torch.Tensor,
        torch.Tensor,
        torch.LongTensor,
        torch.Tensor,
    ]:
        if hidden.ndim != 2 or weight.ndim != 2:
            raise ValueError("Chunked vocabulary tensors must be matrices.")
        if hidden.shape[1] != weight.shape[1]:
            raise ValueError("Hidden and vocabulary projection dimensions differ.")
        if temperature <= 0 or chunk_size <= 0:
            raise ValueError("Temperature and vocabulary chunk size must be positive.")
        labels = labels.to(device=hidden.device, dtype=torch.long)
        if labels.shape != (hidden.shape[0],):
            raise ValueError("Labels must contain one value per hidden row.")

        rows = hidden.shape[0]
        row_indices = torch.arange(rows, device=hidden.device)
        valid = labels.ge(0)
        raw_log_z = torch.full(
            (rows,), -torch.inf, device=hidden.device, dtype=torch.float32
        )
        sample_log_z = torch.full_like(raw_log_z, -torch.inf)
        gold_score = torch.zeros_like(raw_log_z)
        best_gumbel = torch.full_like(raw_log_z, -torch.inf)
        selected_score = torch.zeros_like(raw_log_z)
        selected = torch.zeros(rows, device=hidden.device, dtype=torch.long)
        greedy_score = torch.full_like(raw_log_z, -torch.inf)
        greedy = torch.zeros(rows, device=hidden.device, dtype=torch.long)
        moment_max = torch.full_like(raw_log_z, -torch.inf)
        moment_sum = torch.zeros_like(raw_log_z)
        moment_weighted = torch.zeros_like(raw_log_z)
        vocab_size = weight.shape[0]

        use_constrained = top_k is not None and top_k > 0
        chunk_cand_scores = []
        chunk_cand_tokens = []

        with torch.no_grad():
            for start in range(0, vocab_size, int(chunk_size)):
                stop = min(start + int(chunk_size), vocab_size)
                raw = F.linear(hidden, weight[start:stop])
                scores = (
                    torch.tanh(raw.float() / float(softcap))
                    * float(softcap)
                )
                raw_log_z = torch.logaddexp(
                    raw_log_z,
                    torch.logsumexp(scores, dim=-1),
                )
                sample_scores = scores / float(temperature)
                sample_log_z = torch.logaddexp(
                    sample_log_z,
                    torch.logsumexp(sample_scores, dim=-1),
                )

                in_chunk = valid & labels.ge(start) & labels.lt(stop)
                local_gold = (labels - start).clamp(0, stop - start - 1)
                gold_score = torch.where(
                    in_chunk,
                    scores[row_indices, local_gold],
                    gold_score,
                )

                if use_constrained:
                    k = min(int(top_k), stop - start)
                    c_score, c_idx = torch.topk(sample_scores, k=k, dim=-1)
                    chunk_cand_scores.append(c_score)
                    chunk_cand_tokens.append(c_idx + start)
                else:
                    uniform = torch.rand(
                        sample_scores.shape,
                        device=sample_scores.device,
                        dtype=torch.float32,
                    ).clamp_(
                        min=torch.finfo(torch.float32).tiny,
                        max=1.0 - torch.finfo(torch.float32).eps,
                    )
                    gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
                    chunk_best, chunk_index = _stable_max(gumbel_scores)
                    replace_best = chunk_best.gt(best_gumbel)
                    candidate_score = sample_scores.gather(
                        1, chunk_index[:, None]
                    ).squeeze(-1)
                    best_gumbel = torch.maximum(best_gumbel, chunk_best)
                    selected = torch.where(
                        replace_best,
                        chunk_index + start,
                        selected,
                    )
                    selected_score = torch.where(
                        replace_best,
                        candidate_score,
                        selected_score,
                    )

                chunk_max, chunk_argmax = _stable_max(sample_scores)
                replace_greedy = chunk_max.gt(greedy_score)
                greedy_score = torch.maximum(greedy_score, chunk_max)
                greedy = torch.where(
                    replace_greedy,
                    chunk_argmax + start,
                    greedy,
                )
                shifted = torch.exp(sample_scores - chunk_max[:, None])
                chunk_sum = shifted.sum(dim=-1)
                chunk_weighted = (shifted * sample_scores).sum(dim=-1)
                merged_max = torch.maximum(moment_max, chunk_max)
                previous_scale = torch.exp(moment_max - merged_max)
                chunk_scale = torch.exp(chunk_max - merged_max)
                moment_sum = (
                    moment_sum * previous_scale + chunk_sum * chunk_scale
                )
                moment_weighted = (
                    moment_weighted * previous_scale
                    + chunk_weighted * chunk_scale
                )
                moment_max = merged_max

        if use_constrained:
            all_scores = torch.cat(chunk_cand_scores, dim=-1)
            all_tokens = torch.cat(chunk_cand_tokens, dim=-1)
            global_k = min(int(top_k), all_scores.shape[-1])
            cand_scores, global_idx = torch.topk(all_scores, k=global_k, dim=-1)
            cand_tokens = all_tokens.gather(dim=-1, index=global_idx)
            if min_p is not None and min_p > 0:
                thresh = greedy_score + math.log(float(min_p))
                valid_mask = cand_scores >= thresh[:, None]
                eligible = cand_scores.masked_fill(~valid_mask, -float("inf"))
            else:
                eligible = cand_scores
            uniform = torch.rand(
                eligible.shape,
                device=eligible.device,
                dtype=torch.float32,
            ).clamp_(
                min=torch.finfo(torch.float32).tiny,
                max=1.0 - torch.finfo(torch.float32).eps,
            )
            gumbel = eligible - torch.log(-torch.log(uniform))
            _, chosen_idx = _stable_max(gumbel)
            selected = cand_tokens.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)
            selected_score = cand_scores.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)

        raw_nll = torch.where(
            valid,
            raw_log_z - gold_score,
            torch.zeros_like(raw_log_z),
        )
        temperature_nll = torch.where(
            valid,
            sample_log_z - gold_score / float(temperature),
            torch.zeros_like(sample_log_z),
        )
        confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
        greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
        entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
            torch.finfo(torch.float32).tiny
        )

        ctx.save_for_backward(
            hidden,
            weight,
            labels,
            selected,
            raw_log_z,
            sample_log_z,
            confidence,
            entropy,
        )
        ctx.softcap = float(softcap)
        ctx.temperature = float(temperature)
        ctx.chunk_size = int(chunk_size)
        ctx.mark_non_differentiable(selected, greedy, greedy_confidence)
        return (
            raw_nll,
            temperature_nll,
            selected,
            confidence,
            entropy,
            greedy,
            greedy_confidence,
            sample_log_z - selected_score,
        )

    @staticmethod
    def backward(
        ctx,
        grad_raw_nll: torch.Tensor | None,
        grad_temperature_nll: torch.Tensor | None,
        grad_selected: torch.Tensor | None,
        grad_confidence: torch.Tensor | None,
        grad_entropy: torch.Tensor | None,
        grad_greedy: torch.Tensor | None,
        grad_greedy_confidence: torch.Tensor | None,
        grad_sampled_nll: torch.Tensor | None,
    ):
        del grad_selected, grad_greedy, grad_greedy_confidence
        (
            hidden,
            weight,
            labels,
            selected,
            raw_log_z,
            sample_log_z,
            confidence,
            entropy,
        ) = ctx.saved_tensors
        grad_raw_nll = (
            torch.zeros_like(raw_log_z)
            if grad_raw_nll is None
            else grad_raw_nll.float()
        )
        grad_temperature_nll = (
            torch.zeros_like(sample_log_z)
            if grad_temperature_nll is None
            else grad_temperature_nll.float()
        )
        grad_confidence = (
            torch.zeros_like(confidence)
            if grad_confidence is None
            else grad_confidence.float()
        )
        valid = labels.ge(0)
        row_indices = torch.arange(hidden.shape[0], device=hidden.device)
        grad_hidden = torch.zeros_like(hidden, dtype=torch.float32)
        confidence_scale = (
            grad_confidence * confidence / ctx.temperature
        )
        if grad_sampled_nll is not None:
            confidence_scale = confidence_scale - grad_sampled_nll.float() / ctx.temperature
        hidden_fp32 = hidden.float()
        for start in range(0, weight.shape[0], ctx.chunk_size):
            stop = min(start + ctx.chunk_size, weight.shape[0])
            # The B16 training shape is 4,096 rows by a 32,768-wide vocabulary
            # chunk.  MPS BF16 GEMM can return NaNs for this backward-only
            # projection even when every saved activation is finite.  Recompute
            # scores and accumulate the exact VJP in FP32; cast the finished
            # hidden gradient back to the model dtype only once below.
            weight_fp32 = weight[start:stop].float()
            scores = torch.tanh(
                F.linear(hidden_fp32, weight_fp32) / ctx.softcap
            ) * ctx.softcap
            # Non-deterministic MPS GEMM can make a recomputed score slightly
            # exceed the saved log-partition.  Clamp the exponent so softmax
            # probabilities stay in [0, 1] instead of overflowing to Inf.
            raw_probability = torch.exp(
                (scores - raw_log_z[:, None]).clamp(max=0.0)
            )
            sample_logit = (
                scores / ctx.temperature - sample_log_z[:, None]
            ).clamp(max=0.0)
            sample_probability = torch.exp(sample_logit)
            score_gradient = (
                raw_probability * (grad_raw_nll * valid)[:, None]
                + sample_probability
                * (grad_temperature_nll * valid)[:, None]
                / ctx.temperature
                - sample_probability * confidence_scale[:, None]
            )
            if grad_entropy is not None:
                grad_entropy_f = grad_entropy.float()
                entropy_grad = -(
                    grad_entropy_f / ctx.temperature
                )[:, None] * sample_probability * (
                    sample_logit + entropy[:, None]
                )
                score_gradient = score_gradient + entropy_grad

            gold_in_chunk = valid & labels.ge(start) & labels.lt(stop)
            gold_local = (labels - start).clamp(0, stop - start - 1)
            score_gradient[row_indices, gold_local] -= (
                grad_raw_nll * gold_in_chunk
                + grad_temperature_nll
                * gold_in_chunk
                / ctx.temperature
            )
            selected_in_chunk = selected.ge(start) & selected.lt(stop)
            selected_local = (selected - start).clamp(0, stop - start - 1)
            score_gradient[row_indices, selected_local] += (
                confidence_scale * selected_in_chunk
            )
            score_gradient *= 1.0 - (scores / ctx.softcap).square()
            score_gradient = torch.nan_to_num(
                score_gradient, nan=0.0, posinf=0.0, neginf=0.0
            )
            chunk_gradient = F.linear(
                score_gradient,
                weight_fp32.transpose(0, 1),
            )
            grad_hidden.add_(chunk_gradient)

        grad_hidden = torch.nan_to_num(
            grad_hidden, nan=0.0, posinf=0.0, neginf=0.0
        )
        return (
            grad_hidden.to(hidden.dtype),
            None,
            None,
            None,
            None,
            None,
            None,
            None,
        )


def chunked_vocab_training_statistics(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    labels: torch.LongTensor,
    *,
    softcap: float,
    temperature: float,
    chunk_size: int,
    top_k: int | None = 40,
    min_p: float | None = 0.05,
    return_sampled_nll: bool = False,
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.LongTensor,
    torch.Tensor,
    torch.Tensor,
    torch.LongTensor,
    torch.Tensor,
]:
    """Return exact losses plus sampled and greedy proposal statistics."""

    shape = labels.shape
    outputs = _ChunkedVocabStatistics.apply(
        hidden.reshape(-1, hidden.shape[-1]),
        weight,
        labels.reshape(-1),
        float(softcap),
        float(temperature),
        int(chunk_size),
        top_k,
        min_p,
    )
    (
        raw_nll,
        temperature_nll,
        proposal,
        confidence,
        entropy,
        greedy,
        greedy_confidence,
        sampled_nll,
    ) = outputs
    result = (
        raw_nll.view(shape),
        temperature_nll.view(shape),
        proposal.view(shape),
        confidence.view(shape),
        entropy.view(shape),
        greedy.view(shape),
        greedy_confidence.view(shape),
    )
    return result + (sampled_nll.view(shape),) if return_sampled_nll else result


@torch.no_grad()
def chunked_vocab_statistics(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    *,
    softcap: float,
    temperature: float,
    chunk_size: int,
    repetition_token_mask: torch.BoolTensor | None = None,
    repetition_penalty: float = 1.0,
    sampling_generators: Sequence[torch.Generator] | None = None,
    top_k: int | None = 40,
    min_p: float | None = 0.05,
) -> tuple[
    torch.LongTensor,
    torch.Tensor,
    torch.Tensor,
    torch.LongTensor,
    torch.Tensor,
]:
    """Return exact inference statistics, optionally penalizing seen tokens."""

    if not math.isfinite(repetition_penalty) or repetition_penalty <= 0:
        raise ValueError("`repetition_penalty` must be a finite positive number.")

    # Preserve the original custom-autograd path exactly when disabled. In
    # particular, this keeps the same chunk-local RNG draws and default output.
    if repetition_penalty != 1.0 or sampling_generators is not None:
        if hidden.ndim < 2 or weight.ndim != 2:
            raise ValueError("Chunked vocabulary tensors must have matrix features.")
        if hidden.shape[-1] != weight.shape[1]:
            raise ValueError("Hidden and vocabulary projection dimensions differ.")
        if temperature <= 0 or chunk_size <= 0:
            raise ValueError("Temperature and vocabulary chunk size must be positive.")
        expected_mask_shape = (hidden.shape[0], weight.shape[0])
        if repetition_penalty != 1.0:
            if repetition_token_mask is None or repetition_token_mask.shape != expected_mask_shape:
                raise ValueError(
                    "`repetition_token_mask` must have shape [batch, vocabulary]."
                )
            if repetition_token_mask.dtype != torch.bool:
                raise ValueError("`repetition_token_mask` must be a boolean tensor.")
        if sampling_generators is not None and len(sampling_generators) != hidden.shape[0]:
            raise ValueError(
                "`sampling_generators` must contain one generator per batch row."
            )

        output_shape = hidden.shape[:-1]
        flat_hidden = hidden.reshape(-1, hidden.shape[-1])
        rows_per_batch = math.prod(hidden.shape[1:-1])
        batch_indices = torch.arange(
            flat_hidden.shape[0], device=hidden.device
        ).div(rows_per_batch, rounding_mode="floor")
        sample_log_z = torch.full(
            (flat_hidden.shape[0],),
            -torch.inf,
            device=hidden.device,
            dtype=torch.float32,
        )
        best_gumbel = torch.full_like(sample_log_z, -torch.inf)
        selected_score = torch.zeros_like(sample_log_z)
        selected = torch.zeros(
            flat_hidden.shape[0], device=hidden.device, dtype=torch.long
        )
        greedy_score = torch.full_like(sample_log_z, -torch.inf)
        greedy = torch.zeros_like(selected)
        moment_max = torch.full_like(sample_log_z, -torch.inf)
        moment_sum = torch.zeros_like(sample_log_z)
        moment_weighted = torch.zeros_like(sample_log_z)

        use_constrained = top_k is not None and top_k > 0
        chunk_cand_scores = []
        chunk_cand_tokens = []

        for start in range(0, weight.shape[0], int(chunk_size)):
            stop = min(start + int(chunk_size), weight.shape[0])
            raw = F.linear(flat_hidden, weight[start:stop])
            scores = torch.tanh(raw.float() / float(softcap)) * float(softcap)
            if repetition_penalty != 1.0:
                assert repetition_token_mask is not None
                seen = repetition_token_mask[:, start:stop].index_select(
                    0, batch_indices
                )
                penalized = torch.where(
                    scores < 0,
                    scores * float(repetition_penalty),
                    scores / float(repetition_penalty),
                )
                scores = torch.where(seen, penalized, scores)
            sample_scores = scores / float(temperature)
            sample_log_z = torch.logaddexp(
                sample_log_z, torch.logsumexp(sample_scores, dim=-1)
            )

            if use_constrained:
                k = min(int(top_k), stop - start)
                c_score, c_idx = torch.topk(sample_scores, k=k, dim=-1)
                chunk_cand_scores.append(c_score)
                chunk_cand_tokens.append(c_idx + start)
            else:
                if sampling_generators is None:
                    uniform = torch.rand(
                        sample_scores.shape,
                        device=sample_scores.device,
                        dtype=torch.float32,
                    )
                else:
                    per_request_shape = (
                        rows_per_batch,
                        sample_scores.shape[-1],
                    )
                    uniform = torch.cat(
                        [
                            torch.rand(
                                per_request_shape,
                                device=sample_scores.device,
                                dtype=torch.float32,
                                generator=generator,
                            )
                            for generator in sampling_generators
                        ],
                        dim=0,
                    )
                uniform = uniform.clamp_(
                    min=torch.finfo(torch.float32).tiny,
                    max=1.0 - torch.finfo(torch.float32).eps,
                )
                gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
                chunk_best, chunk_index = _stable_max(gumbel_scores)
                replace_best = chunk_best.gt(best_gumbel)
                candidate_score = sample_scores.gather(
                    1, chunk_index[:, None]
                ).squeeze(-1)
                best_gumbel = torch.maximum(best_gumbel, chunk_best)
                selected = torch.where(replace_best, chunk_index + start, selected)
                selected_score = torch.where(
                    replace_best, candidate_score, selected_score
                )

            chunk_max, chunk_argmax = _stable_max(sample_scores)
            replace_greedy = chunk_max.gt(greedy_score)
            greedy_score = torch.maximum(greedy_score, chunk_max)
            greedy = torch.where(replace_greedy, chunk_argmax + start, greedy)
            shifted = torch.exp(sample_scores - chunk_max[:, None])
            chunk_sum = shifted.sum(dim=-1)
            chunk_weighted = (shifted * sample_scores).sum(dim=-1)
            merged_max = torch.maximum(moment_max, chunk_max)
            previous_scale = torch.exp(moment_max - merged_max)
            chunk_scale = torch.exp(chunk_max - merged_max)
            moment_sum = moment_sum * previous_scale + chunk_sum * chunk_scale
            moment_weighted = (
                moment_weighted * previous_scale + chunk_weighted * chunk_scale
            )
            moment_max = merged_max

        if use_constrained:
            all_scores = torch.cat(chunk_cand_scores, dim=-1)
            all_tokens = torch.cat(chunk_cand_tokens, dim=-1)
            global_k = min(int(top_k), all_scores.shape[-1])
            cand_scores, global_idx = torch.topk(all_scores, k=global_k, dim=-1)
            cand_tokens = all_tokens.gather(dim=-1, index=global_idx)
            if min_p is not None and min_p > 0:
                thresh = greedy_score + math.log(float(min_p))
                valid_mask = cand_scores >= thresh[:, None]
                eligible = cand_scores.masked_fill(~valid_mask, -float("inf"))
            else:
                eligible = cand_scores
            if sampling_generators is None:
                uniform = torch.rand(
                    eligible.shape,
                    device=eligible.device,
                    dtype=torch.float32,
                )
            else:
                per_request_shape = (
                    rows_per_batch,
                    eligible.shape[-1],
                )
                uniform = torch.cat(
                    [
                        torch.rand(
                            per_request_shape,
                            device=eligible.device,
                            dtype=torch.float32,
                            generator=generator,
                        )
                        for generator in sampling_generators
                    ],
                    dim=0,
                )
            uniform = uniform.clamp_(
                min=torch.finfo(torch.float32).tiny,
                max=1.0 - torch.finfo(torch.float32).eps,
            )
            gumbel = eligible - torch.log(-torch.log(uniform))
            _, chosen_idx = _stable_max(gumbel)
            selected = cand_tokens.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)
            selected_score = cand_scores.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)

        confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
        greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
        entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
            torch.finfo(torch.float32).tiny
        )
        return (
            selected.view(output_shape),
            confidence.view(output_shape),
            entropy.view(output_shape),
            greedy.view(output_shape),
            greedy_confidence.view(output_shape),
        )

    labels = torch.full(
        hidden.shape[:-1],
        -100,
        device=hidden.device,
        dtype=torch.long,
    )
    (
        _,
        _,
        proposal,
        confidence,
        entropy,
        greedy,
        greedy_confidence,
    ) = chunked_vocab_training_statistics(
        hidden,
        weight,
        labels,
        softcap=softcap,
        temperature=temperature,
        chunk_size=chunk_size,
        top_k=top_k,
        min_p=min_p,
    )
    return proposal, confidence, entropy, greedy, greedy_confidence


__all__ = [
    "chunked_vocab_statistics",
    "chunked_vocab_training_statistics",
]