import mlx.core as mx
import numpy as np

from h3_mlx.packed import PackedBF16, PackedLinear
from h3_mlx.weights import Linear


def test_palette12_preserves_all_bf16_bit_patterns():
    # Includes signs, subnormals, signed zero, Inf and every NaN payload.
    raw = mx.arange(65536, dtype=mx.uint16)
    packed = PackedBF16.pack(raw.view(mx.bfloat16))
    np.testing.assert_array_equal(
        np.asarray(packed.unpack().view(mx.uint16)), np.asarray(raw)
    )


def test_palette12_odd_length_zero_and_wide_exponents():
    for raw in (
        [0, 0, 0],
        [0, 0x8000, 1, 0x8001, 0x3F80, 0x7F80, 0xFFFF],
        [0x3800, 0x3B01, 0xBAFF, 0x0080, 0x8000],
    ):
        x = mx.array(raw, mx.uint16)
        packed = PackedBF16.pack(x.view(mx.bfloat16))
        np.testing.assert_array_equal(np.asarray(packed.unpack().view(mx.uint16)), raw)


def test_palette12_realistic_linear_is_exact_and_smaller():
    mx.random.seed(731)
    w = (mx.random.normal((256, 512)) * 0.02).astype(mx.bfloat16)
    b = mx.random.normal((256,)).astype(mx.bfloat16)
    x = mx.random.normal((71, 512)).astype(mx.bfloat16)
    original = Linear(w, b)
    packed = PackedLinear.from_linear(original)
    assert packed.nbytes < original.nbytes * 0.76
    np.testing.assert_array_equal(
        np.asarray(packed(x).view(mx.uint16)), np.asarray(original(x).view(mx.uint16))
    )
    tensors = packed.weight.tensors("weight")
    restored = PackedBF16.from_tensors(tensors, "weight", w.shape)
    np.testing.assert_array_equal(
        np.asarray(restored.unpack().view(mx.uint16)), np.asarray(w.view(mx.uint16))
    )