#!/usr/bin/env python3
"""PTQ baseline (D-56): what llama.cpp's own post-training quantisation of the teacher reaches at the bit widths of
our QAT students, measured the same way as the students themselves (llama.cpp KL divergence against the bf16 teacher).

Only the routed experts are quantised, as in our students; everything else stays bf16. Each width gets plain
round-to-nearest in the students' own GGUF type (= the QAT starting point) and llama.cpp's best importance-matrix
types near the same bits per weight.

    python scripts/ptq_baseline.py texts                  # eval.txt (held-out eval stream) + calib.txt (calib stream)
    python scripts/ptq_baseline.py imatrix                # teacher importance matrix on calib.txt (GPU)
    python scripts/ptq_baseline.py quantize [NAME..]      # the PTQ GGUFs (CPU)
    python scripts/ptq_baseline.py base                   # teacher logits for the KL runs (GPU)
    python scripts/ptq_baseline.py kld [NAME..]           # KL divergence of each GGUF against the teacher (GPU)
    python scripts/ptq_baseline.py report                 # reports/ptq_baseline.md + .json

Students go in as extra GGUFs: `kld --gguf name=path.gguf`.
"""
from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
from pathlib import Path

import numpy as np

REPO = Path(__file__).resolve().parent.parent
DATA = Path("/path/to/data")
OUT = DATA / "ptq"
TEACHER_HF = DATA / "models/Marco-Nano-Instruct"
TEACHER = OUT / "teacher-bf16.gguf"
BIN = Path("/path/to/llama.cpp/build-ptq/bin")
EVAL_SEQS, CALIB_SEQS = 40, 120   # eval: 40 × 4,096 tokens (82K scored); calib: 120 × 4,096 (~490K tokens)
CTX = 4096

EXPS = ("ffn_gate_exps", "ffn_up_exps", "ffn_down_exps")
# name → (gate/up type, down type, uses the importance matrix). down_exps rows are 384 wide, so the 256-block K and
# IQ types cannot hold them; those variants put down in the nearest 32-block type at a similar width instead.
VARIANTS = {
    # 4-bit: our int4 student is Q4_0 everywhere
    "q4_0-rtn":   ("q4_0", "q4_0", False),
    "q4_0-imat":  ("q4_0", "q4_0", True),
    "q4_k-imat":  ("q4_k", "q4_0", True),
    "iq4_xs-imat": ("iq4_xs", "iq4_nl", True),
    # 3-bit: our int3 student is Q3_K gate/up with Q4_0 down
    "q3_k-rtn":   ("q3_k", "q4_0", False),
    "q3_k-imat":  ("q3_k", "q4_0", True),
    "iq3_s-imat": ("iq3_s", "q4_0", True),
    "iq3_xxs-imat": ("iq3_xxs", "q4_0", True),
    # 2-bit: our int2 student (and ternary A0) is Q2_0 everywhere
    "q2_0-rtn":   ("q2_0", "q2_0", False),
    "q2_k-imat":  ("q2_k", "q2_0", True),
    "iq2_s-imat": ("iq2_s", "q2_0", True),
    "iq2_xs-imat": ("iq2_xs", "q2_0", True),
    "iq2_xxs-imat": ("iq2_xxs", "q2_0", True),
    # where the error sits: the same gate/up types with the down projection in non-linear IQ4_NL (32-block, fits 384)
    "q4_0-iq4nldown-imat": ("q4_0", "iq4_nl", True),
    "iq4_nl-imat": ("iq4_nl", "iq4_nl", True),
    "iq4_xs-rtn": ("iq4_xs", "iq4_nl", False),
    "q3_k-iq4nldown-imat": ("q3_k", "iq4_nl", True),
    "iq3_s-iq4nldown-imat": ("iq3_s", "iq4_nl", True),
    "q2_0-imat": ("q2_0", "q2_0", True),
    "q2_k-iq4nldown-imat": ("q2_k", "iq4_nl", True),
    "iq2_s-iq4nldown-imat": ("iq2_s", "iq4_nl", True),
    # Q2N's real competitors (~3.1-3.6 GB): the smaller IQ1/IQ2 codebooks with the down projection in IQ4_NL
    "iq2_xs-iq4nldown-imat": ("iq2_xs", "iq4_nl", True),
    "iq2_xxs-iq4nldown-imat": ("iq2_xxs", "iq4_nl", True),
    "iq1_m-iq4nldown-imat": ("iq1_m", "iq4_nl", True),
    "iq1_s-iq4nldown-imat": ("iq1_s", "iq4_nl", True),
    "tq1_0-rtn": ("tq1_0", "q2_0", False),
    "q2_0-iq4nldown-imat": ("q2_0", "iq4_nl", True),
    "tq2_0-rtn": ("tq2_0", "q2_0", False),
    # ~1.7-bit: our TQ1_0 repack of A0 is TQ1_0 gate/up with Q2_0 down
    "iq1_m-imat": ("iq1_m", "q2_0", True),
    "iq1_s-imat": ("iq1_s", "q2_0", True),
}


