"""T0.3 wrapper: the ternary switch replaces routed experts by their fake-quant, and nothing else."""
import torch
from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM

from lowbit.eval.wrapper import routed_expert_params, ternarise_experts
from lowbit.quant.ternary import QuantSpec, fake_quant


def tiny_moe():
    torch.manual_seed(0)
    cfg = Qwen3MoeConfig(vocab_size=256, hidden_size=128, intermediate_size=256, moe_intermediate_size=128,
                         num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, head_dim=64,
                         num_experts=4, num_experts_per_tok=2, max_position_embeddings=64)
    m = Qwen3MoeForCausalLM(cfg)
    with torch.no_grad():
        for p in m.parameters():
            p.normal_(0, 0.02)
    return m


def test_ternarise_experts_only_touches_experts():
    m = tiny_moe()
    before = {n: p.detach().clone() for n, p in m.named_parameters()}
    expert_names = {n for n, _ in routed_expert_params(m)}
    assert expert_names, "no routed-expert weights found"
    spec = QuantSpec()
    n = ternarise_experts(m, spec)
    assert n == sum(before[k].numel() for k in expert_names)
    for name, p in m.named_parameters():
        if name in expert_names:
            assert torch.equal(p, fake_quant(before[name].float(), spec).to(p.dtype))
            g = p.reshape(-1, 128)  # each g128 group holds at most {-s, 0, +s}
            assert all(len(torch.unique(row.abs())) <= 2 for row in g[:16])
        else:
            assert torch.equal(p, before[name]), name