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

// IntAPRCL.cpp — APRCL (Adleman-Pomerance-Rumely-Cohen-Lenstra) deterministic primality test
//
// Algorithm overview:
//   1. Run the Jacobi sum test over the cyclotomic ring Z[ζ_q]/(n)
//   2. All tests pass → prime factors of n are restricted to the form n^j mod e (e > √n)
//   3. Final factor search: try factors of n within the orbit {n^j mod e}
//
// References:
//   H. Cohen, "A Course in Computational Algebraic Number Theory", §9.1
//   R. Crandall & C. Pomerance, "Prime Numbers: A Computational Perspective", §4.4

#include <math/core/mp/Int/IntPrime.hpp>
#include <math/core/mp/Int/IntModular.hpp>
#include <math/core/mp/Int/IntOps.hpp>
#include <math/core/mp/Int/IntGCD.hpp>
#include <algorithm>
#include <cmath>

namespace sangi {
namespace {

// ── Small-number primality test (for helper primes) ──
bool isSmallPrime(int n) {
    if (n < 2) return false;
    if (n < 4) return true;
    if (n % 2 == 0 || n % 3 == 0) return false;
    for (int i = 5; i * i <= n; i += 6)
        if (n % i == 0 || n % (i + 2) == 0) return false;
    return true;
}

// ── Find the smallest prime r satisfying q | r-1 ──
int findHelperPrime(int q) {
    for (int r = q + 1; ; r += q)
        if (isSmallPrime(r)) return r;
}

// ── Primitive root mod r (r is a small prime) ──
int primitiveRoot(int r) {
    if (r == 2) return 1;
    // Prime factorization of φ(r) = r-1
    int phi = r - 1;
    int factors[16];
    int nf = 0;
    {
        int tmp = phi;
        for (int p = 2; p * p <= tmp; ++p) {
            if (tmp % p == 0) {
                factors[nf++] = p;
                while (tmp % p == 0) tmp /= p;
            }
        }
        if (tmp > 1) factors[nf++] = tmp;
    }
    for (int g = 2; g < r; ++g) {
        bool ok = true;
        for (int fi = 0; fi < nf; ++fi) {
            // Verify g^{phi/factors[fi]} mod r ≠ 1
            long long pw = 1;
            int exp = phi / factors[fi];
            long long base = g;
            for (int e = exp; e > 0; e >>= 1) {
                if (e & 1) pw = pw * base % r;
                base = base * base % r;
            }
            if (pw == 1) { ok = false; break; }
        }
        if (ok) return g;
    }
    return -1;
}

// ── Perfect-power test: is n = a^b (b ≥ 2)? ──
bool isPerfectPower(const Int& n) {
    if (n <= Int(3)) return false;
    int maxExp = static_cast<int>(n.bitLength());

    for (int b = 2; b <= maxExp; ++b) {
        if (!isSmallPrime(b)) continue;

        if (b == 2) {
            Int a = IntOps::sqrt(n);
            if (a * a == n) return true;
            continue;
        }

        // Find a = n^{1/b} by binary search
        size_t aBits = (n.bitLength() + b - 1) / b;
        if (aBits == 0) break;

        Int lo(2);
        Int hi = Int(1) << static_cast<int>(aBits + 1);
        if (hi > n) hi = n;

        while (lo <= hi) {
            Int mid = (lo + hi) >> 1;
            // Compute mid^b
            Int pw(1);
            bool overflow = false;
            for (int i = 0; i < b; ++i) {
                pw *= mid;
                if (pw > n) { overflow = true; break; }
            }
            if (!overflow && pw == n) return true;
            if (overflow || pw > n)
                hi = mid - 1;
            else
                lo = mid + 1;
        }
    }
    return false;
}

// ══════════════════════════════════════════════════════════════════
// UnityZp: element of Z[ζ_q]/(n) (q is prime)
//
// ζ_q is a primitive q-th root of unity. From Φ_q(ζ) = 0, ζ^{q-1} = -(1+ζ+...+ζ^{q-2}).
// An element is represented by q-1 Int coefficients: c[0] + c[1]ζ + ... + c[q-2]ζ^{q-2}
// ══════════════════════════════════════════════════════════════════
class UnityZp {
public:
    int q;           // prime order
    int dim;         // q - 1 (number of coefficients)
    Int mod;         // modulus n
    std::vector<Int> c;  // dim coefficients

