// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// IntPrime.hpp
// Primality testing (Miller-Rabin, Fermat) for Int

#ifndef SANGI_INT_PRIME_HPP
#define SANGI_INT_PRIME_HPP

#include "IntBase.hpp"
#include <vector>
#include <utility>
#include <atomic>

namespace sangi {

/**
 * @brief Primality testing algorithms
 *
 * Provides probabilistic primality testing (Miller-Rabin) and a deterministic small-prime check.
 * Used in RSA cryptography, elliptic-curve cryptography, integer factorization, and similar areas.
 */
class IntPrime {
public:
    /**
     * @brief Probabilistic primality test (Miller-Rabin test)
     *
     * Performs a probabilistic primality test based on the Miller-Rabin algorithm.
     * After k rounds, the probability of a false positive is at most 4^(-k).
     *
     * Algorithm:
     *   Decompose n-1 = 2^s * d (with d odd).
     *   For each of k rounds, with base a (2 <= a <= n-2):
     *   - If a^d mod n == 1, advance to the next round
     *   - If a^(2^r * d) mod n == -1 for some r in 0..s-1, advance to the next round
     *   - Otherwise, n is composite
     *
     * @param n integer to test
     * @param k number of rounds (default 20, false-positive probability <= 2^(-40))
     * @return true: probably prime, false: definitely composite
     *
     * Special cases:
     *   isMillerRabinPrime(n < 2) = false
     *   isMillerRabinPrime(2) = true
     *   isMillerRabinPrime(even) = false
     *   isMillerRabinPrime(NaN) = false
     *   isMillerRabinPrime(inf) = false
     *
     * Notes:
     *   - With k = 20 the false-positive probability is about 2^(-40) ~ 10^(-12)
     *   - For cryptographic use, k >= 40 is recommended (false-positive probability <= 2^(-80))
     *   - This is not a deterministic test; use another algorithm if certainty is required
     */
    static bool isMillerRabinPrime(const Int& n, int k = 20);

    /**
     * @brief Probabilistic primality test based on Fermat's little theorem
     *
     * Fermat's little theorem: if p is prime then a^(p-1) ≡ 1 (mod p) for gcd(a,p)=1.
     * Uses the converse heuristically for primality testing.
     *
     * Notes:
     *   - May incorrectly classify Carmichael numbers (pseudoprimes) as prime
     *   - Miller-Rabin is more reliable, so prefer Miller-Rabin in general
     *   - Provided for educational use and quick checks
     *
     * @param n integer to test
     * @param k number of rounds (default 5)
     * @return true: probably prime, false: definitely composite
     *
     * Special cases:
     *   isFermatPrime(n < 2) = false
     *   isFermatPrime(2) = true
     *   isFermatPrime(even) = false
     *   isFermatPrime(NaN) = false
     */
    static bool isFermatPrime(const Int& n, int k = 5);

    /**
     * @brief General-purpose primality test (recommended)
     *
     * Probabilistic primality test using Miller-Rabin.
     * Alias for isMillerRabinPrime.
     *
     * @param n integer to test
     * @param k number of rounds (default 20)
     * @return true: probably prime, false: definitely composite
     */
    static bool isProbablePrime(const Int& n, int k = 20) {
        return isMillerRabinPrime(n, k);
    }

    /**
     * @brief Deterministic primality test (APRCL algorithm)
     *
     * Deterministic primality proof via the Adleman-Pomerance-Rumely-Cohen-Lenstra method.
     * Unlike Miller-Rabin, it returns a mathematically exact (non-probabilistic) result.
     *
     * Algorithm:
     *   1. Perform Jacobi-sum tests over the cyclotomic ring Z[ζ_q]/(n)
     *   2. If all tests pass, the prime factors of n must lie in a restricted residue class
     *   3. Final factor search: if no non-trivial factor exists, n is confirmed prime
     *
     * @param n integer to test
     * @return true: prime (confirmed), false: composite (confirmed)
     *
     * Notes:
     *   - Substantially slower than Miller-Rabin (seconds to minutes)
     *   - Use isProbablePrime when speed matters (e.g. cryptographic key generation)
     *   - Use this routine when a mathematical proof or verification is required
     */
    static bool isProvablePrime(const Int& n);

    /**
     * @brief Returns the smallest prime greater than n
     *
     * Searches for the next prime by testing n+1, n+2, ... in order.
     *
     * @param n reference integer
     * @param k Miller-Rabin round count (default 20)
     * @return smallest prime greater than n
     *
     * Special cases:
     *   nextPrime(n < 2) = 2
     *   nextPrime(2) = 3
     *   nextPrime(NaN) = NaN
     *   nextPrime(inf) = inf
     *
     * Notes:
     *   - The search may take significant time for very large numbers
     *   - By the prime number theorem, the prime gap near n is roughly ln(n)
     */
    static Int nextPrime(const Int& n, int k = 20);

