"""Double quantization — quantize the per-group scales themselves to int8.

QLoRA paper (Dettmers et al., 2023) observes that the per-group fp16/fp32 scales
of NF4 have their own non-trivial storage overhead (~32 bit per group). Double
quantization compresses the scales: each block of `block_size` consecutive
scales (within one output row) gets one fp32 `block_scale`
(= max(abs(scales_in_block)) / 127), and the scales inside the block are
stored as int8.

Storage reduction (per group):
  without double quant:  scale = fp32 = 32 bit/group
  with double quant:     scale_int8 = 8 bit/group + block_scale fp32 per block
                         = 8 + 32/block_size bit/group
                         for block_size=256: ~8.125 bit/group (3.9x reduction on scales)

The block_scale itself could be quantized further (2nd level), but QLoRA stops
at one level — diminishing returns + added complexity. We follow that.

Quantization is performed PER ROW (each output channel's scales quantized
independently), so dequant restores [out, num_groups] without cross-row mixing.

The last block in each row may be shorter than `block_size` (no padding) —
we store the true num_groups per row and trim on dequant. This avoids the
padding overhead that would otherwise dominate for small layers.
"""

import torch


def double_quantize_scales_2d(
    scales: torch.Tensor,
    block_size: int = 256,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize per-row scales [out, num_groups] to int8 + per-block fp32 scale.

    Each output row is quantized independently: scales in row `o` are split
    into blocks of up to `block_size` consecutive elements (the last block may
    be shorter — no padding is added), each block gets one fp32 block_scale,
    the scales inside the block are int8.

    Args:
        scales: fp32/fp16 tensor [out, num_groups] of positive per-group scales.
        block_size: max number of consecutive scales per double-quant block.

    Returns:
        (scales_int8, block_scale) where:
          scales_int8:  int8 [out, num_groups]            (no padding)
          block_scale:  fp32 [out, num_blocks_per_row]     (ceil(num_groups / block_size))
    """
    s = scales.detach().to(torch.float32)
    out_features, num_groups = s.shape

    # Number of blocks per row (last block may be short)
    num_blocks = (num_groups + block_size - 1) // block_size
    block_scale = torch.empty(out_features, num_blocks, dtype=torch.float32)
    scales_int8 = torch.empty(out_features, num_groups, dtype=torch.int8)

    for b in range(num_blocks):
        start = b * block_size
        end = min(start + block_size, num_groups)
        block = s[:, start:end]  # [out, block_len]
        block_max = block.abs().amax(dim=1).clamp(min=1e-12)  # [out]
        bs = block_max / 127.0
        block_scale[:, b] = bs
        q = torch.clamp(torch.round(block / bs.unsqueeze(1)), min=-127, max=127).to(torch.int8)
        scales_int8[:, start:end] = q

    return scales_int8, block_scale


def dequantize_scales_2d(
    scales_int8: torch.Tensor,
    block_scale: torch.Tensor,
    block_size: int = 256,
    num_groups: int | None = None,
) -> torch.Tensor:
    """Reconstruct per-row scales [out, num_groups] from int8 + per-block scale.

    Args:
        scales_int8: int8 [out, num_groups] (from double_quantize_scales_2d).
        block_scale: fp32 [out, num_blocks_per_row].
        block_size: max block size used in quantization.
        num_groups: original num_groups per row. If None, uses scales_int8.shape[1].

    Returns:
        fp32 tensor [out, num_groups].
    """
    out_features = scales_int8.shape[0]
    ng = scales_int8.shape[1] if num_groups is None else num_groups
    num_blocks = block_scale.shape[1]

    deq = torch.empty(out_features, scales_int8.shape[1], dtype=torch.float32)
    for b in range(num_blocks):
        start = b * block_size
        end = min(start + block_size, scales_int8.shape[1])
        block_i = scales_int8[:, start:end].to(torch.float32)
        bs = block_scale[:, b].to(torch.float32).unsqueeze(1)
        deq[:, start:end] = block_i * bs

    if num_groups is not None and num_groups < deq.shape[1]:
        deq = deq[:, :num_groups]
    return deq