"""Lloyd-Max codec for the TurboQuant *value* path — executable reference.

TurboQuant gives keys a rotation plus a Lloyd-Max optimal scalar quantizer, but
stores values with plain uniform asymmetric quantization.  Measured on this
model's head_dim of 256, values are the dominant error source: at three bits
they contribute a relative attention-output error of 0.291 against 0.171 for
keys, and the published perplexity cost of the uniform 3/3 preset is +20.59 %.

Giving values the key-side treatment costs nothing in memory - it is two bytes
*cheaper* per vector, because a single vector norm replaces the scale/zero pair -
and lowers the attention-output error from 0.3474 to 0.2577 at the same bit
budget.

This module is the reference the Triton kernels must reproduce.  It is written
to be read next to ``triton_turboquant_store.py`` and
``triton_turboquant_decode.py``: the same rotation, the same midpoint search,
the same norm correction.

The decode side never rotates a cached vector back.  With values stored as
``V_rot = V @ R`` and unit-normalised,

    P @ V_hat  ==  ((P * norms) @ V_unit_rot) @ R^T

so the per-token norm folds into the attention weight and the inverse rotation
is a single matmul on the output.  ``fused_attention_reference`` implements that
form; ``dequantize_values`` implements the materialising form used by the
continuation-prefill path, where the inverse rotation is one batched GEMM after
the kernel.
"""

from __future__ import annotations

from dataclasses import dataclass
import math

import torch

SCHEMA = "lomonosov_zenit_tq_value_mse_v1"

# Presets this codec adds. ``value_codec`` is the only field that differs from
# the upstream TurboQuant presets of the same name.
VALUE_MSE_PRESETS: dict[str, dict] = {
    "zenit_k3v3_mse": {
        "key_quant_bits": 3,
        "value_quant_bits": 3,
        "norm_correction": True,
        "value_codec": "lloyd_max",
    },
    "zenit_k3v4_mse": {
        "key_quant_bits": 3,
        "value_quant_bits": 4,
        "norm_correction": True,
        "value_codec": "lloyd_max",
    },
    "zenit_k4v4_mse": {
        "key_quant_bits": 4,
        "value_quant_bits": 4,
        "norm_correction": True,
        "value_codec": "lloyd_max",
    },
}


def value_packed_size(head_dim: int, bits: int) -> int:
    """Bytes per value vector: packed indices plus one fp16 norm.

    The uniform path stores four extra bytes (fp16 scale and fp16 zero); the
    Lloyd-Max path needs only the two-byte vector norm.
    """
    return math.ceil(head_dim * bits / 8) + 2


def slot_size(head_dim: int, key_bits: int, value_bits: int) -> int:
    key = math.ceil(head_dim * key_bits / 8) + 2
    total = key + value_packed_size(head_dim, value_bits)
    return total + (total % 2)


def hadamard(head_dim: int, device=None, dtype=torch.float32) -> torch.Tensor:
    """Orthonormal Sylvester Hadamard matrix, matching ``_build_hadamard``."""
    if head_dim & (head_dim - 1):
        raise ValueError(f"head_dim must be a power of two, got {head_dim}")
    matrix = torch.tensor([[1.0]], dtype=dtype)
    while matrix.shape[0] < head_dim:
        matrix = torch.cat(
            [
                torch.cat([matrix, matrix], dim=1),
                torch.cat([matrix, -matrix], dim=1),
            ],
            dim=0,
        )
    return (matrix / math.sqrt(head_dim)).to(device=device, dtype=dtype)