    /**
     * @brief Returns the largest prime less than n
     *
     * Searches for the previous prime by testing n-1, n-2, ... in order.
     *
     * @param n reference integer
     * @param k Miller-Rabin round count (default 20)
     * @return largest prime less than n
     *
     * Special cases:
     *   prevPrime(n <= 2) = NaN (no prime is smaller than 2)
     *   prevPrime(3) = 2
     *   prevPrime(NaN) = NaN
     *   prevPrime(inf) = NaN
     *
     * Notes:
     *   - The search may take significant time for very large numbers
     */
    static Int prevPrime(const Int& n, int k = 20);

    /**
     * @brief Factor search via Fermat's method
     *
     * Looks for n in the form n = a^2 - b^2 = (a+b)(a-b).
     * Fast when n is the product of two nearby primes.
     *
     * @param n composite (odd, n > 1)
     * @param max_iterations maximum iteration count (default 10000)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_Fermat(const Int& n, uint64_t max_iterations = 10000);

    /**
     * @brief Factor search via Pollard's rho
     *
     * Factor search using the pseudo-random sequence x_{k+1} = x_k^2 + c (mod n).
     * Uses Brent's improved cycle-detection variant.
     * Expected complexity: O(n^{1/4})
     *
     * @param n composite (n > 1)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_PollardRho(const Int& n);

    /**
     * @brief Factor search via Pollard's p-1 method
     *
     * Finds prime factors p such that p-1 is B-smooth (all prime factors <= B).
     *
     * Algorithm:
     *   Build M = lcm(2, 3, ..., B) incrementally and compute
     *   gcd(a^M - 1, n). If p | n and (p-1) | M, then
     *   a^M ≡ 1 (mod p), so gcd > 1 and a factor is found.
     *
     * @param n composite (n > 1)
     * @param B1 Stage 1 smooth bound (default 10000)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_PollardP1(const Int& n, uint64_t B1 = 10000);

    /**
     * @brief Factor search via ECM (Lenstra's elliptic-curve method)
     *
     * Lenstra's elliptic-curve factorization method.
     * Uses Montgomery-form elliptic curves By^2 = x^3 + Ax^2 + x.
     * Effective for factors of roughly 20 to 60 digits.
     * Expected complexity: O(exp(sqrt(2 ln p ln ln p))) (p = smallest prime factor)
     *
     * @param n composite (n > 1)
     * @param curves number of curves to try (default 100)
     * @param B1 Stage 1 smooth bound (default 10000)
     * @param B2 Stage 2 bound (default 100*B1)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_ECM(const Int& n, int curves = 100,
                              uint64_t B1 = 10000, uint64_t B2 = 0,
                              const std::atomic<bool>* cancel = nullptr);


    /**
     * @brief Factor search via Williams' p+1 method
     *
     * Finds prime factors p such that p+1 is B-smooth. The dual of the p-1 method.
     * Compute gcd(V_M - 2, n) using the Lucas sequence V_k(a) mod n.
     *
     * Algorithm:
     *   1. Lucas sequence V_k: V_0=2, V_1=a, V_{k+1}=a*V_k - V_{k-1}
     *   2. Compute V_M mod n quickly via a binary ladder
     *   3. Factor detection via g = gcd(V_M - 2, n)
     *
     * @param n composite (odd, n > 1)
     * @param B1 smooth bound (default 10000)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_WilliamsP1(const Int& n, uint64_t B1 = 10000);

    /**
     * @brief Factor search via SQUFOF (Square Form Factorization)
     *
     * Shanks' square-form factorization method, based on continued-fraction expansion.
     * O(n^{1/4}) time, O(1) memory. Fast for small to mid-range inputs (up to ~18 digits).
     *
     * @param n composite (odd, n > 1, not a perfect square)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_SQUFOF(const Int& n);

    /**
     * @brief Factor search via Pollard rho (Montgomery multiplication variant)
     *
     * A fast Pollard rho using Montgomery modular multiplication.
     *
     * Algorithm:
     *   Pick R = 2^k > n and convert to Montgomery representation.
     *   Compute x_{k+1} = x_k^2 + 1 (mod n) via Montgomery multiplication.
     *   Iterate Floyd's cycle detection until gcd(x-y, n) > 1.
     *
     * @param n composite (odd, n > 1)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_RhoMontgomery(const Int& n);

    /**
     * @brief Factor search based on the Goldbach conjecture
     *
     * For N = p * q (semiprime), estimate sum ≈ p + q by binary search and
     * solve the quadratic t^2 - sum*t + N = 0 to recover the factors.
     * Algorithm credit: Dr. Roger G. Doss, PhD.
     *
     * Note: generally slower than Factor_Fermat. Provided for educational purposes.
     *
     * @param n composite (n > 1; assumed semiprime)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_Goldbach(const Int& n);

    /**
     * @brief Factor search via trial division
     *
     * Try dividing in order by the primes in the SMALL_PRIMES table.
     *
     * @param n composite (n > 1)
     * @return a non-trivial factor of n, or n itself if none is found
     */
    static Int findFactor_TrialDivision(const Int& n);

