#!/usr/bin/env python3
"""D-57 · the teacher's routed experts quantised by our own trainer grids (round-to-nearest with the trainer's scale
search, no training) and written as a GGUF: the QAT starting point, measured in llama.cpp beside its own PTQ files.

    python scripts/rtn_gguf.py --method iq4xs --out teacher-ours-iq4xs-rtn.gguf
    python scripts/rtn_gguf.py --ckpt runs/<run>/full_N/state.pt --method int2 --down-method iq4nl --out x.gguf

With --ckpt, a trained student instead: its fp32 expert latents re-quantised (gate/up with --method, the down
projection with --down-method, default the method's own) and its trained attention and norms; the frozen router
and embedding are the teacher's. The checkpoint's torchao optimiser state is skipped (gen has no torchao).

Experts go through `quant.ternary.export_packed` with `model.student.proj_spec` (so the down projection gets its
32-block grid), everything else is written as the teacher holds it (bf16; norms and the router f32), in the template
GGUF's tensor order; `export.gguf.check_round_trip` then confirms the GGUF dequantises to the packed weights.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch
from safetensors import safe_open

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.export.gguf import check_round_trip, gguf_module, student_tensors, write_gguf  # noqa: E402
from lowbit.model.student import proj_spec  # noqa: E402
from lowbit.quant.importance import load as load_importance  # noqa: E402
from lowbit.quant.ternary import GRIDS, METHODS, QuantSpec, export_packed  # noqa: E402

TEACHER = Path("/path/to/data/models/Marco-Nano-Instruct")
TEMPLATE = Path("/path/to/data/ptq/teacher-bf16.gguf")
LLAMA = Path("/path/to/llama.cpp")


def load_ckpt(path: str) -> dict:
    """torch.load of a full checkpoint with torchao's optimiser classes replaced by inert stubs."""
    import pickle
    import types

    class StubT(torch.Tensor):
        @classmethod
        def __torch_dispatch__(cls, func, types_, args=(), kwargs=None):
            raise RuntimeError("optimiser-state stub")

        def __setstate__(self, st):
            pass

    class StubO:
        def __init__(self, *a, **k):
            pass

        def __setstate__(self, st):
            pass

    class U(pickle.Unpickler):
        def find_class(self, mod, name):
            if mod.startswith("torchao"):
                return StubT if ("Tensor" in name or "State" in name or "Optim" in name) else StubO
            return super().find_class(mod, name)

    pm = types.SimpleNamespace(Unpickler=U, load=pickle.load, __name__="pickle")
    return torch.load(path, map_location="cpu", mmap=True, weights_only=False, pickle_module=pm)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--method", required=True, choices=METHODS)
    p.add_argument("--out", required=True)
    p.add_argument("--ckpt", help="a full checkpoint's state.pt: re-quantise its trained latents")
    p.add_argument("--down-method", choices=METHODS)
    p.add_argument("--importance", help="D-57 quant.importance safetensors: weight the scale search")
    p.add_argument("--dense-as", choices=["bf16", "q8_0"], default="bf16",
                   help="attention and embedding: bf16, or q8_0 (the phone build)")
    a = p.parse_args()
    gguf = gguf_module(LLAMA)
    def mk(m: str) -> QuantSpec:
        return QuantSpec(method=m, group_size=GRIDS[m].group if m in GRIDS else 128)
    spec = mk(a.method)
    down_spec = mk(a.down_method) if a.down_method else proj_spec(spec, "down")
    params = load_ckpt(a.ckpt)["params"] if a.ckpt else {}
    imp = load_importance(a.importance) if a.importance else {}
    cfg = json.loads((TEACHER / "config.json").read_text())
    wmap = json.loads((TEACHER / "model.safetensors.index.json").read_text())["weight_map"]
    files = {f: safe_open(TEACHER / f, "pt") for f in set(wmap.values())}

    def get(name: str) -> torch.Tensor:
        return files[wmap[name]].get_tensor(name)

    dense, packed = {}, {}
    for name in wmap:
        if ".mlp.experts." not in name:
            dense[name] = params[name] if name in params else get(name)
    for i in range(cfg["num_hidden_layers"]):
        pre = f"model.layers.{i}.mlp.experts"
        n = cfg["num_experts"]
        if params:
            gate_up, down = params[f"{pre}.gate_up_latent"], params[f"{pre}.down_latent"]
        else:
            gate = torch.stack([get(f"{pre}.{e}.gate_proj.weight") for e in range(n)])
            up = torch.stack([get(f"{pre}.{e}.up_proj.weight") for e in range(n)])
            down = torch.stack([get(f"{pre}.{e}.down_proj.weight") for e in range(n)])
            gate_up = torch.cat([gate, up], 1)
        packed[f"{pre}.gate_up_proj"] = export_packed(gate_up.float(), proj_spec(spec, "gate_up"),
                                                      importance=imp.get(f"layers.{i}.gate_up"))
        packed[f"{pre}.down_proj"] = export_packed(down.float(), down_spec, importance=imp.get(f"layers.{i}.down"))
        print(f"layer {i} quantised", flush=True)
    info = write_gguf(TEMPLATE, a.out, student_tensors(dense, packed, gguf, dense_as=a.dense_as), gguf)
    rt = check_round_trip(a.out, packed, gguf)
    print(json.dumps({**info, "bit_exact": rt["bit_exact"], "mismatched": rt["mismatched_tensors"][:3]}, indent=1))


if __name__ == "__main__":
    main()