"""T2.2 student model on CPU, with a tiny random Qwen3-MoE (2 layers, 8 experts, group 128)."""
import copy
from pathlib import Path

import pytest
import torch

from lowbit.model import student as st
from lowbit.quant.ternary import QuantSpec, reference_codes_and_scale

REPO = Path(__file__).resolve().parent.parent
A0 = REPO / "configs" / "precision" / "A0.yaml"


def tiny(dtype=torch.float32, seed=0, impl="eager", hidden=128):
    """hidden=256 gives gate/up rows Q3_K can hold (D-51's int3)."""
    from transformers import AutoModelForCausalLM, Qwen3MoeConfig
    cfg = Qwen3MoeConfig(vocab_size=256, hidden_size=hidden, intermediate_size=256, moe_intermediate_size=128,
                         num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, head_dim=32,
                         num_experts=8, num_experts_per_tok=2, norm_topk_prob=True, max_position_embeddings=512,
                         tie_word_embeddings=True, decoder_sparse_step=1, mlp_only_layers=[])
    torch.manual_seed(seed)
    m = AutoModelForCausalLM.from_config(cfg, dtype=dtype, attn_implementation="sdpa", experts_implementation=impl)
    with torch.no_grad():  # random init is near zero; give the experts realistic, distinct magnitudes
        for layer in m.model.layers:
            layer.mlp.experts.gate_up_proj.normal_(0, 0.05)
            layer.mlp.experts.down_proj.normal_(0, 0.05)
    return m.eval()


def ids(b=2, t=48, seed=1):
    return torch.randint(0, 256, (b, t), generator=torch.Generator().manual_seed(seed))


def logits_of(model, x, **kw):
    with torch.no_grad():
        return st.logits(model, st.final_hidden(model, x, **kw)).float()


def test_quant_off_equals_teacher_bitwise():
    teacher = tiny()
    prec = st.Precision.load(A0)
    prec.router_fp32 = True  # the fp32-router path: with fp32 tiny weights the two computations coincide
    student = st.build_student(copy.deepcopy(teacher), prec, quant=False)
    x = ids()
    assert torch.equal(logits_of(student, x), teacher(x).logits.float())


def test_quant_off_bf16_router_matches_hf():
    teacher = tiny(torch.bfloat16)
    prec = st.Precision.load(A0)
    assert not prec.router_fp32  # D-28: A0's router is computed as stock Marco's
    student = st.build_student(copy.deepcopy(teacher), prec, quant=False)
    x = ids()
    with torch.no_grad():
        assert torch.equal(logits_of(student, x), teacher(x).logits.float())


def test_quant_on_matches_reference_path():
    base = tiny()
    student = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
    ref = copy.deepcopy(base)
    spec = QuantSpec()
    with torch.no_grad():
        for layer in ref.model.layers:
            for name in ("gate_up_proj", "down_proj"):
                w = getattr(layer.mlp.experts, name)
                flat = w.reshape(-1, w.shape[-1])
                codes, scales = reference_codes_and_scale(flat, spec)
                w_hat = (codes.float().unflatten(-1, (-1, 128)) * scales.unsqueeze(-1)).flatten(-2)
                w.copy_(w_hat.reshape(w.shape))
    x = ids()
    assert torch.equal(logits_of(student, x), ref(x).logits.float())


def test_packed_reload_reproduces_fake_quant(tmp_path):
    prec = st.Precision.load(A0)
    prec.router_fp32 = True  # exercise the fp32-router export and reload path
    student = st.build_student(tiny(torch.bfloat16), prec, quant=True)
    man = st.export_packed_student(student, tmp_path / "packed")
    assert man["router_fp32"] and len(man["expert_shapes"]) == 4
    packed = st.load_packed_student(tmp_path / "packed")
    x = ids()
    with torch.no_grad():
        assert torch.equal(logits_of(student, x), packed(x).logits.float())


def test_doc_masked_equals_separate_documents():
    student = st.build_student(tiny(), st.Precision.load(A0), quant=True)
    a, b = ids(1, 20, seed=2), ids(1, 28, seed=3)
    packed = torch.cat([a, b], 1)
    whole = logits_of(student, packed, mode="doc_masked", doc_starts=[[0, 20]])
    sep = torch.cat([logits_of(student, a), logits_of(student, b)], 1)
    assert torch.allclose(whole, sep, atol=1e-5)
    causal = logits_of(student, packed, mode="causal_packed")
    assert torch.allclose(causal[:, :20], sep[:, :20], atol=1e-5)  # the first document is the same either way
    assert not torch.allclose(causal[:, 20:], sep[:, 20:], atol=1e-3)  # the second sees the first


