#!/usr/bin/env python3
# =============================================================================
# core.py — Prime Mystery Engine: core primitives
# Author: Simon-Pierre Boucher — contact@spboucher.ai
# =============================================================================
# Building blocks (Phase 0 of the protocol):
#   - primes_upto(n)          : numpy Sieve of Eratosthenes (odd-only)
#   - segmented_primes(lo,hi) : segmented sieve, constant memory per segment
#   - is_prime(n)             : deterministic Miller-Rabin for n < 3.317e24
#   - bpsw(n)                 : Baillie-PSW (MR base 2 + strong Lucas)
#   - prime_count(n)          : pi(n) via sieve (validation helper)
# All routines are pure-Python/numpy, no hidden prime tables beyond the
# 12 deterministic MR bases (which are witnesses, not a prime list).
# =============================================================================

from __future__ import annotations  # keep annotations lazy (Python 3.9 nodes)

import numpy as np

# Deterministic Miller-Rabin witness set: correct for all n < 3.317e24
# (Sorenson & Webster 2015). Documented per the Golden Rule.
_MR_BASES = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37)
_MR_LIMIT = 3_317_044_064_679_887_385_961_981


def primes_upto(n: int) -> np.ndarray:
    """All primes <= n, via odd-only Sieve of Eratosthenes (numpy)."""
    if n < 2:
        return np.array([], dtype=np.int64)
    if n < 3:
        return np.array([2], dtype=np.int64)
    # index i represents the odd number 2i+1; index 0 (=1) is not prime
    size = (n + 1) // 2
    sieve = np.ones(size, dtype=bool)
    sieve[0] = False
    for i in range(1, (int(n ** 0.5) + 1) // 2 + 1):
        if sieve[i]:
            p = 2 * i + 1
            start = (p * p) // 2
            sieve[start::p] = False
    odds = 2 * np.nonzero(sieve)[0].astype(np.int64) + 1
    return np.concatenate(([np.int64(2)], odds))


def sieve_segment(lo: int, hi: int, base_primes: np.ndarray | None = None) -> np.ndarray:
    """Primes in [lo, hi) via segmented sieve. base_primes must cover sqrt(hi)."""
    if hi <= 2:
        return np.array([], dtype=np.int64)
    lo = max(lo, 2)
    if base_primes is None:
        base_primes = primes_upto(int(hi ** 0.5) + 1)
    seg = np.ones(hi - lo, dtype=bool)
    for p in base_primes:
        p = int(p)
        if p * p >= hi:
            break
        start = max(p * p, ((lo + p - 1) // p) * p)
        seg[start - lo::p] = False
    if lo <= 1:
        seg[: 2 - lo] = False
    return lo + np.nonzero(seg)[0].astype(np.int64)


def segmented_primes(lo: int, hi: int, segment_size: int = 10_000_000):
    """Yield numpy arrays of primes covering [lo, hi) segment by segment."""
    base = primes_upto(int(hi ** 0.5) + 1)
    for start in range(lo, hi, segment_size):
        end = min(start + segment_size, hi)
        yield sieve_segment(start, end, base)


def prime_count(n: int) -> int:
    """pi(n), computed by segmented sieve (validation helper)."""
    total = 0
    for chunk in segmented_primes(2, n + 1):
        total += len(chunk)
    return total


def _mr_witness(n: int, a: int, d: int, r: int) -> bool:
    """True if a is a Miller-Rabin witness for compositeness of n."""
    x = pow(a, d, n)
    if x == 1 or x == n - 1:
        return False
    for _ in range(r - 1):
        x = x * x % n
        if x == n - 1:
            return False
    return True


def is_prime(n: int) -> bool:
    """Deterministic primality for n < 3.317e24 (Miller-Rabin, 12 bases).
    For larger n, falls back to BPSW (documented: no known counterexample,
    verified exhaustively below 2^64)."""
    if n < 2:
        return False
    for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
        if n % p == 0:
            return n == p
    d, r = n - 1, 0
    while d % 2 == 0:
        d //= 2
        r += 1
    if n < _MR_LIMIT:
        return not any(_mr_witness(n, a, d, r) for a in _MR_BASES if a % n)
    return bpsw(n)


def _jacobi(a: int, n: int) -> int:
    """Jacobi symbol (a/n), n odd positive."""
    a %= n
    result = 1
    while a:
        while a % 2 == 0:
            a //= 2
            if n % 8 in (3, 5):
                result = -result
        a, n = n, a
        if a % 4 == 3 and n % 4 == 3:
            result = -result
        a %= n
    return result if n == 1 else 0


def _strong_lucas(n: int) -> bool:
    """Strong Lucas probable-prime test (Selfridge parameters). n odd, >2,
    not a perfect square, no small factors."""
    # Selfridge: first D in 5,-7,9,-11,... with Jacobi(D/n) = -1
    D = 5
    while True:
        j = _jacobi(D % n, n)
        if j == -1:
            break
        if j == 0 and abs(D) != n:
            return False
        D = -D - 2 if D > 0 else -D + 2
        if abs(D) > 1_000_000:  # would indicate a square slipped through
            raise ArithmeticError("no suitable D found; n likely a square")
    Q = (1 - D) // 4
    # factor n+1 = d * 2^s
    d, s = n + 1, 0
    while d % 2 == 0:
        d //= 2
        s += 1
    # Lucas sequences U_d, V_d by binary ladder
    U, V, k = 0, 2, 0
    Qk = 1
    inv2 = (n + 1) // 2  # inverse of 2 mod n (n odd)
    for bit in bin(d)[2:]:
        # double: U_{2k}=U_k V_k ; V_{2k}=V_k^2 - 2 Q^k
        U, V = U * V % n, (V * V - 2 * Qk) % n
        Qk = Qk * Qk % n
        if bit == "1":
            # increment: U_{k+1}=(P U + V)/2 ; V_{k+1}=(D U + P V)/2 ; P=1
            U, V = (U + V) * inv2 % n, (D * U + V) * inv2 % n
            Qk = Qk * Q % n
    if U == 0 or V == 0:
        return True
    for _ in range(s - 1):
        V = (V * V - 2 * Qk) % n
        if V == 0:
            return True
        Qk = Qk * Qk % n
    return False


def _is_square(n: int) -> bool:
    r = math_isqrt(n)
    return r * r == n


from math import isqrt as math_isqrt  # noqa: E402


def bpsw(n: int) -> bool:
    """Baillie-PSW: strong MR base 2 + strong Lucas (Selfridge).
    No known counterexample; deterministic below 2^64."""
    if n < 2:
        return False
    for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
        if n % p == 0:
            return n == p
    d, r = n - 1, 0
    while d % 2 == 0:
        d //= 2
        r += 1
    if _mr_witness(n, 2, d, r):
        return False
    if _is_square(n):
        return False
    return _strong_lucas(n)


def gaps_in_range(lo: int, hi: int, segment_size: int = 50_000_000):
    """Iterate consecutive prime gaps over [lo, hi). Yields (p, gap) where
    gap = next_prime(p) - p, for all consecutive pairs with p in [lo, hi)."""
    prev = None
    for chunk in segmented_primes(max(lo, 2), hi, segment_size):
        if len(chunk) == 0:
            continue
        if prev is not None:
            yield int(prev), int(chunk[0] - prev)
        diffs = np.diff(chunk)
        for p, g in zip(chunk[:-1], diffs):
            yield int(p), int(g)
        prev = chunk[-1]