#!/usr/bin/env python3
"""One-time codegen: extract kgrid_2bit_256 from llama.cpp's ggml-quants.c
into kernels/iq2_tables.py, so quant-studio doesn't need the C sources at runtime.

Run: python tools/gen_iq2_tables.py [path/to/ggml-quants.c]
"""

import re
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
src = Path(sys.argv[1]) if len(sys.argv) > 1 else ROOT / ".." / "llama.cpp" / "ggml" / "src" / "ggml-quants.c"

text = src.read_text()
m = re.search(r"kgrid_2bit_256\[256\] = \{(.*?)\};", text, re.S)
if not m:
    sys.exit("kgrid_2bit_256 not found")
vals = [int(v) for v in re.findall(r"\d+", m.group(1))]
assert len(vals) == 256, len(vals)

out = ROOT / "kernels" / "iq2_tables.py"
lines = ["# generated by tools/gen_iq2_tables.py from ggml-quants.c -- do not edit",
         "", "KGRID_2BIT_256 = ["]
for i in range(0, 256, 16):
    lines.append("    " + ", ".join(str(v) for v in vals[i:i + 16]) + ",")
lines += ["]", ""]
out.write_text("\n".join(lines))
print(f"wrote {out} ({len(vals)} entries)")