"""Stacked-layout registry/store check (runs only where a stacked checkpoint such as Qwen3-30B-A3B-4bit is cached)."""
import glob, os, pytest, mlx.core as mx
from titan.registry import load_spec
from titan.store import ExpertStore

def _dir():
    hits = glob.glob(os.path.expanduser("~/.cache/huggingface/hub/models--mlx-community--Qwen3-30B-A3B-Instruct-2507-4bit/snapshots/*"))
    return hits[0] if hits else None

@pytest.mark.skipif(_dir() is None, reason="Qwen3-30B-A3B-Instruct-2507-4bit not cached")
def test_stacked_slices_match_mx_load():
    d = _dir(); s = load_spec(d)
    assert s.layout == "stacked" and s.num_layers == 48 and s.num_experts == 128 and s.top_k == 8
    st = ExpertStore(s)
    for (l, e) in [(0, 0), (23, 77), (47, 127)]:
        t = st.read_expert(l, e)
        f = s.experts[(l, e)].tensors["gate.weight"].file
        ref = mx.load(f)
        for name, arr in t.items():
            m, part = name.split(".")
            r = ref[f"model.layers.{l}.mlp.switch_mlp.{m}_proj.{part}"][e]
            assert arr.shape == r.shape and arr.dtype == r.dtype, (name, arr.shape, r.shape)
            assert bool(mx.array_equal(arr, r)), (l, e, name)
    print("bytes_per_expert", s.bytes_per_expert, "expert_bytes_total_GB", s.expert_bytes_total / 1e9)