    /**
     * @brief Number of primes <= n (prime-counting function π(n))
     *
     * Counts primes up to n using the sieve of Eratosthenes.
     *
     * @param n upper bound (n >= 0)
     * @return number of primes <= n
     *
     * Special cases:
     *   primePi(0) = 0, primePi(1) = 0, primePi(2) = 1
     */
    static int64_t primePi(int64_t n);

    // Extract every factor of remaining that is divisible by a small prime (one-limb batched GCD)
    // Removes small-prime factors from remaining and appends found factors to factors.
    // remaining is updated to the value after dividing out the small primes.
    // Note: assumes the prime factor 2 has already been removed via countTrailingZeros().
    static void extractSmallPrimeFactors(
        Int& remaining,
        std::vector<std::pair<Int, int>>& factors);

private:
    // Small-prime list (the 299 odd primes from 3 to 1987)
    // The prime factor 2 is handled by LSB rather than batched GCD, so it is not listed here
    static constexpr int SMALL_PRIMES[] = {
        3,5,7,11,13,17,19,23,29,
        31,37,41,43,47,53,59,61,67,71,
        73,79,83,89,97,101,103,107,109,113,
        127,131,137,139,149,151,157,163,167,173,
        179,181,191,193,197,199,211,223,227,229,
        233,239,241,251,257,263,269,271,277,281,
        283,293,307,311,313,317,331,337,347,349,
        353,359,367,373,379,383,389,397,401,409,
        419,421,431,433,439,443,449,457,461,463,
        467,479,487,491,499,503,509,521,523,541,
        547,557,563,569,571,577,587,593,599,601,
        607,613,617,619,631,641,643,647,653,659,
        661,673,677,683,691,701,709,719,727,733,
        739,743,751,757,761,769,773,787,797,809,
        811,821,823,827,829,839,853,857,859,863,
        877,881,883,887,907,911,919,929,937,941,
        947,953,967,971,977,983,991,997,1009,1013,
        1019,1021,1031,1033,1039,1049,1051,1061,1063,1069,
        1087,1091,1093,1097,1103,1109,1117,1123,1129,1151,
        1153,1163,1171,1181,1187,1193,1201,1213,1217,1223,
        1229,1231,1237,1249,1259,1277,1279,1283,1289,1291,
        1297,1301,1303,1307,1319,1321,1327,1361,1367,1373,
        1381,1399,1409,1423,1427,1429,1433,1439,1447,1451,
        1453,1459,1471,1481,1483,1487,1489,1493,1499,1511,
        1523,1531,1543,1549,1553,1559,1567,1571,1579,1583,
        1597,1601,1607,1609,1613,1619,1621,1627,1637,1657,
        1663,1667,1669,1693,1697,1699,1709,1721,1723,1733,
        1741,1747,1753,1759,1777,1783,1787,1789,1801,1811,
        1823,1831,1847,1861,1867,1871,1873,1877,1879,1889,
        1901,1907,1913,1931,1933,1949,1951,1973,1979,1987
    };
    static constexpr int SMALL_PRIMES_COUNT = sizeof(SMALL_PRIMES) / sizeof(SMALL_PRIMES[0]);

    // Table for one-limb batched GCD
    // Group consecutive odd primes whose product fits in 64 bits
    struct PrimeBatch {
        uint64_t product;       // product of the primes in the batch
        int      start_index;   // starting index within SMALL_PRIMES[]
        int      count;         // number of primes in the batch
        // preinv for mpn::divmod_1 (avoids recomputing every call)
        uint64_t product_norm;  // product << shift (MSB set)
        uint64_t product_dinv;  // invert_limb(product_norm)
        unsigned shift;         // countl_zero(product)
    };

    // Get the batch table (built lazily on first invocation)
    static const std::vector<PrimeBatch>& getBatches();

    // Divisibility check by a small prime (one-limb batched GCD)
    static bool isDivisibleBySmallPrime(const Int& n);
};

} // namespace sangi

#endif // SANGI_INT_PRIME_HPP