//! Quadratic residues — Jacobi and Kronecker symbols, and modular square roots.

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

/// Jacobi symbol core for odd `n >= 1`, by quadratic reciprocity, iterated:
/// pull factors of 2 from `a` flipping the sign per `(2|n)`,
/// then swap `(a, n)` with the reciprocity sign flip, until `a == 0`.
#[allow(clippy::many_single_char_names)] // a, n, s are the canonical symbol names
fn jacobi_odd(mut a: u64, mut n: u64) -> i8 {
    a %= n;
    let mut s: i8 = 1;
    while a != 0 {
        while a & 1 == 0 {
            a >>= 1;
            // (2|n) = -1 exactly when n ≡ 3 or 5 (mod 8).
            if n % 8 == 3 || n % 8 == 5 {
                s = -s;
            }
        }
        (a, n) = (n, a); // swap for reciprocity (a plain swap; no byte-loop for CBMC)
        // reciprocity: (a|n)(n|a) = -1 exactly when a ≡ n ≡ 3 (mod 4).
        if a % 4 == 3 && n % 4 == 3 {
            s = -s;
        }
        a %= n;
    }
    if n == 1 { s } else { 0 }
}

/// The Jacobi symbol `(a | n)` for odd `n`: `0` when `gcd(a, n) != 1`, otherwise
/// `±1`. When `n` is an odd prime this *is* the Legendre symbol — `+1` iff `a` is
/// a nonzero quadratic residue modulo `n`, `-1` if it is a non-residue.
///
/// Computed by quadratic reciprocity — no factorization of
/// `n`, `O((lg n)^2)` work. `None` when `n` is even or zero, where the symbol is
/// undefined.
///
/// # Examples
///
/// ```
/// use coprime::number::jacobi;
///
/// assert_eq!(jacobi(2, 7), Some(1)); // 2 = 3^2 (mod 7): a residue (Legendre, 7 prime)
/// assert_eq!(jacobi(3, 7), Some(-1)); // 3 is a non-residue mod 7
/// assert_eq!(jacobi(6, 9), Some(0)); // gcd(6, 9) = 3 != 1
/// assert_eq!(jacobi(1, 45), Some(1)); // (1|n) = 1 for every odd n
/// assert_eq!(jacobi(5, 4), None); // even n: the symbol is undefined
/// ```
#[must_use]
pub fn jacobi(a: u64, n: u64) -> Option<i8> {
    if n == 0 || n & 1 == 0 {
        return None;
    }
    Some(jacobi_odd(a, n))
}

/// The Kronecker symbol `(a | n)`, extending [`jacobi`] to **every** `n` —
/// even, zero, and odd. For odd `n` it equals `jacobi(a, n)`; the extension
/// adds the `(a | 2)` supplement and the boundary values `(a | 0)` and `(a | 1)`.
///
/// Writing `n = 2^e * m` with `m` odd, `(a | n) = (a | 2)^e * (a | m)`, where the
/// odd part `(a | m)` is the Jacobi symbol and the supplement is
/// `(a | 2) = 0` if `a` is even, `+1` if `a ≡ ±1 (mod 8)`, `-1` if `a ≡ ±3
/// (mod 8)`. The boundaries are `(a | 0) = 1` iff `a = 1` (else `0`) and
/// `(a | 1) = 1`. Returns `0` exactly when `gcd(a, n) != 1`.
///
/// # Examples
///
/// ```
/// use coprime::number::kronecker;
///
/// assert_eq!(kronecker(3, 2), -1); // 3 ≡ 3 (mod 8): (3|2) = -1
/// assert_eq!(kronecker(7, 2), 1); //  7 ≡ -1 (mod 8): (7|2) = 1
/// assert_eq!(kronecker(5, 12), -1); // (5|4)·(5|3) = 1·(-1)
/// assert_eq!(kronecker(1, 0), 1);
/// assert_eq!(kronecker(2, 0), 0); // gcd(2, 0) = 2 != 1
/// ```
#[must_use]
#[allow(clippy::many_single_char_names)] // a, n, e, m, s, t are the canonical symbol names
pub fn kronecker(a: u64, n: u64) -> i8 {
    if n == 0 {
        return i8::from(a == 1); // (a|0) = 1 iff |a| = 1; for u64 only a = 1
    }
    let e = n.trailing_zeros(); // n = 2^e * m, m odd
    let m = n >> e;
    let mut s: i8 = 1;
    if e > 0 {
        if a & 1 == 0 {
            return 0; // a even and n even ⇒ gcd(a, n) >= 2
        }
        // (a|2) = +1 for a ≡ ±1 (mod 8), -1 for a ≡ ±3 (mod 8); odd `a` here.
        let t: i8 = if a % 8 == 1 || a % 8 == 7 { 1 } else { -1 };
        if e % 2 == 1 {
            s = t; // (a|2)^e: e even squares to +1, e odd is (a|2)
        }
    }
    // The odd part reuses jacobi_odd, which returns 1 for m == 1 — covering (a|1).
    s * jacobi_odd(a, m)
}

