//! Primality — Miller–Rabin rounds, the deterministic u64 test, and prime generation.

use crate::number::modular::mulmod;
use crate::number::{isqrt, jacobi, modpow};

/// One Miller–Rabin round: `true` if `n` is a strong probable prime to base `a`
/// — that is, `n` is prime, or `a` is a *strong liar* for the composite `n` —
/// and `false` only when `a` is a *strong witness*, which proves `n` composite.
///
/// Writing `n - 1 = 2^s * d` with `d` odd, `a` is a strong liar iff
/// `a^d ≡ 1 (mod n)` or `a^(2^j * d) ≡ -1 (mod n)` for some `0 <= j < s`.
/// `n < 2` and even `n > 2` are decided directly; the base is reduced mod `n`,
/// and the trivial bases `a ≡ 0, 1, -1 (mod n)` (which witness nothing) pass.
///
/// # Examples
///
/// One round is not a proof: a single base can be fooled by a *strong liar*.
/// `2047 = 23 * 89` is composite, yet base 2 is a liar for it — base 3 is a
/// witness. [`is_prime`] runs twelve bases precisely so no `u64` composite
/// survives.
///
/// ```
/// use coprime::number::is_strong_probable_prime;
///
/// assert!(is_strong_probable_prime(97, 2)); // 97 is prime: every base passes
/// assert!(is_strong_probable_prime(2047, 2)); // base 2 is a strong *liar* here
/// assert!(!is_strong_probable_prime(2047, 3)); // base 3 *witnesses* 2047 composite
/// ```
#[must_use]
#[allow(clippy::many_single_char_names)] // n, a, s, d, x — canonical MR names
pub fn is_strong_probable_prime(n: u64, a: u64) -> bool {
    if n == 2 || n == 3 {
        return true;
    }
    if n < 2 || n & 1 == 0 {
        return false;
    }
    let a = a % n;
    if a <= 1 || a == n - 1 {
        return true;
    }
    let s = (n - 1).trailing_zeros();
    let d = (n - 1) >> s;
    let mut x = modpow(a, d, n);
    if x == 1 || x == n - 1 {
        return true;
    }
    // This loop records its result in `sprp` and `break`s, rather than the more
    // natural `return true` from inside the body. The two are behaviorally
    // identical, but the Aeneas→Lean verification (verify-lean/) extracts the
    // loop's structured form, and Aeneas's loop fixed-point translation rejects
    // an early `return` that exits a loop *and* the function (it handles `break`
    // cleanly). Keep the `break`/flag form — do not "simplify" back to `return`.
    let mut j = 1u32;
    let mut sprp = false;
    while j < s {
        x = modpow(x, 2, n); // x <- x^2 mod n
        if x == n - 1 {
            sprp = true;
            break;
        }
        j += 1;
    }
    sprp
}

/// Deterministic primality test for every `u64`: `true` iff `n` is prime.
///
/// Miller–Rabin against the first twelve primes `{2, 3, …, 37}`. The smallest
/// composite that is a strong pseudoprime to all twelve exceeds `3.19 * 10^23` —
/// far past `u64::MAX ≈ 1.8 * 10^19` — so no `u64` composite can survive all
/// twelve bases: the test is **deterministic**, not probabilistic. This is a
/// *cited* result (the twelve-base witness-set theorem), not a from-scratch
/// proof, and it holds **only** for this exact base set and **only** for
/// `n < 2^64`; changing the bases or the domain needs a different theorem. Small
/// `n` are settled directly by the trial-division pre-filter against those same
/// primes.
///
/// Verified exhaustively below `1_000_000` against an independent sieve, plus the
/// `ψ_t` / Carmichael adversarial set; the full `u64` guarantee rests on the
/// twelve-base witness-set theorem, and on demand on the exhaustive Feitsma sweep
/// (`make feitsma`). There is no Kani harness: the Miller–Rabin path runs `modpow`
/// in `u128`, which CBMC cannot bit-blast tractably, so the exhaustive sieve is
/// the right verification tier.
///
/// # Examples
///
/// ```
/// use coprime::number::is_prime;
///
/// assert!(is_prime(97));
/// assert!(!is_prime(91)); // 7 * 13
/// assert!(!is_prime(561)); // a Carmichael number — fools Fermat, not this test
/// assert!(is_prime(2_305_843_009_213_693_951)); // the Mersenne prime 2^61 - 1
/// ```
#[must_use]
pub fn is_prime(n: u64) -> bool {
    const WITNESSES: [u64; 12] = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37];
    if n < 2 {
        return false;
    }
    for &p in &WITNESSES {
        if n == p {
            return true;
        }
        if n.is_multiple_of(p) {
            return false;
        }
    }
    // n > 37 and coprime to all twelve bases; the witness loop now decides.
    WITNESSES.iter().all(|&a| is_strong_probable_prime(n, a))
}

