import mlx.core as mx
import numpy as np
import pytest

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


def test_decoder_preserves_all_bf16_payloads_in_an_all_escape_fixture():
    # A deliberately uncompressed synthetic fixture: no packing algorithm.
    raw = mx.arange(65536, dtype=mx.uint16)
    packed = PackedBF16(
        (65536,),
        mx.zeros(65536, mx.uint8),
        mx.zeros(32768, mx.uint8),
        mx.arange(65536, dtype=mx.uint32),
        raw,
        mx.array([56], mx.uint32),
    )
    np.testing.assert_array_equal(
        np.asarray(packed.unpack().view(mx.uint16)), np.asarray(raw)
    )


def test_decoder_common_path_escapes_signs_and_odd_length():
    packed = PackedBF16(
        (7,),
        mx.array([128, 128, 1, 0, 128, 255, 0], mx.uint8),
        mx.array([0xF7, 0, 0, 0], mx.uint8),
        mx.array([2, 3, 4, 5, 6], mx.uint32),
        mx.array([0x3801, 0x8000, 0x7F80, 0xFFFF, 0], mx.uint16),
        mx.array([56], mx.uint32),
    )
    np.testing.assert_array_equal(
        np.asarray(packed.unpack().view(mx.uint16)),
        [0x3F80, 0xBF80, 0x3801, 0x8000, 0x7F80, 0xFFFF, 0],
    )
    zero_base = PackedBF16(
        (3,),
        mx.array([0, 0, 1], mx.uint8),
        mx.array([0x80, 1], mx.uint8),
        mx.array([], mx.uint32),
        mx.array([], mx.uint16),
        mx.array([0], mx.uint32),
    )
    np.testing.assert_array_equal(
        np.asarray(zero_base.unpack().view(mx.uint16)), [0, 0x8000, 0x0101]
    )


def test_prepacked_linear_preserves_bias_epilogue_and_compiled_inputs():
    # Known encoding of a 4x8 matrix of BF16 ones.
    packed = PackedBF16(
        (4, 8),
        mx.full((32,), 128, mx.uint8),
        mx.full((16,), 0x77, mx.uint8),
        mx.array([], mx.uint32),
        mx.array([], mx.uint16),
        mx.array([56], mx.uint32),
    )
    bias = mx.array([0.25, -0.5, 1, 0], mx.bfloat16)
    x = mx.arange(24).reshape(3, 8).astype(mx.bfloat16)
    layer = PackedLinear(packed, bias)
    reference = Linear(mx.ones((4, 8), mx.bfloat16), bias)(x)
    compiled = mx.compile(lambda x, state: layer(x, state=state))
    for actual in (layer(x), compiled(x, packed.state())):
        np.testing.assert_array_equal(
            np.asarray(actual.view(mx.uint16)), np.asarray(reference.view(mx.uint16))
        )


def test_packed_shape_validation():
    tensors = {
        "w.low": mx.zeros(3, mx.uint8),
        "w.nibbles": mx.zeros(1, mx.uint8),
        "w.escape_indices": mx.array([], mx.uint32),
        "w.escape_values": mx.array([], mx.uint16),
        "w.base": mx.array([0], mx.uint32),
    }
    with pytest.raises(ValueError, match="dimensions"):
        PackedBF16.from_tensors(tensors, "w", (3,))