/// A square root of `a` modulo a prime `p`: the smaller `x` in `[0, p)` with
/// `x*x ≡ a (mod p)`, or `None` when `a` is a quadratic non-residue mod `p`.
/// `a = 0` yields `0`.
///
/// Defined only for prime `p`: square roots modulo a composite are as hard as
/// factoring it. This is *enforced* — any non-prime `p`
/// (composite, `0`, or `1`) returns `None` rather than attempt it, so the
/// function is total and runs in bounded time for every `u64` (a non-prime
/// modulus could otherwise drive the non-residue search over `O(p)` candidates).
/// For `p ≡ 3 (mod 4)` the closed form `a^((p+1)/4)` gives a root directly;
/// otherwise the general Tonelli–Shanks loop runs, with a
/// quadratic non-residue found by deterministic upward search (its existence and
/// smallness are guaranteed for prime `p`).
///
/// # Examples
///
/// ```
/// use coprime::number::sqrtmod;
///
/// assert_eq!(sqrtmod(2, 7), Some(3)); // 3^2 = 9 ≡ 2 (mod 7); p ≡ 3 (mod 4)
/// assert_eq!(sqrtmod(2, 17), Some(6)); // 6^2 = 36 ≡ 2 (mod 17); Tonelli–Shanks
/// assert_eq!(sqrtmod(5, 7), None); // 5 is a non-residue mod 7
/// assert_eq!(sqrtmod(0, 13), Some(0));
/// let x = sqrtmod(10, 13).unwrap();
/// assert_eq!(x * x % 13, 10);
/// ```
#[must_use]
#[allow(clippy::many_single_char_names)] // p, a, q, s, m, c, t, r, z, i, b — TS names
pub fn sqrtmod(a: u64, p: u64) -> Option<u64> {
    if p == 2 {
        return Some(a & 1); // x^2 ≡ a (mod 2) ⇔ x ≡ a (mod 2)
    }
    // Enforce the prime-modulus contract by construction: square roots modulo a
    // composite are as hard as factoring it, and a non-prime p
    // would let the non-residue search below run over O(p) candidates. Any
    // non-prime modulus fails closed, keeping the function total and bounded.
    if p < 2 || !is_prime(p) {
        return None;
    }
    let a = a % p;
    if a == 0 {
        return Some(0);
    }
    // p is an odd prime, so a coprime to p has a Legendre symbol of ±1.
    if jacobi_odd(a, p) != 1 {
        return None; // quadratic non-residue
    }
    if p % 4 == 3 {
        let r = modpow(a, (p + 1) / 4, p);
        return Some(r.min(p - r));
    }
    // p ≡ 1 (mod 4): Tonelli–Shanks. Write p - 1 = q * 2^s with q odd.
    let s = (p - 1).trailing_zeros();
    let q = (p - 1) >> s;
    // A quadratic non-residue z; exists and is small for prime p.
    let mut z = 2u64;
    while jacobi_odd(z, p) != -1 {
        z += 1;
        if z == p {
            return None; // p is not prime — no non-residue exists
        }
    }
    let mut m = s;
    let mut c = modpow(z, q, p);
    let mut t = modpow(a, q, p);
    let mut r = modpow(a, q.div_ceil(2), p); // (q+1)/2; q is odd so this is exact
    loop {
        if t == 1 {
            return Some(r.min(p - r));
        }
        // least i in (0, m) with t^(2^i) == 1. The `i + 1 < m` bound replaces an
        // `if i == m { return None }` inside the loop body: a return that crosses
        // two loop levels is control flow the Aeneas extraction cannot translate
        // ("breaks to outer loops"), exactly like the SPRP while-loop rewrite.
        // Result-identical on every reachable state (m >= 2, forced by t != 1
        // with ord(t) | 2^(m-1)): None iff t^(2^j) != 1 for all j = 1..m-1.
        let mut i = 1u32;
        let mut t2 = mulmod(t, t, p);
        while t2 != 1 && i + 1 < m {
            t2 = mulmod(t2, t2, p);
            i += 1;
        }
        if t2 != 1 {
            return None; // unreachable for a true residue mod prime p
        }
        let b = modpow(c, 1u64 << (m - i - 1), p);
        m = i;
        c = mulmod(b, b, p);
        t = mulmod(t, c, p);
        r = mulmod(r, b, p);
    }
}