/// Modular addition `a + b (mod n)` for `a, b < n`, formed without overflow (the
/// `wrapping_add` carry is folded back by the conditional subtract).
fn addmod(a: u64, b: u64, n: u64) -> u64 {
    let r = a.wrapping_add(b);
    if r < a || r >= n {
        r.wrapping_sub(n)
    } else {
        r
    }
}

/// Modular subtraction `a - b (mod n)` for `a, b < n`.
fn submod(a: u64, b: u64, n: u64) -> u64 {
    if a >= b { a - b } else { n - (b - a) }
}

/// `x · 2⁻¹ (mod n)` for odd `n` — `x/2` when `x` is even, else `(x+n)/2` formed
/// as `x/2 + n/2 + 1` (both halves floor; the `+1` carries the two odd low bits),
/// so it never overflows. The Lucas recurrence's increment step divides by 2, and
/// `n` odd makes 2 invertible.
fn half_mod(x: u64, n: u64) -> u64 {
    if x & 1 == 0 {
        x >> 1
    } else {
        (x >> 1) + (n >> 1) + 1
    }
}

/// The strong Lucas probable prime test with Selfridge's *Method A* parameters —
/// the Lucas half of Baillie–PSW (Baillie & Wagstaff, *Lucas Pseudoprimes*, Math.
/// Comp. 1980; Lucas-sequence doubling per Hamburg, eprint 2025/2083).
///
/// Picks the discriminant `D` as the first of `5, −7, 9, −11, …` for which the
/// Jacobi symbol `(D | n) = −1`, sets `P = 1` and `Q = (1 − D)/4`, writes
/// `n + 1 = d·2ˢ` with `d` odd, and computes the Lucas sequences `Uₖ, Vₖ (mod n)`
/// by the `V₂ₖ = Vₖ² − 2Qᵏ`, `U₂ₖ = UₖVₖ` doubling chain over the bits of `d`. `n`
/// passes iff `U_d ≡ 0 (mod n)`, or `V_{d·2ʳ} ≡ 0 (mod n)` for some `0 ≤ r < s` —
/// the *strong* form. Every odd prime passes; the composites that pass are *Lucas
/// pseudoprimes*, in practice disjoint from the base-2 strong pseudoprimes, which
/// is why [`bpsw`] combines the two.
///
/// `n < 2`, even `n`, small prime factors, and perfect squares are decided up
/// front (a square never reaches `(D | n) = −1`, so the Selfridge search would not
/// terminate; trial division keeps the search well-defined for tiny `n`).
///
/// # Examples
///
/// ```
/// use coprime::number::is_lucas_probable_prime;
///
/// assert!(is_lucas_probable_prime(97));
/// assert!(!is_lucas_probable_prime(91)); // 7 · 13
/// assert!(!is_lucas_probable_prime(2047)); // a base-2 strong PRP the Lucas test rejects
/// ```
#[must_use]
#[allow(clippy::many_single_char_names)] // n, r, d, s, u, v — canonical Lucas-test names
pub fn is_lucas_probable_prime(n: u64) -> bool {
    // The small primes that keep Selfridge's search well-defined: any `n` reaching
    // the search is coprime to every `|D|` it tries, so `(D | n) = 0` is then a
    // genuine compositeness proof, not the `n | D` ambiguity that bites tiny `n`.
    const SMALL: [u64; 15] = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47];
    if n < 2 {
        return false;
    }
    for &p in &SMALL {
        if n == p {
            return true;
        }
        if n.is_multiple_of(p) {
            return false;
        }
    }
    // perfect square ⇒ composite (the Selfridge search never reaches (D|n) = −1)
    let r = isqrt(n);
    if r * r == n {
        return false;
    }

    // Selfridge Method A: D = 5, −7, 9, −11, …; first with (D | n) = −1.
    let mut d_abs: u64 = 5;
    let mut d_neg = false;
    let (d_mod, q_mod) = loop {
        // (D | n) = (−1 | n)^[D<0] · (|D| | n); n is odd ⇒ `jacobi` is `Some`.
        let Some(jabs) = jacobi(d_abs, n) else {
            return false; // unreachable: n is odd past the trial division
        };
        if jabs == 0 {
            return false; // gcd(|D|, n) is a proper divisor of n ⇒ composite
        }
        let j = if d_neg && n % 4 == 3 { -jabs } else { jabs };
        if j == -1 {
            // D = ±d_abs ≡ 1 (mod 4); Q = (1 − D)/4 is an integer. Reduce mod n.
            let d_mod = if d_neg { n - d_abs } else { d_abs };
            let q_mod = if d_neg {
                (1 + d_abs) / 4 // D < 0 ⇒ Q = (1 + |D|)/4 > 0
            } else {
                n - (d_abs - 1) / 4 // D > 0 ⇒ Q = −(|D| − 1)/4 < 0
            };
            break (d_mod, q_mod);
        }
        d_abs += 2;
        d_neg = !d_neg;
    };

    // n + 1 = d · 2^s, d odd (u128 so n = u64::MAX does not overflow).
    let np1 = u128::from(n) + 1;
    let s = np1.trailing_zeros();
    #[allow(clippy::cast_possible_truncation)] // d = (n+1) >> s ≤ (n+1)/2 < 2^63 fits u64
    let d = (np1 >> s) as u64;

    // Doubling chain over the bits of d, most significant first, from (U₁, V₁) =
    // (1, P) = (1, 1) and Q¹ = q_mod. P = 1, so the increment step simplifies.
    let mut u = 1u64;
    let mut v = 1u64;
    let mut qk = q_mod;
    let mut bit = 63 - d.leading_zeros();
    while bit > 0 {
        bit -= 1;
        // index k → 2k
        u = mulmod(u, v, n);
        v = submod(mulmod(v, v, n), addmod(qk, qk, n), n); // Vₖ² − 2Qᵏ
        qk = mulmod(qk, qk, n);
        if (d >> bit) & 1 == 1 {
            // 2k → 2k+1: (U, V) ← ((U+V)/2, (D·U + V)/2); Q^(2k+1) = Q^(2k)·Q
            let uu = half_mod(addmod(u, v, n), n);
            let vv = half_mod(addmod(mulmod(d_mod, u, n), v, n), n);
            u = uu;
            v = vv;
            qk = mulmod(qk, q_mod, n);
        }
    }

    // Strong test: U_d ≡ 0, or some V_{d·2ʳ} ≡ 0 for 0 ≤ r < s.
    if u == 0 {
        return true;
    }
    let mut vr = v;
    let mut qr = qk;
    for _ in 0..s {
        if vr == 0 {
            return true;
        }
        vr = submod(mulmod(vr, vr, n), addmod(qr, qr, n), n);
        qr = mulmod(qr, qr, n);
    }
    false
}

