import numpy as np, mlx.core as mx, pytest
from titan.registry import load_spec
from titan.store import ExpertStore

def test_registry_olmoe(model_dir):
    s = load_spec(model_dir)
    if s.model_type != "olmoe":
        pytest.skip("OLMoE-specific constants")
    assert s.num_layers == 16 and s.num_experts == 64 and s.top_k == 8
    assert s.quantization == {"group_size": 64, "bits": 4}
    assert s.bytes_per_expert == 3538944
    assert len(s.experts) == 1024
    assert s.expert_bytes_total == 1024 * 3538944

def test_store_reads_match_mx_load(model_dir):
    """Bytes read with pread at registry offsets must equal what mx.load returns for the same tensor."""
    s = load_spec(model_dir)
    if s.layout.startswith("stacked"):
        pytest.skip("stacked layout covered by test_registry_stacked.py")
    # with the TITAN contiguous layout present, this also verifies the conversion byte-for-byte against the original file
    st = ExpertStore(s)
    ref = mx.load(f"{model_dir}/model.safetensors")
    for (l, e) in [(0, 0), (7, 33), (15, 63)]:
        t = st.read_expert(l, e)
        for name, arr in t.items():
            m, part = name.split(".")
            r = ref[f"model.layers.{l}.mlp.experts.{e}.{m}_proj.{part}"]
            assert arr.shape == r.shape and arr.dtype == r.dtype
            assert bool(mx.array_equal(arr, r)), (l, e, name)
    assert st.stats()["bytes_read"] == 3 * s.bytes_per_expert
    if "titan_contiguous" in s.layout:
        assert st.stats()["reads"] == 3  # one pread per expert
    st.close()