"""Self-contained STQ1 packing for 3:4 ternary Sherry weights."""

from __future__ import annotations

from dataclasses import dataclass

import torch


_DTYPE_NAMES = {
    torch.float16: "float16",
    torch.bfloat16: "bfloat16",
    torch.float32: "float32",
    torch.float64: "float64",
}
_NAME_DTYPES = {name: dtype for dtype, name in _DTYPE_NAMES.items()}


@dataclass(frozen=True)
class PackedSTQ1:
    """Packed 5-bit codes, fp16 group scales, and minimal tensor metadata."""

    data: torch.Tensor
    scales: torch.Tensor
    shape: tuple[int, int]
    group_size: int
    original_dtype: str

    @property
    def num_weights(self) -> int:
        """Number of logical weights represented by the payload."""

        return self.shape[0] * self.shape[1]

    @property
    def effective_bits(self) -> float:
        """Actual stored bits per weight, including fp16 scales and byte padding."""

        return (self.data.numel() * 8 + self.scales.numel() * 16) / self.num_weights

    def to_bytes(self) -> bytes:
        """Return the packed 5-bit code stream as bytes."""

        return bytes(self.data.cpu().tolist())


def _validate_fakequant(weight: torch.Tensor, group_size: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Validate fixed-rate STQ1 invariants absent from the original repository."""

    if weight.ndim != 2 or not weight.is_floating_point():
        raise ValueError("w_fakequant must be a floating-point 2-D tensor")
    if group_size <= 0 or group_size % 4:
        raise ValueError("group_size must be positive and divisible by 4")
    if weight.shape[1] % group_size:
        raise ValueError("in_features must be divisible by group_size")
    if weight.dtype not in _DTYPE_NAMES:
        raise ValueError(f"unsupported fake-quant dtype: {weight.dtype}")

    cpu_weight = weight.detach().cpu().contiguous()
    blocks = cpu_weight.reshape(-1, 4)
    zero_mask = blocks == 0
    zero_counts = zero_mask.sum(dim=-1)
    if not torch.all(zero_counts == 1):
        raise ValueError("each four-weight block must contain exactly one zero")

    grouped = cpu_weight.reshape(-1, group_size)
    nonzero_abs = grouped.abs().masked_fill(grouped == 0, float("nan"))
    scales = torch.nanmean(nonzero_abs.float(), dim=-1)
    expected = scales[:, None].expand_as(grouped)
    survivor_error = (grouped.abs().float() - expected).masked_fill(grouped == 0, 0)
    if not torch.all(survivor_error == 0):
        raise ValueError("all nonzero weights in a group must have one exact shared scale")
    fp16_scales = scales.to(torch.float16)
    canonical_scales = fp16_scales.to(cpu_weight.dtype)
    canonical = canonical_scales[:, None].expand_as(grouped)
    if not torch.all((grouped.abs() == canonical) | (grouped == 0)):
        raise ValueError(
            "fake-quant scales are not exactly representable as fp16; export an fp16 "
            "fake-quant tensor before packing"
        )
    return blocks, fp16_scales


def pack(w_fakequant: torch.Tensor, group_size: int = 128) -> PackedSTQ1:
    """Pack a canonical fake-quant matrix into 5 bits per four weights.

    Bit layout is little-endian within the stream.  Each 5-bit code stores the
    zero lane in bits 0-1, followed by three sign bits in ascending nonzero-lane
    order (one means positive).  Scales are stored once per ``group_size`` as
    fp16, making the theoretical cost ``1.25 + 16/group_size`` bits/weight.

    AngelSlim supplies no exporter.  This implementation validates the 3:4 and
    shared-scale invariants instead of silently packing arbitrary latent weights.
    To keep the roundtrip genuinely bit-exact, inputs whose group scales cannot
    be represented as fp16 are rejected with an instruction to cast the export
    fake-quant tensor first.
    """

    blocks, scales = _validate_fakequant(w_fakequant, group_size)
    zero_pos = (blocks == 0).to(torch.int64).argmax(dim=1)
    codes = zero_pos.clone()
    for lane in range(4):
        sign_offset = lane - (lane > zero_pos).to(torch.int64)
        sign_bits = (blocks[:, lane] > 0).to(torch.int64)
        sign_bits = torch.where(lane == zero_pos, 0, sign_bits)
        codes |= sign_bits << (2 + sign_offset)

    logical_codes = int(codes.numel())
    padding = (-logical_codes) % 8
    if padding:
        codes = torch.cat((codes, torch.zeros(padding, dtype=torch.int64)))
    code_groups = codes.reshape(-1, 8)
    words = torch.zeros(code_groups.shape[0], dtype=torch.int64)
    for index in range(8):
        words |= code_groups[:, index] << (5 * index)
    payload = torch.empty((words.numel(), 5), dtype=torch.uint8)
    for index in range(5):
        payload[:, index] = ((words >> (8 * index)) & 0xFF).to(torch.uint8)
    payload = payload.reshape(-1)[: (logical_codes * 5 + 7) // 8].contiguous()

    return PackedSTQ1(
        data=payload,
        scales=scales.contiguous(),
        shape=(w_fakequant.shape[0], w_fakequant.shape[1]),
        group_size=group_size,
        original_dtype=_DTYPE_NAMES[w_fakequant.dtype],
    )


def unpack(packed: PackedSTQ1) -> torch.Tensor:
    """Decode an STQ1 object to its bit-exact canonical fake-quant matrix.

    This inverse has no counterpart in AngelSlim.  It decodes only the stored
    5-bit symbols and fp16 scales; no latent weights, masks, or calibration state
    are needed.  The returned dtype matches the input metadata.
    """

    if packed.data.dtype != torch.uint8 or packed.data.ndim != 1:
        raise ValueError("packed data must be a one-dimensional uint8 tensor")
    if packed.scales.dtype != torch.float16 or packed.scales.ndim != 1:
        raise ValueError("packed scales must be a one-dimensional float16 tensor")
    rows, columns = packed.shape
    if rows <= 0 or columns <= 0 or columns % packed.group_size:
        raise ValueError("invalid packed shape or group size")
    expected_scales = rows * columns // packed.group_size
    if packed.scales.numel() != expected_scales:
        raise ValueError("scale count does not match packed metadata")
    if packed.original_dtype not in _NAME_DTYPES:
        raise ValueError(f"unsupported original dtype metadata: {packed.original_dtype}")

    number_of_blocks = rows * columns // 4
    payload = packed.data.cpu()
    byte_padding = (-payload.numel()) % 5
    if byte_padding:
        payload = torch.cat((payload, torch.zeros(byte_padding, dtype=torch.uint8)))
    byte_groups = payload.reshape(-1, 5).to(torch.int64)
    words = torch.zeros(byte_groups.shape[0], dtype=torch.int64)
    for index in range(5):
        words |= byte_groups[:, index] << (8 * index)
    codes = torch.empty((words.numel(), 8), dtype=torch.int64)
    for index in range(8):
        codes[:, index] = (words >> (5 * index)) & 0x1F
    codes = codes.reshape(-1)[:number_of_blocks]

    zero_pos = codes & 0b11
    values = torch.empty((number_of_blocks, 4), dtype=torch.float16)
    for lane in range(4):
        sign_offset = lane - (lane > zero_pos).to(torch.int64)
        positive = (codes >> (2 + sign_offset)) & 1
        signed = torch.where(positive != 0, 1.0, -1.0).to(torch.float16)
        values[:, lane] = torch.where(lane == zero_pos, 0.0, signed)

    scale_per_weight = packed.scales.repeat_interleave(packed.group_size)
    decoded = values.reshape(-1) * scale_per_weight
    return decoded.reshape(rows, columns).to(_NAME_DTYPES[packed.original_dtype])