"""D-58: codebook formats against llama.cpp's own files (gguf-py's dequantisers are the reference)."""
from pathlib import Path

import numpy as np
import pytest
import torch

from lowbit.export.gguf import gguf_module
from lowbit.quant import codebook as cb

PTQ = Path("/path/to/data/ptq/gguf")
LLAMA = "/path/to/llama.cpp"
SOURCES = {  # format → (file, tensor)
    "IQ2_XXS": ("teacher-iq2_xxs-imat.gguf", "blk.5.ffn_up_exps.weight"),
    "IQ2_XS": ("teacher-iq2_xs-imat.gguf", "blk.5.ffn_up_exps.weight"),
    "IQ2_S": ("teacher-iq2_s-iq4nldown-imat.gguf", "blk.5.ffn_gate_exps.weight"),
    "IQ3_XXS": ("teacher-iq3_xxs-imat.gguf", "blk.5.ffn_gate_exps.weight"),
    "IQ3_S": ("teacher-iq3_s-iq4nldown-imat.gguf", "blk.5.ffn_gate_exps.weight"),
    "IQ4_XS": ("teacher-iq4_xs-imat.gguf", "blk.5.ffn_up_exps.weight"),
    "IQ4_NL": ("teacher-iq3_s-iq4nldown-imat.gguf", "blk.5.ffn_down_exps.weight"),
}
N = 4096  # blocks per test


@pytest.fixture(scope="module")
def gguf():
    return gguf_module(LLAMA)


def _blocks(gguf, name, n=N):
    path, tensor = SOURCES[name]
    if not (PTQ / path).exists():
        pytest.skip(f"{path} not on this machine")
    t = next(t for t in gguf.GGUFReader(str(PTQ / path)).tensors if t.name == tensor)
    fmt = cb.FORMATS[name]
    assert t.tensor_type == getattr(gguf.GGMLQuantizationType, name)
    return fmt, np.asarray(t.data).reshape(-1, fmt.nbytes)[:n].copy()


@pytest.mark.parametrize("name", list(SOURCES))
def test_fields_round_trip(gguf, name):
    fmt, b = _blocks(gguf, name)
    out = cb.write_blocks(fmt, b, cb.get_d(b), cb.get_s(fmt, b))
    assert np.array_equal(out, b)  # byte-identical


@pytest.mark.parametrize("name", list(SOURCES))
def test_decode_matches_ggml(gguf, name):
    fmt, b = _blocks(gguf, name)
    ref = gguf.quants.dequantize(b, getattr(gguf.GGMLQuantizationType, name)).reshape(len(b), fmt.block)
    v = cb.values(fmt, b, gguf)
    e = cb.eff_np(fmt, cb.get_d(b), cb.get_s(fmt, b))
    got = (v.astype(np.float32).reshape(len(b), -1, fmt.sub) * e[..., None]).reshape(len(b), -1)
    assert np.array_equal(got, ref)
    # the torch form at its start (dlog = 0, slat = s) is the same weights
    d0 = torch.from_numpy(cb.get_d(b).astype(np.float32))
    s0 = torch.from_numpy(cb.get_s(fmt, b).astype(np.float32))
    w = cb.weight(fmt, torch.from_numpy(v.reshape(-1)), d0, torch.zeros_like(d0),
                  s0.reshape(-1) if fmt.name != "IQ4_NL" else None)
    assert torch.equal(w, torch.from_numpy(ref.reshape(-1)))


@pytest.mark.parametrize("name", list(SOURCES))
def test_new_scales_dequantise_as_trained(gguf, name):
    fmt, b = _blocks(gguf, name)
    v = torch.from_numpy(cb.values(fmt, b, gguf).reshape(-1))
    d0 = torch.from_numpy(cb.get_d(b).astype(np.float32))
    g = torch.Generator().manual_seed(0)
    dlog = 0.05 * torch.randn(d0.shape, generator=g)
    slat = None
    if fmt.name != "IQ4_NL":
        s0 = torch.from_numpy(cb.get_s(fmt, b).astype(np.float32)).reshape(-1)
        slat = s0 + 1.5 * torch.randn(s0.shape, generator=g)
    w = cb.weight(fmt, v, d0, dlog, slat)
    d, s = cb.trained_fields(fmt, d0, dlog, slat)
    nb = cb.write_blocks(fmt, b, d, s.reshape(len(b), -1))
    ref = gguf.quants.dequantize(nb, getattr(gguf.GGMLQuantizationType, name)).reshape(-1)
    assert torch.equal(w, torch.from_numpy(ref))


