"""fp4 — FP4 (E2M1) 4-bit floating-point format primitives.

FP4 E2M1: 1 sign + 2 exponent + 1 mantissa. Bias=1. 16 levels (including
signed zero). Per-group FP8 (E4M3) scale (group_size=16) + F32 global scale.

NVFP4 (NVIDIA): weights AND activations in FP4, per-16 FP8 scales, F32 global
weight_scale_2 + input_scale. Storage: 0.5 byte/weight + 1 byte FP8 scale per
16 weights + 4 bytes global.

MXFP4 (OCP MX): FP4 E2M1 + E8M0 (8-bit exponent only) shared block scale
(block_size=32). Single scale per block.
"""

from __future__ import annotations

import torch


# FP4 E2M1: sign(1) + exp(2) + mant(1). bias=1.
# 16 levels: code 0..15, bit 3=sign, bits 2-1=exp, bit 0=mant.
# Special: code 0 = +0, code 8 = -0.
# exp=0 (codes 0,1,8,9): subnormal: sign * (mant/2) * 2^(1-bias) = sign * mant/2 * 2^0 = sign*mant/2
#   code 0: +0, code 1: +0.5, code 8: -0, code 9: -0.5
# exp=1 (codes 2,3,10,11): sign * (1 + mant/2) * 2^(1-1) = sign * (1+mant/2)
#   code 2: +1.0, code 3: +1.5, code 10: -1.0, code 11: -1.5
# exp=2 (codes 4,5,12,13): sign * (1+mant/2) * 2^(2-1) = sign * (1+mant/2)*2
#   code 4: +2.0, code 5: +3.0, code 12: -2.0, code 13: -3.0
# exp=3 (codes 6,7,14,15): sign * (1+mant/2) * 2^(3-1) = sign * (1+mant/2)*4
#   code 6: +4.0, code 7: +6.0, code 14: -4.0, code 15: -6.0
FP4_E2M1_LUT: torch.Tensor = torch.tensor([
    0.0,    # 0: +0
    0.5,    # 1: +0.5 (subnormal)
    1.0,    # 2: +1.0
    1.5,    # 3: +1.5
    2.0,    # 4: +2.0
    3.0,    # 5: +3.0
    4.0,    # 6: +4.0
    6.0,    # 7: +6.0
    -0.0,   # 8: -0
    -0.5,   # 9: -0.5
    -1.0,   # 10: -1.0
    -1.5,   # 11: -1.5
    -2.0,   # 12: -2.0
    -3.0,   # 13: -3.0
    -4.0,   # 14: -4.0
    -6.0,   # 15: -6.0
], dtype=torch.float32)


def quantize_fp4(w: torch.Tensor, group_size: int = 16) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize weights to FP4 E2M1 indices + per-group absmax scale.

    Args:
        w: float [out, in] (or flat). in is padded to be divisible by group_size.
        group_size: elements per scale group (NVFP4=16, MXFP4=32).

    Returns:
        (idx, scale): idx int64 [out, in_padded] values [0,15], scale fp32 [out, num_groups].
    """
    w = w.detach().float()
    out_features = w.shape[0]
    in_features = w.shape[1] if w.dim() > 1 else w.numel()
    if w.dim() > 1:
        flat = w
    else:
        flat = w.reshape(1, -1)
        out_features = 1
    pad = (group_size - (flat.shape[1] % group_size)) % group_size
    if pad > 0:
        flat = torch.nn.functional.pad(flat, (0, pad))
    in_padded = flat.shape[1]
    num_groups = in_padded // group_size
    grouped = flat.reshape(out_features, num_groups, group_size)
    # Per-group absmax scale.
    scale = grouped.abs().amax(dim=2).clamp(min=1e-8)
    scale_exp = scale.unsqueeze(2).expand_as(grouped)
    w_norm = grouped / scale_exp
    # Nearest FP4 level.
    lut = FP4_E2M1_LUT.to(w_norm.device)
    diff = w_norm.unsqueeze(-1) - lut
    idx = diff.abs().argmin(dim=-1).to(torch.int64)
    idx = idx.reshape(out_features, in_padded)
    return idx, scale


def dequantize_fp4(
    idx: torch.Tensor, scale: torch.Tensor, group_size: int = 16, in_features: int | None = None
) -> torch.Tensor:
    """Reconstruct float weight from FP4 indices + per-group scale."""
    out_features = idx.shape[0]
    in_padded = idx.shape[1]
    num_groups = in_padded // group_size
    lut = FP4_E2M1_LUT.to(idx.device)
    w_norm = lut[idx]  # [out, in_padded]
    w_grouped = w_norm.reshape(out_features, num_groups, group_size)
    scale_exp = scale.to(torch.float32).unsqueeze(2).expand_as(w_grouped)
    w_deq = (w_grouped * scale_exp).reshape(out_features, in_padded)
    if in_features is not None and in_features < in_padded:
        w_deq = w_deq[:, :in_features]
    return w_deq


def pack_fp4(idx: torch.Tensor) -> torch.Tensor:
    """Pack two 4-bit FP4 indices into one int8 byte."""
    assert idx.shape[-1] % 2 == 0, "in_features must be even for packing"
    i = idx.to(torch.int16)
    packed = (i[..., 0::2] << 4) | (i[..., 1::2] & 0x0F)
    return packed.to(torch.int8)


def unpack_fp4(packed: torch.Tensor) -> torch.Tensor:
    """Unpack int8 bytes into two 4-bit FP4 indices."""
    p = packed.to(torch.int16)
    high = (p >> 4) & 0x0F
    low = p & 0x0F
    out = torch.stack([high, low], dim=-1).reshape(*p.shape[:-1], -1)
    return out.to(torch.int64)


# E8M0: 8-bit exponent only (no sign, no mantissa). Used as MXFP block scale.
# Value = 2^(code - 127). code 0..255.
def _make_e8m0_lut() -> torch.Tensor:
    lut = torch.zeros(256, dtype=torch.float64)
    for code in range(256):
        lut[code] = 2.0 ** (code - 127)
    return lut.float()


E8M0_LUT: torch.Tensor = _make_e8m0_lut()