#!/usr/bin/env python3
# =============================================================================
# explore_gaps.py — Cycle 1 EXPLORER: prime-gap statistics up to a bound
# Author: Simon-Pierre Boucher — contact@spboucher.ai
# =============================================================================
# Fully vectorized (numpy) segmented scan of [2, LIMIT). Produces:
#   data/maximal_gaps_<L>.csv   : record (maximal) gaps with merit and CSG ratio
#   data/gap_histogram_<L>.csv  : count of each even gap g
#   data/gap_mod6_start_<L>.csv : histogram of (start prime mod 6, gap mod 6)
#   data/summary_<L>.json       : headline numbers (deterministic, no seed used)
# Method: segmented Sieve of Eratosthenes (core.py), np.diff per segment,
# np.bincount accumulation. No per-prime Python loop.
# Usage: python3 explore_gaps.py 1e8
# =============================================================================

import json
import sys
import time
from math import log
from pathlib import Path

import numpy as np

from core import primes_upto, sieve_segment

DATA = Path(__file__).resolve().parent.parent / "data"
DATA.mkdir(exist_ok=True)


def scan(limit: int, segment_size: int = 50_000_000):
    base = primes_upto(int(limit ** 0.5) + 1)
    hist = np.zeros(4000, dtype=np.int64)          # gap -> count
    mod6 = np.zeros((6, 6), dtype=np.int64)        # (p mod 6, gap mod 6) -> count
    maximal = []                                   # (gap, start prime) records
    best = 0
    prev = None                                    # last prime of previous segment
    n_gaps = 0
    t0 = time.time()
    for lo in range(2, limit, segment_size):
        hi = min(lo + segment_size, limit)
        primes = sieve_segment(lo, hi, base)
        if len(primes) == 0:
            continue
        if prev is not None:
            primes = np.concatenate(([prev], primes))
        if len(primes) >= 2:
            gaps = np.diff(primes)
            starts = primes[:-1]
            n_gaps += len(gaps)
            hist += np.bincount(gaps, minlength=len(hist))[: len(hist)]
            np.add.at(mod6, (starts % 6, gaps % 6), 1)
            # record gaps within this segment
            seg_best = int(gaps.max())
            if seg_best > best:
                idx = np.nonzero(gaps > best)[0]
                for i in idx:
                    g, p = int(gaps[i]), int(starts[i])
                    if g > best:
                        best = g
                        maximal.append((g, p))
        prev = primes[-1]
    dt = time.time() - t0
    return hist, mod6, maximal, n_gaps, dt


def main():
    limit = int(float(sys.argv[1])) if len(sys.argv) > 1 else 10**8
    tag = f"{limit:.0e}".replace("+0", "").replace("+", "")
    hist, mod6, maximal, n_gaps, dt = scan(limit)

    # maximal gaps with merit and Cramér-Shanks-Granville ratio
    with open(DATA / f"maximal_gaps_{tag}.csv", "w") as f:
        f.write("gap,start_prime,merit,csg_ratio\n")
        for g, p in maximal:
            f.write(f"{g},{p},{g/log(p):.6f},{g/log(p)**2:.6f}\n")

    nz = np.nonzero(hist)[0]
    with open(DATA / f"gap_histogram_{tag}.csv", "w") as f:
        f.write("gap,count\n")
        for g in nz:
            f.write(f"{g},{hist[g]}\n")

    with open(DATA / f"gap_mod6_start_{tag}.csv", "w") as f:
        f.write("p_mod6,gap_mod6,count\n")
        for a in range(6):
            for b in range(6):
                if mod6[a, b]:
                    f.write(f"{a},{b},{mod6[a, b]}\n")

    # headline stats
    champion = int(nz[np.argmax(hist[nz])])
    g_arr = nz.astype(float)
    c_arr = hist[nz].astype(float)
    total = c_arr.sum()
    frac_mod6 = {r: float(c_arr[g_arr % 6 == r].sum() / total) for r in range(6)}
    gmax, pmax = maximal[-1]
    best_merit = max((g / log(p), g, p) for g, p in maximal if p > 100)
    best_csg = max((g / log(p) ** 2, g, p) for g, p in maximal if p > 100)
    summary = {
        "limit": limit,
        "n_gaps": int(n_gaps),
        "scan_seconds": round(dt, 1),
        "largest_gap": {"gap": gmax, "after_prime": pmax},
        "n_maximal_gaps": len(maximal),
        "jumping_champion": champion,
        "gap_fraction_mod6": frac_mod6,
        "best_merit": {"merit": round(best_merit[0], 5), "gap": best_merit[1], "after_prime": best_merit[2]},
        "best_csg": {"csg": round(best_csg[0], 5), "gap": best_csg[1], "after_prime": best_csg[2]},
        "method": "segmented sieve of Eratosthenes, numpy, deterministic (no seed)",
    }
    with open(DATA / f"summary_{tag}.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(json.dumps(summary, indent=2))
    print(f"\nmaximal gaps ({len(maximal)}):")
    for g, p in maximal:
        print(f"  gap {g:4d} after {p:>15d}  merit {g/log(p):7.4f}  csg {g/log(p)**2:.4f}")


if __name__ == "__main__":
    main()