def lloyd_max_tables(head_dim: int, bits: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Centroids and midpoints, from vLLM's own solver so the tables match."""
    from vllm.model_executor.layers.quantization.turboquant.centroids import (
        solve_lloyd_max,
    )

    centroids, midpoints = solve_lloyd_max(head_dim, bits)
    return centroids.float(), midpoints.float()


@dataclass(frozen=True)
class EncodedValues:
    """What the store kernel writes for one value vector per head per token."""

    indices: torch.Tensor  # [..., head_dim] int32, in [0, 2**bits)
    norms: torch.Tensor  # [...] float32, cast to fp16 in the cache
    bits: int

    @property
    def bytes_per_vector(self) -> int:
        return value_packed_size(self.indices.shape[-1], self.bits)


def encode_values(
    values: torch.Tensor,
    bits: int,
    *,
    rotation: torch.Tensor | None = None,
    midpoints: torch.Tensor | None = None,
) -> EncodedValues:
    """Rotate, unit-normalise, and bucketize - the key path applied to values.

    ``values`` is [..., head_dim] in the model's own coordinates.
    """
    head_dim = values.shape[-1]
    rotation = hadamard(head_dim, values.device) if rotation is None else rotation
    if midpoints is None:
        _, midpoints = lloyd_max_tables(head_dim, bits)
        midpoints = midpoints.to(values.device)

    rotated = values.float() @ rotation
    norms = rotated.norm(dim=-1, keepdim=True).clamp(min=1e-12)
    unit = rotated / norms
    indices = torch.bucketize(unit, midpoints).to(torch.int32)
    return EncodedValues(indices=indices, norms=norms.squeeze(-1), bits=bits)


def decode_unit_values(
    encoded: EncodedValues,
    *,
    centroids: torch.Tensor | None = None,
    norm_correction: bool = True,
) -> torch.Tensor:
    """Unit-norm value vectors **in rotated space** - what the kernel holds."""
    head_dim = encoded.indices.shape[-1]
    if centroids is None:
        centroids, _ = lloyd_max_tables(head_dim, encoded.bits)
        centroids = centroids.to(encoded.indices.device)
    unit = centroids[encoded.indices.long()]
    if norm_correction:
        unit = unit / unit.norm(dim=-1, keepdim=True).clamp(min=1e-12)
    return unit


def dequantize_values(
    encoded: EncodedValues,
    *,
    rotation: torch.Tensor | None = None,
    centroids: torch.Tensor | None = None,
    norm_correction: bool = True,
) -> torch.Tensor:
    """Materialise values in model coordinates.

    Used by the continuation-prefill path, where the inverse rotation is one
    batched GEMM applied to the whole dequantised tensor rather than per vector.
    """
    head_dim = encoded.indices.shape[-1]
    rotation = (
        hadamard(head_dim, encoded.indices.device) if rotation is None else rotation
    )
    unit = decode_unit_values(
        encoded, centroids=centroids, norm_correction=norm_correction
    )
    return (unit * encoded.norms.unsqueeze(-1)) @ rotation.T


def fused_attention_reference(
    probabilities: torch.Tensor,
    encoded: EncodedValues,
    *,
    rotation: torch.Tensor | None = None,
    centroids: torch.Tensor | None = None,
    norm_correction: bool = True,
) -> torch.Tensor:
    """The decode-kernel form: no cached vector is ever rotated back.

    ``probabilities`` is [queries, context] post-softmax. The per-token norm
    folds into the weight and the inverse rotation is applied once to the
    output.
    """
    head_dim = encoded.indices.shape[-1]
    rotation = (
        hadamard(head_dim, encoded.indices.device) if rotation is None else rotation
    )
    unit = decode_unit_values(
        encoded, centroids=centroids, norm_correction=norm_correction
    )
    weighted = probabilities * encoded.norms.unsqueeze(0)
    return (weighted @ unit) @ rotation.T


def uniform_reference(values: torch.Tensor, bits: int) -> torch.Tensor:
    """The codec being replaced, for A/B comparison at identical bit width."""
    levels = (1 << bits) - 1
    low = values.amin(dim=-1, keepdim=True)
    high = values.amax(dim=-1, keepdim=True)
    scale = ((high - low) / levels).to(torch.float16).float().clamp(min=1e-12)
    zero = low.to(torch.float16).float()
    quantized = torch.round((values - zero) / scale).clamp(0, levels)
    return quantized * scale + zero


__all__ = [
    "SCHEMA",
    "VALUE_MSE_PRESETS",
    "EncodedValues",
    "decode_unit_values",
    "dequantize_values",
    "encode_values",
    "fused_attention_reference",
    "hadamard",
    "lloyd_max_tables",
    "slot_size",
    "uniform_reference",
    "value_packed_size",
]