#!/usr/bin/env python3
"""D-58: a codebook student's start, read from a llama.cpp GGUF of the teacher (`scripts/ptq_baseline.py`).

    python scripts/cb_import.py --gguf $MT_DATA/ptq/gguf/teacher-iq3_s-iq4nldown-imat.gguf --out $MT_DATA/cbinit/iq3s

Writes <out>/experts.safetensors with, per layer and projection (gate, up, down): `blocks` (the GGUF's raw blocks,
uint8 [E, rows, nblk, nbytes]), `v` (int8 [E, rows, cols]: codebook entry × sign per weight), `d` (fp16
[E, rows, nblk]) and `s` (uint8 [E, rows, nsub]); and <out>/manifest.json with each tensor's format and the
source file's sha256. Everything outside the routed experts comes from the teacher, as in every PTQ baseline
(attention, embeddings bf16; norms, router f32), so this file and the GGUF describe the same model.
"""
from __future__ import annotations

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

import numpy as np
import torch

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.export.gguf import gguf_module  # noqa: E402
from lowbit.quant import codebook as cb  # noqa: E402

LLAMA = "/path/to/llama.cpp"


def sha256(p: Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 24), b""):
            h.update(chunk)
    return h.hexdigest()


def write_grids(out: Path, names: set, gguf) -> None:
    """grids.safetensors: each format's codebook (`quant.codebook.grid_from_gguf`), for free-code training."""
    from safetensors.torch import save_file as save_pt
    save_pt({n: cb.grid_from_gguf(cb.FORMATS[n], gguf).contiguous() for n in sorted(names)}, str(out / "grids.safetensors"))


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--gguf", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--llama", default=LLAMA, help="a directory holding llama.cpp's gguf-py (spark: third_party)")
    p.add_argument("--grids-only", action="store_true", help="only (re)write grids.safetensors for the manifest's formats")
    a = p.parse_args()
    from safetensors.numpy import save_file
    gguf = gguf_module(a.llama)
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    if a.grids_only:
        write_grids(out, set(json.loads((out / "manifest.json").read_text())["formats"].values()), gguf)
        return
    tensors, formats = {}, {}
    for t in gguf.GGUFReader(a.gguf).tensors:
        m = re.fullmatch(r"blk\.(\d+)\.ffn_(gate|up|down)_exps\.weight", t.name)
        if not m:
            continue
        name = gguf.GGMLQuantizationType(t.tensor_type).name
        if name not in cb.FORMATS:
            sys.exit(f"{t.name} is {name}, not a codebook format ({sorted(cb.FORMATS)})")
        fmt = cb.FORMATS[name]
        E, rows, cols = (int(x) for x in reversed(t.shape))  # gguf-py lists ne0 first
        nblk = cols // fmt.block
        b = np.asarray(t.data).reshape(-1, fmt.nbytes)
        key = f"layers.{m.group(1)}.{m.group(2)}"
        formats[key] = name
        tensors[f"{key}.blocks"] = b.reshape(E, rows, nblk, fmt.nbytes).copy()
        tensors[f"{key}.v"] = cb.values(fmt, b, gguf).reshape(E, rows, cols)
        tensors[f"{key}.d"] = cb.get_d(b).reshape(E, rows, nblk)
        tensors[f"{key}.s"] = cb.get_s(fmt, b).reshape(E, rows, -1)
        print(key, name, (E, rows, cols), flush=True)
    save_file(tensors, str(out / "experts.safetensors"))
    write_grids(out, set(formats.values()), gguf)
    (out / "manifest.json").write_text(json.dumps({"source": str(a.gguf), "source_sha256": sha256(Path(a.gguf)),
                                                   "formats": formats}, indent=1))
    print(json.dumps({"out": str(out), "tensors": len(formats),
                      "bytes": (out / "experts.safetensors").stat().st_size}))


if __name__ == "__main__":
    main()