    UnityZp() : q(0), dim(0) {}

    UnityZp(int q_, const Int& n_)
        : q(q_), dim(q_ - 1), mod(n_), c(q_ - 1) {}

    // Multiplicative identity (1)
    static UnityZp one(int q, const Int& n) {
        UnityZp u(q, n);
        u.c[0] = Int(1);
        return u;
    }

    // Multiplication: product within Z[ζ_q]/(n)
    UnityZp operator*(const UnityZp& rhs) const {
        UnityZp result(q, mod);

        // Wrap the polynomial product with x^q = 1 → q coefficients
        std::vector<Int> buf(q);
        for (int i = 0; i < dim; ++i) {
            if (c[i].isZero()) continue;
            for (int j = 0; j < dim; ++j) {
                if (rhs.c[j].isZero()) continue;
                int pos = (i + j) % q;
                buf[pos] += c[i] * rhs.c[j];
            }
        }

        // Normalize mod n
        for (auto& x : buf)
            x = IntModular::mod(x, mod);

        // Φ_q reduction: subtract buf[q-1] from all coefficients
        // (ζ^{q-1} = -(1 + ζ + ... + ζ^{q-2}))
        for (int j = 0; j < dim; ++j)
            result.c[j] = IntModular::mod(buf[j] - buf[q - 1], mod);

        return result;
    }

    // Exponentiation: left-to-right binary method
    UnityZp pow(const Int& exp) const {
        if (exp.isZero()) return one(q, mod);

        int bits = static_cast<int>(exp.bitLength());
        UnityZp result = one(q, mod);

        for (int i = bits - 1; i >= 0; --i) {
            result = result * result;
            if (exp.getBit(i))
                result = result * (*this);
        }
        return result;
    }

    // Galois automorphism σ_s: ζ → ζ^s
    UnityZp sigma(int s) const {
        // Map coefficient c[i] to position (i*s) mod q
        std::vector<Int> buf(q);
        for (int i = 0; i < dim; ++i) {
            if (c[i].isZero()) continue;
            int pos = static_cast<int>((static_cast<long long>(i) * s) % q);
            if (pos < 0) pos += q;
            buf[pos] += c[i];
        }

        // Φ_q reduction + mod n
        UnityZp result(q, mod);
        for (int j = 0; j < dim; ++j)
            result.c[j] = IntModular::mod(buf[j] - buf[q - 1], mod);
        return result;
    }

    bool operator==(const UnityZp& rhs) const {
        if (q != rhs.q) return false;
        for (int i = 0; i < dim; ++i)
            if (c[i] != rhs.c[i]) return false;
        return true;
    }