def test_ste_reaches_scales():
    fmt = cb.FORMATS["IQ3_S"]
    v = torch.randint(-15, 16, (2, 256), dtype=torch.int8)
    d0 = torch.full((2, 1), 0.01)
    dlog = torch.zeros(2, 1, requires_grad=True)
    slat = torch.full((2, 8), 5.0, requires_grad=True)
    cb.weight(fmt, v, d0, dlog, slat).sum().backward()
    assert dlog.grad.abs().sum() > 0 and slat.grad.abs().sum() > 0


def _fake_init(gguf, tmp_path, cfg):
    """A cb_import-style directory for a tiny model, from real llama.cpp blocks (IQ3_S gate/up, IQ4_NL down)."""
    from safetensors.numpy import save_file
    import json
    E, I, H = cfg.num_experts, cfg.moe_intermediate_size, cfg.hidden_size
    f3, b3 = _blocks(gguf, "IQ3_S")
    fn, bn = _blocks(gguf, "IQ4_NL", 4 * N)
    need3, needn = E * I * (H // 256), E * H * (I // 32)
    assert len(b3) >= 2 * need3 and len(bn) >= needn
    t, formats = {}, {}
    for i in range(cfg.num_hidden_layers):
        for k, (fmt, src, rows, cols) in {"gate": (f3, b3[:need3], I, H), "up": (f3, b3[need3:2 * need3], I, H),
                                          "down": (fn, bn[:needn], H, I)}.items():
            key = f"layers.{i}.{k}"
            formats[key] = fmt.name
            t[f"{key}.blocks"] = src.reshape(E, rows, cols // fmt.block, fmt.nbytes).copy()
            t[f"{key}.v"] = cb.values(fmt, src, gguf).reshape(E, rows, cols)
            t[f"{key}.d"] = cb.get_d(src).reshape(E, rows, -1)
            t[f"{key}.s"] = cb.get_s(fmt, src).reshape(E, rows, -1)
    save_file(t, str(tmp_path / "experts.safetensors"))
    (tmp_path / "manifest.json").write_text(json.dumps({"formats": formats}))
    return t


def test_student_round_trip(gguf, tmp_path):
    """The codebook student's weights are the source blocks' dequantisation; after a training-like change the packed
    export dequantises to the student's weights; gate/up keep llama.cpp's order (gate first)."""
    from test_student import tiny
    from lowbit.model import student as st
    from lowbit.quant.ternary import QuantSpec
    from lowbit.export.gguf import packed_blocks, student_tensors
    base = tiny(hidden=256)
    (tmp_path / "init").mkdir()
    t = _fake_init(gguf, tmp_path / "init", base.config)
    prec = st.Precision.load(Path(__file__).resolve().parent.parent / "configs/precision/cb_iq3s.yaml")
    prec.codebook = str(tmp_path / "init")
    prec.experts = QuantSpec(method="codebook", group_size=32)
    m = st.build_student(base, prec, quant=True)
    T = gguf.GGMLQuantizationType
    ex = m.model.layers[1].mlp.experts
    deq = lambda b, q: gguf.quants.dequantize(b, q).reshape(*b.shape[:2], -1)
    ref_gu = np.concatenate([deq(t[f"layers.1.{k}.blocks"], T.IQ3_S) for k in ("gate", "up")], 1)
    assert np.array_equal(st._expert_weight(ex, "gate_up").detach().numpy(), ref_gu)
    ref_d = deq(t["layers.1.down.blocks"], T.IQ4_NL)
    assert np.array_equal(st._expert_weight(ex, "down").detach().numpy(), ref_d)
    # forward/backward reach the scales
    m.train()
    x = torch.randint(0, 256, (2, 32), generator=torch.Generator().manual_seed(1))
    st.logits(m, st.final_hidden(m, x)).float().pow(2).mean().backward()
    assert ex.gate_up_dlog.grad.abs().sum() > 0 and ex.gate_up_slat.grad.abs().sum() > 0
    assert ex.down_dlog.grad.abs().sum() > 0 and not hasattr(ex, "down_slat")
    # untrained export: the source blocks unchanged, gate first
    st.export_packed_student(m, tmp_path / "p0")
    sd, packed, man = st.read_packed(tmp_path / "p0")
    tens = student_tensors({}, packed, gguf)
    for k in ("gate", "up", "down"):
        assert tens[f"blk.1.ffn_{k}_exps.weight"][1].tobytes() == t[f"layers.1.{k}.blocks"].tobytes()
    # a training-like change: the export dequantises to the student's weights
    g = torch.Generator().manual_seed(2)
    with torch.no_grad():
        for layer in m.model.layers:
            e = layer.mlp.experts
            e.gate_up_dlog.add_(0.05 * torch.randn(e.gate_up_dlog.shape, generator=g))
            e.gate_up_slat.add_(1.5 * torch.randn(e.gate_up_slat.shape, generator=g))
            e.down_dlog.add_(0.05 * torch.randn(e.down_dlog.shape, generator=g))
    st.export_packed_student(m, tmp_path / "p1")
    sd, packed, man = st.read_packed(tmp_path / "p1")
    for name in ("gate_up", "down"):
        key = f"model.layers.1.mlp.experts.{name}_proj"
        qt, raw = packed_blocks(packed[key], gguf)
        got = gguf.quants.dequantize(raw, qt).reshape(packed[key]["shape"])
        assert np.array_equal(got, st._expert_weight(ex, name).detach().numpy())
        assert torch.equal(st.dequant_codebook(key, packed[key], torch.float32), st._expert_weight(ex, name).detach())
    assert man["codebook_formats"]["model.layers.0.mlp.experts.gate_up_proj"] == "IQ3_S"


FREE = ("IQ3_S", "IQ2_XXS", "IQ4_NL")


@pytest.mark.parametrize("name", FREE)
def test_free_codes_start_is_the_file(gguf, name):
    """Latents at llama.cpp's dequantised weights find llama.cpp's own codes: the re-encoded blocks are identical."""
    fmt, b = _blocks(gguf, name)
    ref = gguf.quants.dequantize(b, getattr(gguf.GGMLQuantizationType, name)).reshape(len(b), fmt.block)
    grid = cb.grid_from_gguf(fmt, gguf)
    d0 = torch.from_numpy(cb.get_d(b).astype(np.float32))[:, None]
    s0 = torch.from_numpy(cb.get_s(fmt, b).astype(np.float32))
    lat = torch.from_numpy(ref.copy())
    w = cb.weight_free(fmt, lat, d0, torch.zeros_like(d0), None if name == "IQ4_NL" else s0, grid)
    assert torch.equal(w, lat)
    e = cb.eff_torch(fmt, d0, s0).unsqueeze(-1)
    idx, vals = cb.free_codes(fmt, (lat.unflatten(-1, (-1, fmt.sub)) / e).flatten(-2), grid)
    out = cb.encode(fmt, idx, vals, cb.get_d(b), cb.get_s(fmt, b))
    assert np.array_equal(out.reshape(b.shape), b)


@pytest.mark.parametrize("name", FREE)
def test_free_codes_moved_dequantise_as_trained(gguf, name):
    fmt, b = _blocks(gguf, name)
    ref = gguf.quants.dequantize(b, getattr(gguf.GGMLQuantizationType, name)).reshape(len(b), fmt.block)
    grid = cb.grid_from_gguf(fmt, gguf)
    g = torch.Generator().manual_seed(3)
    lat = torch.from_numpy(ref.copy()) * (1 + 0.3 * torch.randn(ref.shape, generator=g))
    d0 = torch.from_numpy(cb.get_d(b).astype(np.float32))[:, None]
    dlog = 0.05 * torch.randn(d0.shape, generator=g)
    slat = None if name == "IQ4_NL" else torch.from_numpy(cb.get_s(fmt, b).astype(np.float32)) + torch.randn(len(b), fmt.block // fmt.sub, generator=g)
    w = cb.weight_free(fmt, lat, d0, dlog, slat, grid)
    d, s = cb.trained_fields(fmt, d0, dlog, slat)
    e = cb.eff_torch(fmt, torch.from_numpy(d.astype(np.float32)), torch.from_numpy(s.astype(np.float32))).unsqueeze(-1)
    idx, vals = cb.free_codes(fmt, (lat.unflatten(-1, (-1, fmt.sub)) / e).flatten(-2), grid)
    out = cb.encode(fmt, idx, vals, d, s)
    got = gguf.quants.dequantize(out.reshape(b.shape), getattr(gguf.GGMLQuantizationType, name)).reshape(w.shape)
    assert torch.equal(w, torch.from_numpy(got))
    assert (out.reshape(b.shape) != b).any()  # codes really moved


def test_student_free_round_trip(gguf, tmp_path):
    """Free codes: the start is the file; latents get gradients; a moved student exports to its own weights."""
    from test_student import tiny
    from lowbit.model import student as st
    from lowbit.quant.ternary import QuantSpec
    from lowbit.export.gguf import packed_blocks
    from safetensors.torch import save_file
    base = tiny(hidden=256)
    (tmp_path / "init").mkdir()
    t = _fake_init(gguf, tmp_path / "init", base.config)
    save_file({n: cb.grid_from_gguf(cb.FORMATS[n], gguf) for n in ("IQ3_S", "IQ4_NL")}, str(tmp_path / "init/grids.safetensors"))
    prec = st.Precision.load(Path(__file__).resolve().parent.parent / "configs/precision/cb_iq3s.yaml")
    prec.codebook, prec.codebook_mode = str(tmp_path / "init"), "free"
    prec.experts = QuantSpec(method="codebook", group_size=32)
    m = st.build_student(base, prec, quant=True)
    ex = m.model.layers[1].mlp.experts
    T = gguf.GGMLQuantizationType
    deq = lambda b, q: gguf.quants.dequantize(b, q).reshape(*b.shape[:2], -1)
    ref = np.concatenate([deq(t[f"layers.1.{k}.blocks"], T.IQ3_S) for k in ("gate", "up")], 1)
    assert np.array_equal(st._expert_weight(ex, "gate_up").detach().numpy(), ref)
    m.train()
    x = torch.randint(0, 256, (2, 32), generator=torch.Generator().manual_seed(1))
    st.logits(m, st.final_hidden(m, x)).float().pow(2).mean().backward()
    assert ex.gate_up_latent.grad.abs().sum() > 0 and ex.down_latent.grad.abs().sum() > 0
    assert ex.gate_up_slat.grad.abs().sum() > 0
    g = torch.Generator().manual_seed(4)
    with torch.no_grad():
        for layer in m.model.layers:
            e = layer.mlp.experts
            for n in ("gate_up", "down"):
                lat = getattr(e, f"{n}_latent")
                lat.mul_(1 + 0.3 * torch.randn(lat.shape, generator=g))
    st.export_packed_student(m, tmp_path / "p")
    sd, packed, man = st.read_packed(tmp_path / "p")
    for name in ("gate_up", "down"):
        key = f"model.layers.1.mlp.experts.{name}_proj"
        qt, raw = packed_blocks(packed[key], gguf)
        got = gguf.quants.dequantize(raw, qt).reshape(packed[key]["shape"])
        assert np.array_equal(got, st._expert_weight(ex, name).detach().numpy())


def test_codebook_export_resume(gguf, tmp_path):
    """A student loaded from its own export exports the same file again (blocks and dense tensors)."""
    from test_student import tiny
    from lowbit.model import student as st
    from lowbit.quant.ternary import QuantSpec
    from safetensors.torch import load_file
    base = tiny(hidden=256)
    (tmp_path / "init").mkdir()
    _fake_init(gguf, tmp_path / "init", base.config)
    prec = st.Precision.load(Path(__file__).resolve().parent.parent / "configs/precision/cb_iq3s.yaml")
    prec.codebook = str(tmp_path / "init")
    prec.experts = QuantSpec(method="codebook", group_size=32)
    m = st.build_student(base, prec, quant=True)
    g = torch.Generator().manual_seed(5)
    with torch.no_grad():
        for layer in m.model.layers:
            e = layer.mlp.experts
            e.gate_up_dlog.add_(0.05 * torch.randn(e.gate_up_dlog.shape, generator=g))
            e.gate_up_slat.add_(1.5 * torch.randn(e.gate_up_slat.shape, generator=g))
            e.down_dlog.add_(0.05 * torch.randn(e.down_dlog.shape, generator=g))
            layer.self_attn.q_proj.weight.add_(0.01)
    st.export_packed_student(m, tmp_path / "a")
    m2 = st.build_student(tiny(hidden=256), prec, quant=True)
    st.load_codebook_export(m2, tmp_path / "a")
    st.export_packed_student(m2, tmp_path / "b")
    a, b = load_file(str(tmp_path / "a/model.safetensors")), load_file(str(tmp_path / "b/model.safetensors"))
    assert a.keys() == b.keys() and all(torch.equal(a[k], b[k]) for k in a)