#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Susi Lehtola
#
# Generate the prime-table sources consumed by src/primes.c and primes.h.
#
# Usage (from the repository root):
#   python3 tools/gen_prime_table.py
#
# Two files are written:
#   src/prime_table.inc          — data definitions (g_nprimes, g_primes[],
#                                  g_pi_table[]); included once by
#                                  src/primes.c.
#   src/prime_table_macros.h     — preprocessor constants
#                                  (PRIME_SIEVE_LIMIT, MAX_PRIME_COUNT,
#                                  MAX_FACTORIAL_ARG); included by
#                                  src/primes.h so all consumers see the
#                                  same numeric ceilings.
#
# Adjust LIMIT below to enlarge the prime table; everything else is
# derived.  The sieve is extended slightly past LIMIT so the script can
# locate the smallest prime > LIMIT and set MAX_FACTORIAL_ARG to
# (next_prime - 1), the largest N for which pfrac_mul_factorial(N) is
# correct (every prime p <= N must be in the table).

import os

LIMIT = 20011

# Sieve a little past LIMIT so we can find the next prime above it.
SIEVE_TOP = LIMIT + 1024
sieve = bytearray(SIEVE_TOP + 1)
sieve[0] = sieve[1] = 1
i = 2
while i * i <= SIEVE_TOP:
    if not sieve[i]:
        j = i * i
        while j <= SIEVE_TOP:
            sieve[j] = 1
            j += i
    i += 1

primes = [p for p in range(2, LIMIT + 1) if not sieve[p]]
next_prime_above_limit = next(
    p for p in range(LIMIT + 1, SIEVE_TOP + 1) if not sieve[p])
MAX_FACTORIAL_ARG = next_prime_above_limit - 1

# pi(n) = number of primes <= n.  Indexed by n in [0 .. LIMIT].  Used by
# pfrac_mul_factorial / pfrac_div_factorial to size their per-factorial
# vector add without a per-call binary search of g_primes; also serves
# as the lookup table for prime_index_of() in tests, replacing the
# separate g_prime_index table that used to be emitted here.
pi_table = [0] * (LIMIT + 1)
running = 0
for n in range(LIMIT + 1):
    if not sieve[n] and n >= 2:
        running += 1
    pi_table[n] = running

# Resolve output paths relative to the repository root (which is the
# parent of the directory holding this script).
script_dir = os.path.dirname(os.path.abspath(__file__))
repo_root  = os.path.dirname(script_dir)
inc_path    = os.path.join(repo_root, "src", "prime_table.inc")
macros_path = os.path.join(repo_root, "src", "prime_table_macros.h")

# ── prime_table_macros.h ─────────────────────────────────────────────────────
with open(macros_path, "w") as f:
    f.write("/* SPDX-License-Identifier: BSD-3-Clause\n")
    f.write(" * Copyright (c) 2026 Susi Lehtola\n")
    f.write(" *\n")
    f.write(" * Generated by tools/gen_prime_table.py — do not edit by hand.\n")
    f.write(" *\n")
    f.write(" * Numeric ceilings for the prime table compiled into libwignernj.\n")
    f.write(" * Included by src/primes.h so every translation unit sees the\n")
    f.write(" * same constants.  Regenerate by running tools/gen_prime_table.py\n")
    f.write(" * after adjusting LIMIT in that script. */\n")
    f.write("#ifndef WIGNERNJ_PRIME_TABLE_MACROS_H\n")
    f.write("#define WIGNERNJ_PRIME_TABLE_MACROS_H\n")
    f.write("\n")
    f.write(f"#define PRIME_SIEVE_LIMIT  {LIMIT}\n")
    f.write(f"#define MAX_PRIME_COUNT    {len(primes)}\n")
    f.write(f"#define MAX_FACTORIAL_ARG  {MAX_FACTORIAL_ARG}\n")
    f.write("\n")
    f.write("#endif /* WIGNERNJ_PRIME_TABLE_MACROS_H */\n")

# ── prime_table.inc ──────────────────────────────────────────────────────────
with open(inc_path, "w") as f:
    f.write("/* Generated by tools/gen_prime_table.py — do not edit by hand.\n")
    f.write(f" * PRIME_SIEVE_LIMIT = {LIMIT}, prime count = {len(primes)},\n")
    f.write(f" * MAX_FACTORIAL_ARG = {MAX_FACTORIAL_ARG}"
            f" (next prime above LIMIT is {next_prime_above_limit}). */\n")
    f.write("\n")
    f.write(f"const int g_nprimes = {len(primes)};\n")
    f.write("\n")

    # g_primes — 16 values per row
    f.write("const int g_primes[MAX_PRIME_COUNT] = {\n")
    row = []
    for i, p in enumerate(primes):
        row.append(f"{p:5d}")
        if len(row) == 16 or i == len(primes) - 1:
            suffix = "," if i < len(primes) - 1 else ""
            f.write("    " + ",".join(row) + suffix + "\n")
            row = []
    f.write("};\n")
    f.write("\n")

    # g_pi_table — 20 values per row
    f.write("const short g_pi_table[PRIME_SIEVE_LIMIT + 1] = {\n")
    row = []
    for i, v in enumerate(pi_table):
        row.append(f"{v:5d}")
        if len(row) == 20 or i == len(pi_table) - 1:
            suffix = "," if i < len(pi_table) - 1 else ""
            f.write("    " + ",".join(row) + suffix + "\n")
            row = []
    f.write("};\n")

print(f"wrote {macros_path} (LIMIT = {LIMIT}, MAX_FACTORIAL_ARG = {MAX_FACTORIAL_ARG})")
print(f"wrote {inc_path} ({len(primes)} primes)")
