"""CPU smoke test of the apex port on a tiny random Qwen3.5-MoE text model: REAP's observer with the
3D-expert hook, REAP's prune (keep half), save_pretrained, reload, and the saved weight names."""

import json
import tempfile
from pathlib import Path
from types import SimpleNamespace

import torch
from safetensors import safe_open
from transformers import AutoModelForCausalLM
from transformers.models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeTextConfig

from reap.main import _setup_observer
from reap.prune import prune

torch.manual_seed(0)
cfg = Qwen3_5MoeTextConfig(
    vocab_size=128, hidden_size=64, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2,
    head_dim=16, moe_intermediate_size=32, shared_expert_intermediate_size=32, num_experts=8,
    num_experts_per_tok=2, layer_types=["linear_attention", "full_attention"], linear_num_key_heads=2,
    linear_num_value_heads=2, linear_key_head_dim=16, linear_value_head_dim=16,
)
model = AutoModelForCausalLM.from_config(cfg).eval().cuda()
print(type(model).__name__)
observer = _setup_observer(model, SimpleNamespace(renormalize_router_weights=True, record_pruning_metrics_only=True))
ids = torch.randint(0, 128, (1, 48)).cuda()
mask = torch.ones_like(ids)
mask[0, 20] = 0
with torch.no_grad(), observer.set_attention_mask(mask):
    model(input_ids=ids, attention_mask=mask)
state = {k: {m: (v.mean if hasattr(v, "mean") and not torch.is_tensor(v) else v) for m, v in s.items()} for k, s in observer.state.items()}
print("layers observed", list(state), "reap[0]", [round(float(x), 4) for x in state[0]["reap"]], "freq[0]", state[0]["expert_frequency"].tolist())
observer.close_hooks()
with tempfile.TemporaryDirectory() as d:
    out = prune(state, model, SimpleNamespace(prune_method="reap", perserve_super_experts=False, perserve_outliers=False), 4, Path(d) / "pruned")
    keys = []
    for f in Path(out).glob("*.safetensors"):
        with safe_open(f, "pt") as fh:
            keys += list(fh.keys())
    print("saved config num_experts", json.load(open(Path(out) / "config.json"))["num_experts"])
    print("saved expert keys", sorted(k for k in keys if "experts" in k)[:6], "...", len([k for k in keys if "experts." in k]))
    re = AutoModelForCausalLM.from_pretrained(out).eval().cuda()
    print("reloaded experts", re.model.layers[0].mlp.experts.gate_up_proj.shape, re.model.layers[0].mlp.gate.weight.shape)
    with torch.no_grad():
        print("reloaded forward ok", tuple(re(input_ids=ids).logits.shape))