"""fp8 — FP8 (E4M3 / E5M2) floating-point 8-bit format primitives.

FP8 E4M3: 1 sign + 4 exponent + 3 mantissa. Bias=7. Wider mantissa → more
precision, narrower range. Used for weights/activations forward (NVIDIA H100+).

FP8 E5M2: 1 sign + 5 exponent + 2 mantissa. Bias=15. Wider range, less precision.
Used for gradients / heavy-tailed distributions.

Both encodable/decodable in pure PyTorch (no hardware FP8 tensor cores needed).
Dequant → fp32/fp16 matmul (like int8/NF4 strategies).
"""

from __future__ import annotations

import math

import torch


def _make_fp8_lut(exp_bits: int, mant_bits: int, bias: int) -> torch.Tensor:
    """Build the 256-entry lookup table for an FP8 format.

    Returns fp32 tensor [256] of dequantized values for each 8-bit code.
    Code layout: bit 7 = sign, bits 6..(7-exp) = exponent, low bits = mantissa.
    Special: all-zero exponent with zero mantissa = +0; all-zero exp with nonzero
    mantissa = subnormal; all-ones exponent = NaN/Inf (we map to ±448 for E4M3,
    ±57344 for E5M2 — max finite, to avoid NaN in dequant).
    """
    lut = torch.zeros(256, dtype=torch.float32)
    exp_mask = (1 << exp_bits) - 1
    mant_mask = (1 << mant_bits) - 1
    max_exp = exp_mask  # all-ones exponent
    max_finite = 0.0
    # Compute max finite value (exp = max_exp - 1, mant = all ones).
    emax = (max_exp - 1) - bias
    max_finite = (2.0 ** emax) * (1.0 + mant_mask / (2.0 ** mant_bits))
    for code in range(256):
        sign = -1.0 if (code & 0x80) else 1.0
        exp = (code >> mant_bits) & exp_mask
        mant = code & mant_mask
        if exp == 0:
            if mant == 0:
                val = 0.0  # signed zero
            else:
                # Subnormal: val = sign * mant / 2^mant_bits * 2^(1-bias)
                val = sign * (mant / (2.0 ** mant_bits)) * (2.0 ** (1 - bias))
        elif exp == max_exp:
            # Inf/NaN — map to max finite (avoid NaN propagation).
            val = sign * max_finite
        else:
            # Normalized: val = sign * (1 + mant/2^mb) * 2^(exp - bias)
            val = sign * (1.0 + mant / (2.0 ** mant_bits)) * (2.0 ** (exp - bias))
        lut[code] = val
    return lut


# Precomputed LUTs (256 entries each, fp32).
FP8_E4M3_LUT: torch.Tensor = _make_fp8_lut(exp_bits=4, mant_bits=3, bias=7)
FP8_E5M2_LUT: torch.Tensor = _make_fp8_lut(exp_bits=5, mant_bits=2, bias=15)


def quantize_fp8(w: torch.Tensor, lut: torch.Tensor, scale: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize a float tensor to FP8 codes + per-group/per-channel scale.

    Args:
        w: float tensor [out, in] (or any shape).
        lut: 256-entry FP8 LUT (FP8_E4M3_LUT or FP8_E5M2_LUT).
        scale: precomputed scale [out] (per-channel) or scalar. If None, computed
               as absmax(w) / max(abs(lut)).

    Returns:
        (codes, scale) where codes: uint8 [out, in], scale: fp32 [out] or scalar.
    """
    w = w.detach().float()
    if scale is None:
        max_lut = lut.abs().amax().clamp(min=1e-12)
        if w.dim() > 1:
            max_abs = w.abs().amax(dim=tuple(range(1, w.dim())))
            scale = (max_abs / max_lut).clamp(min=1e-12)
        else:
            scale = (w.abs().amax() / max_lut).clamp(min=1e-12).reshape(1)
    # Normalize by scale (broadcast depends on scale shape).
    if scale.numel() == 1:
        w_norm = w / scale
    else:
        reshape = [1] * w.dim()
        reshape[0] = w.shape[0]
        w_norm = w / scale.reshape(reshape)
    # Nearest LUT entry.
    # Clip to LUT range to avoid out-of-range.
    lut_max = lut.abs().amax().item()
    w_norm = w_norm.clamp(-lut_max, lut_max)
    diff = w_norm.unsqueeze(-1) - lut.to(w_norm.device)
    codes = diff.abs().argmin(dim=-1).to(torch.uint8)
    return codes, scale.to(torch.float32)


def dequantize_fp8(codes: torch.Tensor, scale: torch.Tensor, lut: torch.Tensor) -> torch.Tensor:
    """Reconstruct float tensor from FP8 codes + scale.

    Args:
        codes: uint8 [out, in] (or any shape).
        scale: fp32 [out] (per-channel) or scalar.
        lut: 256-entry FP8 LUT.

    Returns:
        fp32 tensor same shape as codes.
    """
    lut_dev = lut.to(codes.device)
    w = lut_dev[codes.long()]  # [shape], fp32
    if scale.numel() == 1:
        return w * scale.to(torch.float32)
    reshape = [1] * w.dim()
    reshape[0] = w.shape[0]
    return w * scale.to(torch.float32).reshape(reshape)


def pack_fp8(codes: torch.Tensor) -> torch.Tensor:
    """FP8 codes are already 1 byte each — no packing needed. Returns int8 view."""
    return codes.view(torch.int8)


def unpack_fp8(packed: torch.Tensor) -> torch.Tensor:
    """Unpack int8 view back to uint8 codes."""
    return packed.view(torch.uint8)