#!/usr/bin/env python3
# =============================================================================
# analyze_gaps.py — Cycle 1 EXPLORER/ADVERSARY: gap statistics vs Hardy-Littlewood
# Author: Simon-Pierre Boucher — contact@spboucher.ai
# =============================================================================
# One segmented scan up to LIMIT with checkpoints. Produces (in /data):
#   gapstats_<L>.npz : per-checkpoint gap histograms, consecutive-gap pair
#                      matrix (g_n, g_{n+1}) for gaps < 400, correlation sums
#   gapstats_<L>.json: HL-normalized ratios, champion tables, correlations
# HL model used (documented):
#   N(g, x) ~ C(g) * INT_2^x exp(-g/ln t) / ln^2 t dt,
#   C(g) = 2*C2 * prod_{p | g, p odd} (p-1)/(p-2),  C2 = 0.6601618158468696.
# Deterministic; no randomness.
# Usage: python3 analyze_gaps.py 1e8 [checkpoints default 1e6,1e7,1e8,1e9,4e9 <= limit]
# =============================================================================

import json
import sys
import time
from math import exp, 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)

C2 = 0.6601618158468696
GMAX = 400  # pair matrix covers gaps < GMAX


def singular(g: int) -> float:
    """C(g) = 2*C2 * prod_{p|g, p>2} (p-1)/(p-2)."""
    c = 2 * C2
    d, m = 3, g
    while d * d <= m:
        if m % d == 0:
            if d > 2:
                c *= (d - 1) / (d - 2)
            while m % d == 0:
                m //= d
        d += 2 if d > 2 else 1
    if m > 2:
        c *= (m - 1) / (m - 2)
    return c


def hl_integral(g: int, x: int) -> float:
    """INT_2^x exp(-g/ln t)/ln^2 t dt, trapezoid on a log grid."""
    ts = np.exp(np.linspace(log(3), log(x), 4000))
    ys = np.exp(-g / np.log(ts)) / np.log(ts) ** 2
    return float(np.trapezoid(ys, ts))


def scan(limit: int, checkpoints, segment_size: int = 50_000_000):
    base = primes_upto(int(limit ** 0.5) + 1)
    hist = np.zeros(4000, dtype=np.int64)
    pair = np.zeros((GMAX, GMAX), dtype=np.int64)
    # streaming sums for Pearson correlation of (g_n, g_{n+1})
    S = dict(n=0, x=0.0, y=0.0, xx=0.0, yy=0.0, xy=0.0)
    cp_hists, cp_corr = {}, {}
    cps = sorted(checkpoints)
    # segment boundaries aligned on checkpoints so each checkpoint histogram
    # covers exactly [2, checkpoint)
    bounds = sorted({limit, *cps, *range(2, limit, segment_size)} - {2})
    prev = None       # last prime seen
    prev_gap = None   # gap ending at prev
    t0 = time.time()
    lo = 2
    for hi in bounds:
        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)
            hist += np.bincount(gaps, minlength=len(hist))[: len(hist)]
            # consecutive pairs, including the junction pair across segments
            if prev_gap is not None:
                gl = np.concatenate(([prev_gap], gaps))
            else:
                gl = gaps
            a, b = gl[:-1].astype(np.float64), gl[1:].astype(np.float64)
            am, bm = gl[:-1], gl[1:]
            mask = (am < GMAX) & (bm < GMAX)
            np.add.at(pair, (am[mask], bm[mask]), 1)
            S["n"] += len(a)
            S["x"] += a.sum(); S["y"] += b.sum()
            S["xx"] += (a * a).sum(); S["yy"] += (b * b).sum()
            S["xy"] += (a * b).sum()
            prev_gap = int(gaps[-1])
        prev = primes[-1]
        while cps and hi >= cps[0]:
            c = cps.pop(0)
            cp_hists[c] = hist.copy()
            cp_corr[c] = dict(S)
            print(f"  checkpoint {c:.1e} reached at {time.time()-t0:.1f}s")
        lo = hi
    return hist, pair, cp_hists, cp_corr, time.time() - t0


def pearson(S):
    n = S["n"]
    cov = S["xy"] / n - (S["x"] / n) * (S["y"] / n)
    vx = S["xx"] / n - (S["x"] / n) ** 2
    vy = S["yy"] / n - (S["y"] / n) ** 2
    return cov / (vx * vy) ** 0.5


def main():
    limit = int(float(sys.argv[1])) if len(sys.argv) > 1 else 10**8
    cps = [int(c) for c in (10**6, 10**7, 10**8, 10**9, 4 * 10**9) if c <= limit]
    tag = f"{limit:.0e}".replace("+0", "").replace("+", "")
    hist, pair, cp_hists, cp_corr, dt = scan(limit, cps)

    np.savez_compressed(
        DATA / f"gapstats_{tag}.npz",
        hist=hist, pair=pair,
        checkpoints=np.array(sorted(cp_hists)),
        **{f"hist_{c}": h for c, h in cp_hists.items()},
    )

    out = {"limit": limit, "scan_seconds": round(dt, 1), "checkpoints": {}}
    for c in sorted(cp_hists):
        h = cp_hists[c]
        nz = np.nonzero(h)[0]
        even = nz[nz % 2 == 0]
        # HL comparison for even gaps up to 120
        ratios = {}
        for g in [int(v) for v in even if 2 <= v <= 120]:
            pred = singular(g) * hl_integral(g, c)
            ratios[g] = round(int(h[g]) / pred, 4) if pred > 0 else None
        champ = int(nz[np.argmax(h[nz])])
        # largest G such that every multiple of 6 up to G is a strict local max
        # of the gap histogram (h[g] > h[g-2] and h[g] > h[g+2])
        mult6_localmax_upto = 0
        for g in range(6, int(even.max()) - 2, 6):
            if h[g] > h[g - 2] and h[g] > h[g + 2]:
                mult6_localmax_upto = g
            else:
                break
        out["checkpoints"][str(c)] = {
            "jumping_champion": champ,
            "N2": int(h[2]), "N4": int(h[4]), "N6": int(h[6]),
            "N2_gt_N4": bool(h[2] > h[4]),
            "mult6_local_max_up_to": int(mult6_localmax_upto),
            "pearson_consecutive_gaps": round(pearson(cp_corr[c]), 5),
            "HL_ratio_obs_over_pred": ratios,
        }
    with open(DATA / f"gapstats_{tag}.json", "w") as f:
        json.dump(out, f, indent=2)
    # console digest
    for c in sorted(cp_hists):
        o = out["checkpoints"][str(c)]
        r = o["HL_ratio_obs_over_pred"]
        rv = [v for v in r.values() if v]
        print(f"x={c:.0e}: champ={o['jumping_champion']} N2={o['N2']} N4={o['N4']} "
              f"N2>N4={o['N2_gt_N4']} 6|g-localmax-upto={o['mult6_local_max_up_to']} "
              f"rho={o['pearson_consecutive_gaps']} HLratio[min,max]=[{min(rv):.3f},{max(rv):.3f}]")


if __name__ == "__main__":
    main()