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

from h3_mlx.weights import WeightStore


@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16, mx.float32])
def test_selected_rows_preserve_order_duplicates_and_bypass_full_loader(
    tmp_path, monkeypatch, dtype
):
    values = mx.arange(128, dtype=mx.float32).reshape(16, 8).astype(dtype)
    path = tmp_path / "embedding.safetensors"
    mx.save_safetensors(str(path), {"embedding": values})
    ids = [15, 0, 3, 3, 1, 0]
    reference = np.asarray(values.astype(mx.float32))[ids]
    store = WeightStore(path, reader="mlx", cache_gb=1)

    def forbidden(*_args, **_kwargs):
        raise AssertionError("Row reads must not materialize the vocabulary")

    monkeypatch.setattr(mx, "load", forbidden)
    selected = store.get("embedding", rows=ids)
    np.testing.assert_array_equal(np.asarray(selected.astype(mx.float32)), reference)
    assert selected.dtype == dtype
    assert not store.cache
    assert store.bytes_read == len(ids) * 8 * values.itemsize


def test_selected_bf16_rows_preserve_payload_bits(tmp_path):
    raw = mx.array(
        [0, 0x8000, 0x7FC1, 0xFFFF, 1, 0x3F80, 0x7F80, 0xFF80], mx.uint16
    ).reshape(4, 2)
    path = tmp_path / "payload.safetensors"
    mx.save_safetensors(str(path), {"embedding": raw.view(mx.bfloat16)})
    selected = WeightStore(path).get("embedding", rows=[3, 0, 1, 1])
    np.testing.assert_array_equal(
        np.asarray(selected.view(mx.uint16)), np.asarray(raw)[[3, 0, 1, 1]]
    )