def run(cmd: list[str], log: Path) -> None:
    print(" ".join(map(str, cmd)), flush=True)
    with open(log, "w") as f:
        r = subprocess.run([str(c) for c in cmd], stdout=f, stderr=subprocess.STDOUT)
    if r.returncode:
        sys.exit(f"failed ({r.returncode}); see {log}")


def texts(a) -> None:
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(TEACHER_HF)
    for stream, n in (("eval", EVAL_SEQS), ("calib", CALIB_SEQS)):
        t = np.memmap(DATA / f"shards/{stream}/tokens.u32", dtype=np.uint32, mode="r")
        seqs = [tok.decode(t[i * CTX:(i + 1) * CTX].tolist(), skip_special_tokens=False) for i in range(n)]
        (OUT / f"{stream}.txt").write_text("".join(seqs))
        print(stream, n, "sequences,", sum(map(len, seqs)), "chars")


def imatrix(a) -> None:
    # on a Q8_0 copy of the teacher, the usual practice: the bf16 one (16 GB) does not fit the 16 GB card, and with
    # its experts on the CPU a pass took 7 s per 512-token chunk (~2 h for the calib text)
    q8 = OUT / "teacher-q8_0.gguf"
    if not q8.exists():
        run([BIN / "llama-quantize", TEACHER, q8, "q8_0", "16"], OUT / "logs/quantize-q8_0.log")
    run([BIN / "llama-imatrix", "-m", q8, "-f", OUT / "calib.txt", "-o", OUT / "imatrix.gguf", "-c", "512",
         "-ngl", "99", "-t", "16", "-b", "512"], OUT / "logs/imatrix.log")


def quantize(a) -> None:
    (OUT / "gguf").mkdir(exist_ok=True)
    for name in a.names or VARIANTS:
        gu, down, imat = VARIANTS[name]
        dst = OUT / "gguf" / f"teacher-{name}.gguf"
        if dst.exists() and not a.force:
            continue
        cmd = [BIN / "llama-quantize", "--allow-requantize"]
        if imat:
            cmd += ["--imatrix", OUT / "imatrix.gguf"]
        # a bf16 target makes llama-quantize copy every tensor unchanged, overrides included; so the target is Q8_0
        # and every 2-D non-expert tensor is pinned back to bf16 (norms and the router are never quantised)
        cmd += ["--tensor-type", f"ffn_gate_exps={gu}", "--tensor-type", f"ffn_up_exps={gu}",
                "--tensor-type", f"ffn_down_exps={down}"]
        for t in ("attn_q", "attn_k", "attn_v", "attn_output"):
            cmd += ["--tensor-type", f"{t}=bf16"]
        cmd += ["--output-tensor-type", "bf16", "--token-embedding-type", "bf16", TEACHER, dst, "q8_0", "16"]
        run(cmd, OUT / f"logs/quantize-{name}.log")


