import mlx.core as mx
import numpy as np

from h3_mlx.config import DiTConfig
from h3_mlx.dit import DiTBlock
from h3_mlx.weights import WeightStore


def test_packed_block_cache_and_compiled_bf16_boundaries(tmp_path):
    mx.random.seed(729)
    c = DiTConfig(
        hidden_size=32, num_attention_heads=2, attention_head_dim=16, ffn_hidden_size=64
    )
    tensors = {}
    for name, shape in {
        "norm1": (32,),
        "norm2": (32,),
        "attn.q_norm": (16,),
        "attn.k_norm": (16,),
        "attn.qkv_proj": (96, 32),
        "attn.out_proj": (32, 32),
        "mlp.fc1": (128, 32),
        "mlp.fc2": (32, 64),
    }.items():
        tensors[f"blocks.0.{name}.weight"] = (mx.random.normal(shape) * 0.02).astype(
            mx.bfloat16
        )
        if len(shape) == 2:
            tensors[f"blocks.0.{name}.bias"] = (
                mx.random.normal((shape[0],)) * 0.01
            ).astype(mx.bfloat16)
    original = tmp_path / "original"
    original.mkdir()
    mx.save_safetensors(str(original / "model.safetensors"), tensors)
    store = WeightStore(original)
    native = DiTBlock(store, "blocks.0", c)
    packed = DiTBlock(store, "blocks.0", c, packed_cache=tmp_path / "packed")
    restored = DiTBlock(store, "blocks.0", c, packed_cache=tmp_path / "packed")
    x = mx.random.normal((17, 32)).astype(mx.bfloat16)
    params = tuple(
        (mx.random.normal((3, 32)) * 0.1).astype(mx.bfloat16) for _ in range(6)
    )
    ids = mx.arange(17) % 3
    expected = native(x, params=params, indices=ids, mlp_chunk=0)
    state = tuple(
        p.weight.state()
        for p in (restored.qkv, restored.out, restored.fc1, restored.fc2)
    )

    def forward(x, params, ids, state):
        return restored(
            x, params=params, indices=ids, mlp_chunk=0, projection_weights=state
        )

    compiled = mx.compile(forward)
    for actual in (
        packed(x, params=params, indices=ids, mlp_chunk=0),
        restored(x, params=params, indices=ids, mlp_chunk=0),
        compiled(x, params, ids, state),
    ):
        np.testing.assert_array_equal(
            np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16))
        )