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

Run: python tools/dev_gen_iq2xs_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_iq2xs_tables.py from ggml-quants.c -- do not edit", ""]
for name, count in (("kgrid_2bit_512", 512), ("kgrid_2bit_1024", 1024)):
    m = re.search(rf"{name}\[{count}\] = \{{(.*?)\}};", text, re.S)
    if not m:
        sys.exit(f"{name} not found")
    vals = [int(v) for v in re.findall(r"\d+", m.group(1))]
    assert len(vals) == count, len(vals)
    assert max(vals) < 43692, max(vals)
    lines.append(f"{name.upper()} = [")
    for i in range(0, count, 16):
        lines.append("    " + ", ".join(str(v) for v in vals[i:i + 16]) + ",")
    lines += ["]", ""]

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