//! Integer factorization — Brent–Pollard rho with a certified-prime base case.

use super::modular::mulmod;
use crate::number::{gcd, is_prime};

/// The prime factorization of a `u64`, returned by [`factorize`]: the distinct
/// prime factors paired with their exponents, ascending by prime. A `u64` has at
/// most fifteen distinct prime factors (the product of the first sixteen primes
/// exceeds `2^64`), so the capacity is fixed and no allocation is needed.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Factorization {
    factors: [(u64, u32); 15],
    len: usize,
}

impl Factorization {
    /// The `(prime, exponent)` pairs, ascending by prime. Empty for `0` and `1`.
    #[must_use]
    pub fn as_slice(&self) -> &[(u64, u32)] {
        &self.factors[..self.len]
    }

    /// Record one occurrence of prime `p`, coalescing with any existing entry.
    /// A `u64` has at most fifteen distinct prime factors, so `push` adds a new
    /// entry at most fifteen times: the highest index written is `14`, and `len`
    /// reaches `15` (= capacity) only after the final push, never before — so the
    /// `self.factors[self.len]` write is always in bounds. This buffer is
    /// exact-fit (15 for 15), unlike `cfrac`'s headroom array (93 for 92).
    fn push(&mut self, p: u64) {
        let mut i = 0;
        while i < self.len {
            if self.factors[i].0 == p {
                self.factors[i].1 += 1;
                return;
            }
            i += 1;
        }
        self.factors[self.len] = (p, 1);
        self.len += 1;
    }
}

/// One step of the rho iteration polynomial: `x² + c (mod n)`, formed in `u128`
/// so it never overflows; the result is `< n <= u64::MAX` so the cast is exact.
#[allow(clippy::cast_possible_truncation)] // x²+c mod n < n <= u64::MAX fits u64
fn rho_step(x: u64, c: u64, n: u64) -> u64 {
    ((u128::from(x) * u128::from(x) + u128::from(c)) % u128::from(n)) as u64
}

/// Advance `y0` by `count` rho steps of `x² + c (mod n)`. Returns `(y0, y)` —
/// the unmoved starting point (the round anchor `x` in [`brent`]) alongside the
/// advanced value.
fn rho_advance(y0: u64, count: u64, c: u64, n: u64) -> (u64, u64) {
    let mut y = y0;
    let mut i = 0u64;
    while i < count {
        y = rho_step(y, c, n);
        i += 1;
    }
    (y0, y)
}

/// One Brent batch of `lim` rho steps from `y0`, folding each `|x − y|` into the
/// running product `q (mod n)`. Returns `(y0, y, q)` — the start of the batch
/// (the checkpoint [`brent`]'s collision backtrack restarts from) and the
/// advanced values.
#[allow(clippy::many_single_char_names)] // x, y, q, c, n, j — canonical rho names
fn rho_batch(x: u64, y0: u64, q: u64, lim: u64, c: u64, n: u64) -> (u64, u64, u64) {
    let mut y = y0;
    let mut q = q;
    let mut j = 0u64;
    while j < lim {
        y = rho_step(y, c, n);
        q = mulmod(q, x.abs_diff(y), n);
        j += 1;
    }
    (y0, y, q)
}

/// Brent's batched-gcd pass for one doubling round `r`: walk `y` up to `r` steps
/// in batches of (at most) 128, folding each batch's distances into `q` and
/// testing `gcd(q, n)` after each batch; stop early once the gcd is nontrivial.
/// Entered only with `g == 1`, and `r >= 1` means at least one batch runs.
/// Returns `(y, ys, q, g)` — `ys` is the start of the last batch walked, the
/// checkpoint [`brent`]'s collision backtrack restarts from.
#[allow(clippy::many_single_char_names)] // x, y, r, q, g, c, k — canonical rho names
fn rho_round(x: u64, y: u64, q: u64, r: u64, c: u64, n: u64) -> (u64, u64, u64, u64) {
    let mut y = y;
    let mut ys = y;
    let mut q = q;
    let mut g = 1u64;
    let mut k = 0u64;
    while k < r && g == 1 {
        let lim = core::cmp::min(128, r - k);
        let (y_start, y_next, q_next) = rho_batch(x, y, q, lim, c, n);
        ys = y_start;
        y = y_next;
        q = q_next;
        g = gcd(q, n);
        k += 128;
    }
    (y, ys, q, g)
}

