#!/usr/bin/env python3
"""One-time codegen: extract the iq3 grids kgrid_256/kgrid_512 (iq3xs_init_impl)
from llama.cpp's ggml-quants.c into kernels/iq3_tables.py.

Run: python tools/dev_gen_iq3_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()
lines = ["# generated by tools/dev_gen_iq3_tables.py from ggml-quants.c -- do not edit", ""]
for size in (256, 512):
    m = re.search(rf"kgrid_{size}\[{size}\] = \{{(.*?)\}};", text, re.S)
    if not m:
        sys.exit(f"kgrid_{size} not found")
    vals = [int(v) for v in re.findall(r"\d+", m.group(1))]
    assert len(vals) == size, len(vals)
    lines.append(f"KGRID_3BIT_{size} = [")
    for i in range(0, size, 16):
        lines.append("    " + ", ".join(str(v) for v in vals[i:i + 16]) + ",")
    lines += ["]", ""]

out = ROOT / "kernels" / "iq3_tables.py"
out.write_text("\n".join(lines))
print(f"wrote {out}")