def base(a) -> None:
    # the reference stays exact bf16: the experts of layers 0-8 (~4.9 GB) on the CPU, the rest on the card
    run([BIN / "llama-perplexity", "-m", TEACHER, "-f", OUT / "eval.txt", "-c", str(CTX), "-ngl", "99",
         "-ot", r"blk\.[0-8]\.ffn_.*_exps=CPU", "-t", "16", "--kl-divergence-base", OUT / "teacher.kld"],
        OUT / "logs/base.log")


def kld(a) -> None:
    targets = {n: OUT / "gguf" / f"teacher-{n}.gguf" for n in (a.names or ([] if a.gguf else VARIANTS))}
    extra = dict(g.split("=", 1) for g in a.gguf or [])
    if extra:  # remembered for the report
        reg = OUT / "students.json"
        reg.write_text(json.dumps({**(json.loads(reg.read_text()) if reg.exists() else {}), **extra}, indent=1))
    targets.update(extra)
    for name, path in targets.items():
        log = OUT / f"logs/kld-{name}.log"
        if log.exists() and "Mean    KLD" in log.read_text() and not a.force:
            continue
        run([BIN / "llama-perplexity", "-m", path, "-c", str(CTX), "-ngl", "99", "-t", "16",
             "--kl-divergence-base", OUT / "teacher.kld", "--kl-divergence"], log)


def expert_bpw(path: Path) -> float:
    sys.path.insert(0, "/path/to/llama.cpp/gguf-py")
    import gguf
    r = gguf.GGUFReader(str(path))
    b = n = 0
    for t in r.tensors:
        if any(e in t.name for e in EXPS):
            b += int(t.n_bytes)
            n += int(np.prod(t.shape))
    return 8 * b / n


def report(a) -> None:
    rows = []
    for log in sorted((OUT / "logs").glob("kld-*.log")):
        txt = log.read_text()
        name = log.stem[4:]
        m = {k: re.search(rf"{k}\s*:\s*([-\d.]+)", txt) for k in ("Mean    KLD", "Same top p", "Mean PPL\\(Q\\)")}
        if not m["Mean    KLD"]:
            continue
        src = OUT / "gguf" / f"teacher-{name}.gguf"
        gg = src if src.exists() else Path(json.loads((OUT / "students.json").read_text())[name])
        rows.append({"name": name, "kld": float(m["Mean    KLD"].group(1)),
                     "same_top1_pct": float(m["Same top p"].group(1)) if m["Same top p"] else None,
                     "ppl": float(m["Mean PPL\\(Q\\)"].group(1)) if m["Mean PPL\\(Q\\)"] else None,
                     "expert_bpw": round(expert_bpw(gg), 3), "file_gb": round(gg.stat().st_size / 1e9, 2)})
    rows.sort(key=lambda r: r["expert_bpw"])
    (REPO / "reports/ptq_baseline.json").write_text(json.dumps(rows, indent=1) + "\n")
    lines = ["| model | expert bits/weight | file GB | KL vs teacher | same top-1 % | PPL |", "|---|---|---|---|---|---|"]
    lines += [f"| {r['name']} | {r['expert_bpw']} | {r['file_gb']} | {r['kld']:.4f} | {r['same_top1_pct']} | "
              f"{r['ppl']} |" for r in rows]
    print("\n".join(lines))


def main() -> None:
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="op", required=True)
    sub.add_parser("texts")
    sub.add_parser("imatrix")
    q = sub.add_parser("quantize")
    q.add_argument("names", nargs="*")
    q.add_argument("--force", action="store_true")
    sub.add_parser("base")
    k = sub.add_parser("kld")
    k.add_argument("names", nargs="*")
    k.add_argument("--gguf", action="append", help="name=path of an extra GGUF (a student)")
    k.add_argument("--force", action="store_true")
    sub.add_parser("report")
    a = p.parse_args()
    (OUT / "logs").mkdir(parents=True, exist_ok=True)
    {"texts": texts, "imatrix": imatrix, "quantize": quantize, "base": base, "kld": kld, "report": report}[a.op](a)


if __name__ == "__main__":
    main()