/// One Brent–Pollard rho attempt on the odd composite `n` with the polynomial
/// `x² + c (mod n)`: returns a nontrivial factor, or `None` if this `c` fails
/// (Pollard's rho with Brent 1980's batched-gcd cycle detection). The
/// `x² + c` is formed in `u128` so it never overflows; the per-step `|x - y|` is
/// a `u64` absolute difference (no underflow), accumulated into `q` mod `n`.
///
/// **Partial-correctness postcondition:** if this returns `Some(d)` then
/// `1 < d < n` and `d ∣ n`. This holds *by construction* and is independent of
/// the iteration polynomial, the doubling schedule, and the batch size: every
/// candidate is `gcd(_, n)`, the `while g == 1` loop exits only on `g ≠ 1`, and
/// the final `g == n` guard maps the degenerate case to `None`. So those
/// parameters steer only *which* divisor is found and *how fast* — never whether
/// a returned value divides `n`, the sole property [`factorize`] relies on. That
/// is precisely why mutating them is invisible to any output test, and why those
/// mutants are dispositioned by this contract rather than chased (see Section B
/// of `.cargo/mutants.toml`), not silenced for want of a test.
///
/// **Termination is probabilistic, not proven.** rho is a Las Vegas algorithm:
/// it finds a factor in heuristically `O(n^¼)` expected steps, but no
/// deterministic worst-case bound is established here, and the differential
/// `factorize` tests demonstrate success on hard inputs without certifying
/// termination. The deterministic safety net is structural: a failed attempt
/// returns `None` and [`pollard_split`] retries with a fresh constant.
///
/// The rho machinery is split into [`rho_step`] / [`rho_advance`] /
/// [`rho_batch`] / [`rho_round`] so each function carries at most one loop
/// and every checkpoint (`x`, `ys`) comes back as a callee return value
/// rather than an in-loop variable copy — shapes the Aeneas symbolic
/// interpreter's loop fixed-point computation handles (the same extraction
/// medicine as the SPRP and `sqrtmod` while-loop rewrites). Result-identical:
/// the step sequence, batch boundaries, gcd schedule, and checkpoints are
/// unchanged on every input.
#[allow(clippy::many_single_char_names)] // x, y, r, q, g, c, k — canonical rho names
fn brent(n: u64, c: u64) -> Option<u64> {
    let mut y = 2u64;
    let mut x = y;
    let mut ys = y;
    let mut r = 1u64;
    let mut q = 1u64;
    let mut g = 1u64;
    while g == 1 {
        let (anchor, y_adv) = rho_advance(y, r, c, n);
        x = anchor;
        y = y_adv;
        let (y_round, checkpoint, q_round, g_round) = rho_round(x, y, q, r, c, n);
        y = y_round;
        ys = checkpoint;
        q = q_round;
        g = g_round;
        r *= 2;
    }
    if g == n {
        // A batch hid two factors in `q`; re-walk one step at a time to recover one.
        loop {
            ys = rho_step(ys, c, n);
            g = gcd(x.abs_diff(ys), n);
            if g > 1 {
                break;
            }
        }
    }
    if g == n { None } else { Some(g) }
}

/// A nontrivial divisor of the odd composite `n`, retrying the rho polynomial
/// constant on failure (any `c != 0` and `c != -2`).
///
/// **Postcondition:** returns `d` with `1 < d < n` and `d ∣ n` — inherited
/// directly from [`brent`]'s partial-correctness contract. **Termination** rests
/// on rho's probabilistic argument plus this retry loop, not a deterministic
/// bound (see [`brent`]). No realized `factorize` input drives the retry past
/// `c = 1` (the cofactor reaching rho always splits on the first constant), so
/// the unit probe `pollard_split_retries_when_first_constant_fails` reaches the
/// arm directly with `n = 25` and pins recovery-on-failure. The finer
/// `c`-trajectory mutants (the exact increment, the forbidden-`c` skip) remain
/// live MISSED signal rather than excluded: they change only *which* constant
/// succeeds, which this contract leaves free — and `c += 1 → c -= 1` is a latent
/// concern (it can reach the forbidden constant 0) that no input triggers, kept
/// visible deliberately.
fn pollard_split(n: u64) -> u64 {
    let mut c = 1u64;
    loop {
        if let Some(d) = brent(n, c) {
            return d;
        }
        c += 1;
        if c == n - 2 {
            c += 1; // skip the forbidden c ≡ -2 (mod n)
        }
    }
}

/// Recursively split `n` into primes, pushing each into `acc`. Base case:
/// [`is_prime`] (exact for all `u64`), so every emitted factor is certified
/// prime. Depth `<= 64`: each split strictly shrinks the operand.
fn factor_recurse(n: u64, acc: &mut Factorization) {
    if is_prime(n) {
        acc.push(n);
        return;
    }
    let d = pollard_split(n);
    factor_recurse(d, acc);
    factor_recurse(n / d, acc);
}

/// The prime factorization of `n` as ascending `(prime, exponent)` pairs (see
/// [`Factorization`]). Small primes are peeled by trial division; the surviving
/// cofactor is split by Brent–Pollard rho with [`is_prime`] as the base case, so
/// every returned factor is certified prime. `0` and `1` have no prime factors
/// and yield an empty factorization.
///
/// # Examples
///
/// ```
/// use coprime::number::factorize;
///
/// assert_eq!(factorize(360).as_slice(), &[(2, 3), (3, 2), (5, 1)]); // 2³·3²·5
/// assert_eq!(factorize(1_000_003).as_slice(), &[(1_000_003, 1)]); // prime
/// assert!(factorize(1).as_slice().is_empty());
/// ```
#[must_use]
#[allow(clippy::many_single_char_names)] // n, d, s, i, j — loop/index locals
pub fn factorize(n: u64) -> Factorization {
    let mut result = Factorization {
        factors: [(0, 0); 15],
        len: 0,
    };
    let mut n = n;
    if n < 2 {
        return result; // 0 and 1: the empty product
    }
    while n.is_multiple_of(2) {
        result.push(2);
        n /= 2;
    }
    let mut d = 3u64;
    while d < 1 << 10 && d <= n / d {
        while n.is_multiple_of(d) {
            result.push(d);
            n /= d;
        }
        d += 2;
    }
    if n > 1 {
        factor_recurse(n, &mut result);
    }
    sort_factors(&mut result.factors[..result.len]);
    result
}

