import math

import torch

from memory_native import RMSCounterLinear
from memory_native.actquant import (
    effective_bits,
    pack_int4,
    quantize_codes,
    stochastic_quantize,
    unpack_int4,
)


def test_pack_int4_roundtrip_and_halves_bytes():
    torch.manual_seed(0)
    codes = torch.randint(-7, 8, (133,), dtype=torch.int8)  # odd length on purpose
    packed = pack_int4(codes)
    assert packed.dtype == torch.uint8
    assert packed.numel() == (codes.numel() + 1) // 2          # ~0.5 byte/elem
    back = unpack_int4(packed, codes.numel())
    assert torch.equal(back, codes)


def test_quantizer_is_unbiased():
    torch.manual_seed(0)
    x = torch.randn(500, 32)
    # average many stochastic quantizations -> should converge to x (E[Q(x)|x]=x). 2000 draws put
    # the Monte-Carlo noise floor at ~2.3e-3, so a 5e-3 bound actually bites: a small systematic
    # bias (e.g. a rounding-direction or clamp error) would push this well past it. (The old 0.02
    # bound was ~5x the noise floor and would have passed a biased quantizer.)
    avg = torch.stack([stochastic_quantize(x, 4) for _ in range(2000)]).mean(0)
    assert (avg - x).abs().mean().item() < 5e-3


def test_codes_fit_bit_width():
    torch.manual_seed(0)
    x = torch.randn(8, 64)
    for bits in (3, 4, 8):
        codes, scale = quantize_codes(x, bits)
        lim = (1 << (bits - 1)) - 1
        assert int(codes.abs().max()) <= lim
        assert scale.shape == (8, 1)


def test_effective_bits_includes_scale():
    # one fp16 scale amortized over a row of length 32 -> +0.5 bit/elem
    assert abs(effective_bits(4, 32) - 4.5) < 1e-9


def test_low_bit_saved_activation_still_recovers_teacher():
    """The counter update needs only E[Q(x)|x]=x, so int8/4/3 saved activations must train
    as well as fp (the activation-memory lever from the deep-v2 research)."""
    def recover(bits):
        torch.manual_seed(0)
        n, N, C = 16, 256, 11
        ts = 0.25
        base = math.sqrt(3.0 / (2.0 * n))
        tw = torch.randint(-1, 2, (n, n)).float()
        x = torch.randn(N, n)
        y = x @ (ts * tw).t()
        lay = RMSCounterLinear(n, n, C=C, lr=0.005, lr_scale=0.0,
                               init_gain=ts / base, act_save_bits=bits).train()
        for _ in range(400):
            ((lay(x) - y) ** 2).mean().backward()
        with torch.no_grad():
            return ((lay(x) - y) ** 2).mean().item()

    fp = recover(0)
    for bits in (8, 4, 3):
        mse = recover(bits)
        assert mse < 5e-3, (bits, mse)
        assert mse <= fp + 2e-3, (bits, mse, fp)