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

// IntAPRCL.cpp — APRCL (Adleman-Pomerance-Rumely-Cohen-Lenstra) 決定的素数判定
//
// アルゴリズム概要:
//   1. 円分環 Z[ζ_q]/(n) 上で Jacobi 和テストを実施
//   2. 全テスト通過 → n の素因数は n^j mod e の形に限定 (e > √n)
//   3. 最終因数探索: 軌道 {n^j mod e} 内で n の因数を試す
//
// 参考文献:
//   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 calx {
namespace {

// ── 小数の素数判定 (ヘルパー素数用) ──
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;
}

// ── q | r-1 を満たす最小の素数 r を見つける ──
int findHelperPrime(int q) {
    for (int r = q + 1; ; r += q)
        if (isSmallPrime(r)) return r;
}

// ── 原始根 mod r (r は小さい素数) ──
int primitiveRoot(int r) {
    if (r == 2) return 1;
    // φ(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) {
            // 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;
}

// ── 完全べき乗判定: 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;
        }

        // 二分探索で a = n^{1/b} を求める
        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;
            // 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: Z[ζ_q]/(n) の元 (q は素数)
//
// ζ_q は 1 の原始 q 乗根。Φ_q(ζ) = 0 より ζ^{q-1} = -(1+ζ+...+ζ^{q-2}).
// 元は q-1 個の Int 係数で表現: c[0] + c[1]ζ + ... + c[q-2]ζ^{q-2}
// ══════════════════════════════════════════════════════════════════
class UnityZp {
public:
    int q;           // 素数位数
    int dim;         // q - 1 (係数の数)
    Int mod;         // 法 n
    std::vector<Int> c;  // dim 個の係数

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

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

    // 乗法単位元 (1)
    static UnityZp one(int q, const Int& n) {
        UnityZp u(q, n);
        u.c[0] = Int(1);
        return u;
    }

    // 乗算: Z[ζ_q]/(n) 内での積
    UnityZp operator*(const UnityZp& rhs) const {
        UnityZp result(q, mod);

        // 多項式積を x^q = 1 で巻き込む → q 個の係数
        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];
            }
        }

        // mod n で正規化
        for (auto& x : buf)
            x = IntModular::mod(x, mod);

        // Φ_q 簡約: buf[q-1] を全係数から引く
        // (ζ^{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;
    }

    // べき乗: 左→右二進法
    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;
    }

    // ガロア自己同型 σ_s: ζ → ζ^s
    UnityZp sigma(int s) const {
        // 係数 c[i] を位置 (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 簡約 + 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); }
};

// ── Jacobi 和 J(χ,χ) の計算 ──
// χ は位数 q の指標 mod ヘルパー素数 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);

    // 離散対数テーブル: 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;
    }

    // 原始多項式の係数 (位置 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
        int e = (ind[a] + ind[b]) % q;
        raw[e] += 1;
    }

    // Φ_q 簡約
    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;
}

// ── パラメータ選択 ──
struct APRCLParam {
    int q;            // 素数
    int helperPrime;  // q | helperPrime - 1 を満たす素数
};

// 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;

    // 素数を小さい順に追加 (まず小さい素数 → 効率的)
    // q=2 は特殊 (Z[ζ_2] = Z, 自明) なので 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 は素数なので q ∤ e なら e*q, そうでなければ e
        Int qi(q);
        Int g = IntGCD::gcd(e, qi);
        e = e * qi / g;
    }

    return params;
}

} // anonymous namespace

// ══════════════════════════════════════════════════════════════════
// isProvablePrime — APRCL 決定的素数判定
// ══════════════════════════════════════════════════════════════════
bool IntPrime::isProvablePrime(const Int& n) {
    // ── 特殊ケース ──
    if (n.isSpecialState()) return false;
    if (n <= Int(1)) return false;
    if (n == Int(2) || n == Int(3)) return true;
    if (n.isEven()) return false;

    // ── 小さい数: Miller-Rabin で決定的 ──
    // n < 3.3×10^24 (≈ 82 bits) では特定の底で決定的
    if (n.bitLength() <= 82) {
        return isMillerRabinPrime(n, 40);
    }

    // ── 完全べき乗チェック ──
    if (isPerfectPower(n)) return false;

    // ── パラメータ選択 ──
    auto params = selectParams(n);

    // 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;
    }

    // ── 前処理: gcd(n, e) チェック ──
    Int ge = IntGCD::gcd(n, e);
    if (ge != Int(1)) {
        // n は小さい素数で割り切れる
        return n == ge;
    }

    // ── Jacobi 和テスト ──
    for (auto& param : params) {
        // gcd(n, helperPrime) チェック
        Int gr(param.helperPrime);
        Int nmod_r = IntModular::mod(n, gr);
        if (nmod_r.isZero()) {
            return n == gr;
        }

        // Jacobi 和 J(χ,χ) を計算
        UnityZp J = jacobiSum(param.q, param.helperPrime, n);

        // J^n mod n を計算
        UnityZp Jn = J.pow(n);

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

        // テスト: J^n ≡ σ_s(J) (mod n)
        if (Jn != sigmaJ) {
            return false;  // 合成数確定
        }
    }

    // ── 最終因数探索 ──
    // Jacobi 和テスト通過 → n の全素因数 p は p ≡ n^j (mod e) を満たす
    // e > √n なので p < e → p = n^j mod e
    // 軌道 {n^j mod e : j ≥ 0} 内で n の因数を探す
    Int nmod = IntModular::mod(n, e);
    Int cur(1);  // n^0 mod e
    Int one_val(1);

    // 軌道の長さ = ord_e(n), 最大 φ(e)
    // 安全上限を設定 (超大数では ECPP を使うべき)
    constexpr size_t MAX_ORBIT = 10000000;

    for (size_t j = 0; j < MAX_ORBIT; ++j) {
        if (j > 0 && cur == one_val) break;  // 一周した

        if (cur > one_val && cur < n) {
            // cur が n の因数か確認
            Int g = IntGCD::gcd(cur, n);
            if (g > one_val && g < n)
                return false;  // 非自明因数発見 → 合成数
        }

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

    return true;  // 因数なし → 素数確定
}

} // namespace calx