/// Sort `factors` in place, ascending by prime (`.0`), by insertion sort.
///
/// **Postcondition:** on return `factors` is non-decreasing by `.0` and is a
/// permutation of its argument. The trial-division peels arrive already
/// ascending and Brent–Pollard rho appends at most a few larger primes, so the
/// realized inputs are nearly sorted — but the postcondition holds for arbitrary
/// initial order, which is what the proof establishes.
///
/// Extracted as a named unit precisely so it has a postcondition a proof can
/// discharge: `sort_factors_ascending_perm` (`proofs/factor_sort.rs`) proves both
/// clauses over symbolic arrays. That proof is what settles the surviving
/// insertion-sort mutants by *reason* rather than by a single KAT's luck — the
/// inner-loop decrement `j -= 1` is load-bearing (a mutation leaving `j`
/// unchanged half-finishes a multi-step bubble and breaks sortedness, so the
/// proof fails), while the guard `>` vs `>=` is provably equivalent (both leave
/// the output sorted), so no test can or should distinguish it.
#[allow(clippy::many_single_char_names)] // i, j — canonical insertion-sort indices
pub(crate) fn sort_factors(factors: &mut [(u64, u32)]) {
    let mut i = 1;
    while i < factors.len() {
        let mut j = i;
        while j > 0 && factors[j - 1].0 > factors[j].0 {
            factors.swap(j - 1, j);
            j -= 1;
        }
        i += 1;
    }
}

// ── White-box probes (private-function tests) ─────────────────────────────────
//
// The output-level tests in tests.rs prove that `factorize(n)` returns the right
// factors, but a multi-path algorithm can produce the right output with a broken
// internal branch: `pollard_split`'s c=2 retry masks a broken g==n recovery in
// `brent`, so the output KAT cannot tell the two apart. These probes call the
// private functions directly to pin the correctness-bearing branches the output
// tests cannot reach.
#[cfg(test)]
mod tests {
    use super::{brent, pollard_split};

    /// brent's `g == n` batch-collision recovery (the `loop { ys = f(ys); … }`
    /// backtrack) returns a nontrivial factor rather than `None`. n = 1031·1039
    /// drives g==n on the first rho constant c=1. The output-level
    /// `factorize(1_071_209)` KAT cannot kill the `156:10 ==→!=` mutation because
    /// `pollard_split` retries with c=2 and still yields the correct factors;
    /// this probe is the only check that pins the recovery branch itself.
    #[test]
    fn brent_recovers_from_batch_collision() {
        let n = 1_071_209u64; // 1031 · 1039; c=1 triggers the g==n batch-collision path
        let d = brent(n, 1).expect("g==n backtrack must recover a nontrivial factor, not None");
        assert!(
            d > 1 && d < n && n.is_multiple_of(d),
            "d={d} must be a proper divisor of n={n}"
        );
    }

    /// `pollard_split`'s retry arm: when `brent(n, 1)` fails (returns `None`), the
    /// loop must try another constant and still return a proper divisor. n = 25
    /// (= 5²) is the smallest n with `brent(25, 1) == None`, so it is the only
    /// kind of input that executes the retry at all — `factorize` peels 25 by
    /// trial division and never calls `pollard_split`, leaving the arm dead to
    /// every output test. This probe pins exactly the behaviour that matters:
    /// recovery on first-constant failure.
    ///
    /// It does NOT claim to kill the retry's *trajectory* mutants (the precise
    /// `c += 1` step at line 190, the forbidden-`c` skip at 191–192). Those alter
    /// only WHICH constant eventually succeeds — `pollard_split`'s contract leaves
    /// that free — so they survive as live signal rather than being excluded; the
    /// `c += 1 → c -= 1` variant is a latent concern (it can reach the forbidden
    /// constant 0) that no realized input triggers, kept visible on purpose. Only
    /// `c *= 1` (spin on the failing c=1) is killed here, by `--timeout`.
    #[test]
    fn pollard_split_retries_when_first_constant_fails() {
        let n = 25u64;
        assert!(
            brent(n, 1).is_none(),
            "precondition: c=1 must fail for n={n} so the retry arm is exercised"
        );
        let d = pollard_split(n);
        assert!(
            d > 1 && d < n && n.is_multiple_of(d),
            "pollard_split({n})={d} must be a proper divisor, found via the c>1 retry"
        );
    }
}