/// The **Baillie–PSW** primality test: a base-2 strong Fermat (Miller–Rabin) round
/// composed with the strong Lucas test [`is_lucas_probable_prime`]. The two
/// pseudoprime sets are, in practice, disjoint — and Baillie, Fiori & Wagstaff
/// (Math. Comp. 2021) verified that **no** composite below `2⁶⁴` passes both — so
/// over `u64` this is a fast, **ceiling-free** primality test that needs neither
/// the deterministic twelve-base witness set nor the Feitsma enumeration that
/// [`is_prime`] relies on.
///
/// Its full-domain correctness is, like those, an empirical/conjectural result
/// rather than a from-scratch proof; a *constructive* primality proof is the
/// Pocklington-certificate route. `bpsw` and [`is_prime`] agree on every `u64`.
///
/// # Examples
///
/// ```
/// use coprime::number::bpsw;
///
/// assert!(bpsw(2_305_843_009_213_693_951)); // the Mersenne prime 2^61 − 1
/// assert!(!bpsw(2047)); // 23 · 89: a base-2 strong PRP, but the Lucas leg rejects it
/// assert!(!bpsw(561)); // a Carmichael number
/// ```
#[must_use]
pub fn bpsw(n: u64) -> bool {
    if n == 2 {
        return true;
    }
    if n < 2 || n & 1 == 0 {
        return false;
    }
    is_strong_probable_prime(n, 2) && is_lucas_probable_prime(n)
}

