"""Precompute Lloyd-Max codebooks for common LLM head dimensions.

Generates turboquant/codebook_data.py with exact Beta(d/2, d/2) codebooks
for dim < 256, avoiding scipy overhead at import time.

Usage:
    python scripts/precompute_codebooks.py
"""

from turboquant.codebook import _lloyd_max_beta

DIMS = [32, 48, 64, 80, 96, 128, 160, 192]
BITS = [1, 2, 3, 4]


def main():
    codebooks = {}
    for dim in DIMS:
        for bit_width in BITS:
            print(f"Computing Beta({dim / 2}, {dim / 2}) codebook for {bit_width}-bit...")
            cb = _lloyd_max_beta(bit_width, dim)
            codebooks[(bit_width, dim)] = {
                "boundaries": cb.boundaries.tolist(),
                "centroids": cb.centroids.tolist(),
            }

    out_path = "turboquant/codebook_data.py"
    with open(out_path, "w") as f:  # noqa: PTH123
        f.write('"""Precomputed Lloyd-Max codebooks for Beta distributions.\n\n')
        f.write("Generated by scripts/precompute_codebooks.py — do not edit manually.\n")
        f.write('"""\n\n')
        f.write("PRECOMPUTED_BETA_CODEBOOKS: dict[tuple[int, int], dict[str, list[float]]] = {\n")
        for (bit_width, dim), data in sorted(codebooks.items()):
            f.write(f"    ({bit_width}, {dim}): {{\n")
            f.write(f'        "boundaries": {data["boundaries"]},\n')
            f.write(f'        "centroids": {data["centroids"]},\n')
            f.write("    },\n")
        f.write("}\n")

    print(f"\nWrote {len(codebooks)} codebooks to {out_path}")


if __name__ == "__main__":
    main()