def test_gradients_reach_trainable_classes_only():
    student = st.build_student(tiny(), st.Precision.load(A0), quant=True)
    classes = st.param_classes(student)
    assert len(classes["experts"]) == 4 and not any(p.requires_grad for p in classes["router"] + classes["embedding"])
    x = ids(1, 32)
    h = st.final_hidden(student, x)
    st.logits(student, h).float().logsumexp(-1).mean().backward()
    for p in classes["experts"] + classes["attention"] + classes["norms"]:
        assert p.grad is not None and p.grad.abs().sum() > 0
    assert all(p.grad is None for p in classes["router"] + classes["embedding"])
    ex = student.model.layers[0].mlp.experts
    assert ex.gate_up_proj is None and ex.down_proj is None  # no bf16 expert weights kept between forwards


@pytest.mark.parametrize("early_stop", [True, False])
def test_checkpointed_backward_keeps_no_expert_weights(early_stop):
    student = st.build_student(tiny(torch.bfloat16), st.Precision.load(A0), quant=True)
    if early_stop:  # HF's default with non-reentrant checkpointing: the recompute stops before _post_hook
        student.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    else:
        st.enable_checkpointing(student)
    student.train()
    h = st.final_hidden(student, ids(1, 32))
    st.logits(student, h).float().logsumexp(-1).mean().backward()
    kept = [l.mlp.experts.gate_up_proj is not None for l in student.model.layers]
    assert any(kept) if early_stop else not any(kept)


@pytest.mark.parametrize("method", ["absmean", "twn"])
def test_compiled_quantiser_gives_the_eager_codes(method):
    from lowbit.quant.ternary import fake_quant
    spec = QuantSpec(method=method)
    w = tiny().model.layers[0].mlp.experts.gate_up_proj.detach().float()
    w = torch.cat([w, w * 37.3, w * 1e-3])  # several scale ranges

    def quant(x):
        return fake_quant(x, spec)

    with torch.no_grad():  # fp32 sums and division moved ~0.09% of codes under compile before the fix
        assert torch.equal(quant(w), torch.compile(quant, dynamic=False)(w))


def test_compiled_layers_match_eager():
    base = tiny()
    plain = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
    comp = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
    for m in (plain, comp):
        m.config._experts_implementation = "grouped_mm"
    st.compile_layers(comp)
    x = ids()
    assert torch.allclose(logits_of(plain, x), logits_of(comp, x), atol=1e-5)


def test_grouped_mm_matches_eager():
    base = tiny()
    eager = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
    grouped = st.build_student(copy.deepcopy(base), st.Precision.load(A0), quant=True)
    grouped.config._experts_implementation = "grouped_mm"
    x = ids()
    assert torch.allclose(logits_of(eager, x), logits_of(grouped, x), atol=1e-4)


def test_a2_ternary_attention_and_a4_int8_layers(tmp_path):
    for arm in ("A2", "A4"):
        prec = st.Precision.load(REPO / "configs" / "precision" / f"{arm}.yaml")
        if arm == "A4":
            prec.int8_layers = (1,)  # the tiny model has 2 layers
        stu = st.build_student(tiny(torch.bfloat16), prec, quant=True)
        x = ids()
        if arm == "A2":
            q = stu.model.layers[0].self_attn.q_proj
            q(torch.zeros(1, q.in_features, dtype=torch.bfloat16))
            assert q.weight is None  # rebuilt per forward, not kept
            assert len(st.param_classes(stu)["experts"]) == 4 + 4  # experts + q/o latents of 2 layers
        h = st.final_hidden(stu, ids(1, 32))
        st.logits(stu, h).float().logsumexp(-1).mean().backward()
        lat = q.weight_latent if arm == "A2" else stu.model.layers[1].mlp.experts.gate_up_latent
        assert lat.grad is not None and lat.grad.abs().sum() > 0  # the STE reaches the latents
        man = st.export_packed_student(stu, tmp_path / arm)
        assert man["linear_shapes" if arm == "A2" else "int8_shapes"]
        packed = st.load_packed_student(tmp_path / arm)
        with torch.no_grad():
            assert torch.equal(logits_of(stu, x), packed(x).logits.float()), arm