#!/usr/bin/env python3
# =============================================================================
# model_c4.py — Cycle 2 PROVER: HL triple-correlation model for the C4 constant
# Author: Simon-Pierre Boucher — contact@spboucher.ai
# =============================================================================
# Heuristic model for the joint law of consecutive gaps (g1, g2) at scale
# lambda = ln x:
#     P(g1, g2)  PROPORTIONAL TO  W(g1, g2) * exp(-(g1+g2)/lambda)
# where W is the Hardy-Littlewood singular series of the triple {0, g1, g1+g2}
# restricted to primes p <= diameter (the tail p > g1+g2 is constant across
# pairs and cancels in the normalization):
#     W = prod_{p odd, p <= G} [ (1 - nu_p/p) / (1 - 1/p)^3 ] * [p=2 term const]
#     nu_p = #{0, g1, g1+g2 mod p}  (in {1,2,3}; equals 3 unless p | g1,
#            p | g2 or p | g1+g2)
# KNOWN LIMITATION (documented): no inclusion-exclusion for "no prime strictly
# inside the gaps" — this is the crude first-order model. If W == 1 (Cramer),
# rho == 0 identically; all predicted anticorrelation is singular-series
# coupling (the p | g1+g2 factor prevents factorization W = u(g1) v(g2)).
# Output: model rho(lambda) * lambda for lambdas of interest + large-lambda
# trend, printed and saved to data/model_c4.json. Deterministic.
# =============================================================================

import json
from math import exp, log
from pathlib import Path

import numpy as np

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


def small_primes(n):
    s = np.ones(n + 1, dtype=bool); s[:2] = False
    for i in range(2, int(n ** 0.5) + 1):
        if s[i]:
            s[i * i:: i] = False
    return np.nonzero(s)[0]


def model_rho(lam: float, gmax: int) -> float:
    """Pearson correlation of (g1,g2) under the HL-coupled exponential model."""
    gs = np.arange(2, gmax + 1, 2, dtype=np.int64)
    g1 = gs[:, None]
    g2 = gs[None, :]
    gsum = g1 + g2
    logW = np.zeros((len(gs), len(gs)), dtype=np.float64)
    for p in small_primes(2 * gmax):
        if p == 2:
            continue  # constant factor for even gaps (nu=1), cancels
        # nu_p on the grid
        d1 = (g1 % p == 0)
        d2 = (g2 % p == 0)
        ds = (gsum % p == 0)
        nu = np.full(logW.shape, 3, dtype=np.int64)
        nu[d1 | d2 | ds] = 2
        nu[d1 & d2] = 1
        logW += np.log(1 - nu / p) - 3 * log(1 - 1 / p)
    P = np.exp(logW - gsum / lam)
    P /= P.sum()
    m1 = (P * g1).sum(); m2 = (P * g2).sum()
    v1 = (P * (g1 - m1) ** 2).sum(); v2 = (P * (g2 - m2) ** 2).sum()
    cov = (P * (g1 - m1) * (g2 - m2)).sum()
    return float(cov / (v1 * v2) ** 0.5)


def main():
    out = {"model": "P(g1,g2) ~ W_HL(0,g1,g1+g2) * exp(-(g1+g2)/lambda)",
           "predictions": {}}
    print(f"{'lambda':>8} {'x=e^lam':>10} {'rho_model':>12} {'rho*lambda':>12}")
    for lam, label in [(log(1e8), "1e8"), (log(1e9), "1e9"), (log(1e10), "1e10"),
                       (log(4e10), "4e10"), (log(1e12), "1e12"),
                       (log(1e15), "1e15"), (log(1e20), "1e20"),
                       (log(1e30), "1e30"), (log(1e40), "1e40")]:
        gmax = int(14 * lam)          # truncation: weight e^-14 ~ 8e-7
        r = model_rho(lam, gmax)
        out["predictions"][label] = {"lambda": round(lam, 4),
                                     "rho": round(r, 6),
                                     "rho_times_lambda": round(r * lam, 5)}
        print(f"{lam:8.3f} {label:>10} {r:12.6f} {r*lam:12.5f}")
    with open(DATA / "model_c4.json", "w") as f:
        json.dump(out, f, indent=2)
    print("written:", DATA / "model_c4.json")


if __name__ == "__main__":
    main()