    bool operator!=(const UnityZp& rhs) const { return !(*this == rhs); }
};

// ── Computation of the Jacobi sum J(χ,χ) ──
// χ is a character of order q mod the helper prime r
// J(χ,χ) = Σ_{a=2}^{r-1} ζ_q^{ind(a) + ind(1-a)}
UnityZp jacobiSum(int q, int r, const Int& n) {
    int g = primitiveRoot(r);

    // Discrete logarithm table: ind[a] = log_g(a) mod r
    std::vector<int> ind(r, -1);
    long long ga = 1;
    for (int i = 0; i < r - 1; ++i) {
        ind[static_cast<int>(ga)] = i;
        ga = ga * g % r;
    }

    // Coefficients of the raw polynomial (accumulated at positions 0..q-1)
    std::vector<Int> raw(q);
    for (int a = 2; a < r; ++a) {
        int b = ((1 - a) % r + r) % r;
        if (b == 0) continue;  // χ(0) = 0 (excluded)
        int e = (ind[a] + ind[b]) % q;
        raw[e] += 1;
    }

    // Φ_q reduction
    UnityZp J(q, n);
    for (int j = 0; j < q - 1; ++j)
        J.c[j] = IntModular::mod(raw[j] - raw[q - 1], n);

    return J;
}

// ── Parameter selection ──
struct APRCLParam {
    int q;            // prime
    int helperPrime;  // prime satisfying q | helperPrime - 1
};

// Add primes until e = lcm(q_i) > √n
std::vector<APRCLParam> selectParams(const Int& n) {
    std::vector<APRCLParam> params;
    Int e(1);
    Int target = IntOps::sqrt(n) + 1;

    // Add primes in increasing order (small primes first → efficient)
    // q=2 is special (Z[ζ_2] = Z, trivial), so start from q=3
    for (int q = 3; q < 10000; q += 2) {
        if (e > target) break;
        if (!isSmallPrime(q)) continue;

        int r = findHelperPrime(q);
        params.push_back({q, r});

        // e = lcm(e, q) — q is prime, so e*q if q ∤ e, otherwise e
        Int qi(q);
        Int g = IntGCD::gcd(e, qi);
        e = e * qi / g;
    }

    return params;
}

} // anonymous namespace

// ══════════════════════════════════════════════════════════════════
// isProvablePrime — APRCL deterministic primality test
// ══════════════════════════════════════════════════════════════════
bool IntPrime::isProvablePrime(const Int& n) {
    // ── Special cases ──
    if (n.isSpecialState()) return false;
    if (n <= 1) return false;
    if (n == 2 || n == 3) return true;
    if (n.isEven()) return false;

    // ── Small numbers: deterministic via Miller-Rabin ──
    // For n < 3.3×10^24 (≈ 82 bits), deterministic with specific bases
    if (n.bitLength() <= 82) {
        return isMillerRabinPrime(n, 40);
    }

    // ── Perfect-power check ──
    if (isPerfectPower(n)) return false;

    // ── Parameter selection ──
    auto params = selectParams(n);

    // Compute e = lcm(q_i)
    Int e(1);
    for (auto& p : params) {
        Int qi(p.q);
        Int g = IntGCD::gcd(e, qi);
        e = e * qi / g;
    }

    // ── Preprocessing: gcd(n, e) check ──
    Int ge = IntGCD::gcd(n, e);
    if (!ge.isOne()) {
        // n is divisible by a small prime
        return n == ge;
    }

    // ── Jacobi sum test ──
    for (auto& param : params) {
        // gcd(n, helperPrime) check
        Int gr(param.helperPrime);
        Int nmod_r = IntModular::mod(n, gr);
        if (nmod_r.isZero()) {
            return n == gr;
        }

        // Compute the Jacobi sum J(χ,χ)
        UnityZp J = jacobiSum(param.q, param.helperPrime, n);

        // Compute J^n mod n
        UnityZp Jn = J.pow(n);

        // Compute σ_{n mod q}(J)
        Int nmodq = IntModular::mod(n, Int(param.q));
        int s = nmodq.toInt();
        UnityZp sigmaJ = J.sigma(s);

        // Test: J^n ≡ σ_s(J) (mod n)
        if (Jn != sigmaJ) {
            return false;  // confirmed composite
        }
    }

    // ── Final factor search ──
    // Jacobi sum test passed → every prime factor p of n satisfies p ≡ n^j (mod e)
    // Since e > √n, p < e → p = n^j mod e
    // Search for factors of n within the orbit {n^j mod e : j ≥ 0}
    Int nmod = IntModular::mod(n, e);
    Int cur(1);  // n^0 mod e
    Int one_val(1);

    // Orbit length = ord_e(n), at most φ(e)
    // Set a safety bound (ECPP should be used for very large numbers)
    constexpr size_t MAX_ORBIT = 10000000;

    for (size_t j = 0; j < MAX_ORBIT; ++j) {
        if (j > 0 && cur == one_val) break;  // completed a full cycle

        if (cur > one_val && cur < n) {
            // Check whether cur is a factor of n
            Int g = IntGCD::gcd(cur, n);
            if (g > one_val && g < n)
                return false;  // nontrivial factor found → composite
        }

        cur = IntModular::mod(cur * nmod, e);
    }

    return true;  // no factor → confirmed prime
}

} // namespace sangi