/// The smallest prime `>= n`, or `None` if there is none in `u64` (i.e. `n`
/// exceeds the largest `u64` prime). Deterministic and RNG-free: it scans odd
/// candidates with [`is_prime`], so the result is a certified prime.
///
/// # Examples
///
/// ```
/// use coprime::number::next_prime;
///
/// assert_eq!(next_prime(14), Some(17));
/// assert_eq!(next_prime(17), Some(17)); // n itself, when n is prime
/// assert_eq!(next_prime(0), Some(2));
/// ```
#[must_use]
pub fn next_prime(n: u64) -> Option<u64> {
    if n <= 2 {
        return Some(2); // 2 is the only even prime and the least prime
    }
    let mut c = n | 1; // least odd >= n
    loop {
        if is_prime(c) {
            return Some(c);
        }
        c = c.checked_add(2)?; // next odd; None on overflow = no prime <= u64::MAX
    }
}

/// A source of uniformly random `u64` words. coprime is RNG-agnostic — *mechanism,
/// not policy*: the caller supplies the source. For anything secret (key
/// material), inject a cryptographically secure RNG; for tests or non-secret
/// values any PRNG will do. The security of generated values is the source's
/// responsibility, not coprime's.
///
/// # Examples
///
/// Implement it over any word stream — here a tiny LCG (a PRNG, *not* for
/// secrets):
///
/// ```
/// use coprime::number::Rng;
///
/// struct Lcg(u64);
/// impl Rng for Lcg {
///     fn next_u64(&mut self) -> u64 {
///         self.0 = self.0.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
///         self.0
///     }
/// }
///
/// let mut rng = Lcg(0);
/// assert_eq!(rng.next_u64(), 1); // 0 * a + 1
/// ```
pub trait Rng {
    /// The next uniformly random `u64`.
    fn next_u64(&mut self) -> u64;
}

/// A uniformly random prime of exactly `bits` bits (`2 <= bits <= 64`), drawn by
/// rejection: sample a `bits`-bit odd candidate (top and bottom bits set) and
/// retry until [`is_prime`] accepts it. `None` if `bits` is out of range.
///
/// Termination is guaranteed: by Bertrand's postulate the interval
/// `[2^(bits-1), 2^bits)` always contains a prime for `bits >= 2`, and by the
/// prime number theorem the expected number of trials is `O(bits)`. The result
/// is certified prime (`is_prime` is exact for `u64`); its *unpredictability*
/// is exactly as good as the injected `rng`.
///
/// # Examples
///
/// ```
/// use coprime::number::{is_prime, random_prime, Rng};
///
/// struct Lcg(u64);
/// impl Rng for Lcg {
///     fn next_u64(&mut self) -> u64 {
///         self.0 = self.0.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
///         self.0
///     }
/// }
///
/// let mut rng = Lcg(0xDEAD_BEEF);
/// let p = random_prime(&mut rng, 16).expect("16 is in 2..=64");
/// assert!(is_prime(p)); // certified, not probable
/// assert_eq!(64 - p.leading_zeros(), 16); // exactly 16 bits
/// assert_eq!(random_prime(&mut rng, 1), None); // out of the 2..=64 range
/// ```
#[must_use]
pub fn random_prime<R: Rng>(rng: &mut R, bits: u32) -> Option<u64> {
    if !(2..=64).contains(&bits) {
        return None;
    }
    let top = 1u64 << (bits - 1);
    let mask = if bits == 64 {
        u64::MAX
    } else {
        (1u64 << bits) - 1
    };
    loop {
        let n = (rng.next_u64() & mask) | top | 1; // bits-bit, high bit + odd
        if is_prime(n) {
            return Some(n);
        }
    }
}
