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

// IntFactorization.cpp
// 完全素因数分解ディスパッチャの実装
// 移植元: lib++20 整数.素因数分解.cpp

#include "math/core/mp/Int/IntFactorization.hpp"
#include "math/core/mp/Int/IntOps.hpp"
#include "math/core/mp/Int/IntSqrt.hpp"
#include "math/core/mp/ThreadPool.hpp"

#include <map>
#include <cstdint>
#include <atomic>
#include <mutex>
#include <future>
#include <cmath>

namespace calx {

// ============================================================================
// uint64 専用素因数分解 — Int オブジェクトを一切使わない高速パス
// ============================================================================

namespace {

// 最小素因数テーブル (SPF) — 奇数のみ、uint16 (合成数の SPF ≤ √TABLE_LIMIT)
// TABLE_LIMIT = 1,000,000 → √ ≈ 1000 → uint16 で十分
// メモリ: 500,000 × 2 bytes = 1 MB
static constexpr int TABLE_LIMIT = 1000000;

struct SpfTable {
    uint16_t data[(TABLE_LIMIT + 1) / 2];  // data[n/2] = SPF of odd n (0 = prime)

    SpfTable() {
        // 全エントリを 0 (素数) で初期化
        std::memset(data, 0, sizeof(data));
        data[0] = 1;  // n=1 は特殊

        int sqrt_limit = static_cast<int>(std::sqrt(static_cast<double>(TABLE_LIMIT)));
        for (int p = 3; p <= sqrt_limit; p += 2) {
            if (data[p >> 1] == 0) {  // p は素数
                // p の奇数倍に対して SPF を記録 (未設定の場合のみ → 最小素因数)
                for (int64_t m = static_cast<int64_t>(p) * p; m <= TABLE_LIMIT; m += 2 * p) {
                    if (data[m >> 1] == 0) {
                        data[m >> 1] = static_cast<uint16_t>(p);
                    }
                }
            }
        }
    }
};

const SpfTable& getSpfTable() {
    static const SpfTable table;
    return table;
}

// プロセス起動時にテーブルを構築 (初回 factorize() 呼び出しの遅延を回避)
static const auto& spf_init_ = getSpfTable();

// pollardRho64 の前方宣言 (IntPrime.cpp の無名名前空間に定義済みだがリンクできないため)
uint64_t mulmod64(uint64_t a, uint64_t b, uint64_t m) {
#ifdef _MSC_VER
    uint64_t hi;
    uint64_t lo = _umul128(a, b, &hi);
    uint64_t rem;
    _udiv128(hi, lo, m, &rem);
    return rem;
#else
    return static_cast<uint64_t>(
        (static_cast<unsigned __int128>(a) * b) % m);
#endif
}

uint64_t gcd64(uint64_t a, uint64_t b) {
    if (a == 0) return b;
    if (b == 0) return a;
#ifdef _MSC_VER
    unsigned long shift;
    _BitScanForward64(&shift, a | b);
    unsigned long sa; _BitScanForward64(&sa, a); a >>= sa;
    unsigned long sb; _BitScanForward64(&sb, b); b >>= sb;
#else
    int shift = __builtin_ctzll(a | b);
    a >>= __builtin_ctzll(a);
    b >>= __builtin_ctzll(b);
#endif
    while (true) {
        if (a > b) { uint64_t t = a; a = b; b = t; }
        b -= a;
        if (b == 0) return a << shift;
#ifdef _MSC_VER
        unsigned long ctz; _BitScanForward64(&ctz, b); b >>= ctz;
#else
        b >>= __builtin_ctzll(b);
#endif
    }
}

uint64_t pollardRho64_local(uint64_t n) {
    if (n % 2 == 0) return 2;
    if (n < 4) return n;

    for (uint64_t c = 1; c <= 20; ++c) {
        uint64_t x = 2, y = 2, d = 1;
        auto f = [&](uint64_t v) { return (mulmod64(v, v, n) + c) % n; };

        uint64_t r = 1, q = 1;
        while (d == 1) {
            x = y;
            for (uint64_t i = 0; i < r; ++i) y = f(y);

            uint64_t k = 0;
            while (k < r && d == 1) {
                uint64_t ys = y;
                uint64_t m = (r - k < 128) ? (r - k) : 128;
                for (uint64_t i = 0; i < m; ++i) {
                    y = f(y);
                    uint64_t diff = (x > y) ? (x - y) : (y - x);
                    q = mulmod64(q, diff, n);
                }
                d = gcd64(q, n);
                k += m;

                if (d == n) {
                    // バックトラック
                    y = ys;
                    d = 1;
                    while (d == 1) {
                        y = f(y);
                        uint64_t diff = (x > y) ? (x - y) : (y - x);
                        d = gcd64(diff, n);
                    }
                }
            }
            r <<= 1;
            q = 1;
            if (r > 1000000) break;
        }
        if (d != 1 && d != n) return d;
    }
    return n;  // 見つからなかった
}

// 決定的 Miller-Rabin (uint64 範囲)
// bases {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37} で 2^64 未満は決定的
bool isPrime64(uint64_t n) {
    if (n < 2) return false;
    if (n < 4) return true;
    if (n % 2 == 0) return false;
    if (n % 3 == 0) return n == 3;
    if (n % 5 == 0) return n == 5;
    if (n % 7 == 0) return n == 7;

    // n-1 = d * 2^r
    uint64_t d = n - 1;
    int r = 0;
    while ((d & 1) == 0) { d >>= 1; ++r; }

    // 冪剰余 (binary method)
    auto powmod = [](uint64_t base, uint64_t exp, uint64_t mod) -> uint64_t {
        uint64_t result = 1;
        base %= mod;
        while (exp > 0) {
            if (exp & 1) result = mulmod64(result, base, mod);
            exp >>= 1;
            base = mulmod64(base, base, mod);
        }
        return result;
    };

    // n < 3,215,031,751 なら bases {2, 3, 5, 7} で十分
    // n < 2^64 なら {2,3,5,7,11,13,17,19,23,29,31,37} で決定的
    static const uint64_t witnesses[] = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37};
    int num_witnesses = (n < 3215031751ULL) ? 4 : 12;

    for (int i = 0; i < num_witnesses; ++i) {
        uint64_t a = witnesses[i];
        if (a >= n) continue;

        uint64_t x = powmod(a, d, n);
        if (x == 1 || x == n - 1) continue;

        bool composite = true;
        for (int j = 0; j < r - 1; ++j) {
            x = mulmod64(x, x, n);
            if (x == n - 1) { composite = false; break; }
        }
        if (composite) return false;
    }
    return true;
}

// factors に素因数を追加するヘルパー
void addFactor64(std::vector<std::pair<Int, int>>& factors, uint64_t p, int exp) {
    Int pi(static_cast<int64_t>(p));
    for (auto& [fp, fe] : factors) {
        if (fp == pi) { fe += exp; return; }
    }
    factors.push_back({std::move(pi), exp});
}

// SPF テーブルで分解 (n ≤ TABLE_LIMIT, n は奇数)
void factorize_spf(uint64_t n, std::vector<std::pair<Int, int>>& factors) {
    const auto& spf = getSpfTable();
    while (n > 1) {
        if ((n & 1) == 0) {
            int exp = 0;
            while ((n & 1) == 0) { n >>= 1; ++exp; }
            addFactor64(factors, 2, exp);
            continue;
        }
        uint16_t s = spf.data[n >> 1];
        if (s == 0) {
            // n は素数
            addFactor64(factors, n, 1);
            return;
        }
        uint64_t p = s;
        int exp = 0;
        while (n % p == 0) { n /= p; ++exp; }
        addFactor64(factors, p, exp);
    }
}

// uint64 Fermat 法 — 均衡半素数を即座に発見
uint64_t fermat64(uint64_t n, int max_iter = 1000) {
    // ceil(sqrt(n)) — double 近似で十分 (64-bit 範囲で誤差 ≤ 1)
    uint64_t a = static_cast<uint64_t>(std::ceil(std::sqrt(static_cast<double>(n))));
    // double の精度不足を補正
    while (a * a < n) ++a;

    for (int i = 0; i < max_iter; ++i) {
        uint64_t a2 = a * a;
        uint64_t b2 = a2 - n;
        uint64_t b = static_cast<uint64_t>(std::sqrt(static_cast<double>(b2)));
        // double 近似の補正
        while (b * b < b2) ++b;
        if (b * b == b2) {
            uint64_t f = a - b;
            if (f > 1 && f < n) return f;
        }
        ++a;
    }
    return n;
}

// uint64 専用素因数分解 — Int オブジェクトを一切使わない高速パス
void factorize_u64(uint64_t n, std::vector<std::pair<Int, int>>& factors) {
    if (n <= 1) return;

    // 偶数の処理
    if ((n & 1) == 0) {
        int exp = 0;
        while ((n & 1) == 0) { n >>= 1; ++exp; }
        addFactor64(factors, 2, exp);
        if (n <= 1) return;
    }

    // テーブル範囲内: SPF テーブル引きで即座に分解
    if (n <= TABLE_LIMIT) {
        factorize_spf(n, factors);
        return;
    }

    // テーブル範囲外 (n > 10^6):
    // 1. SPF テーブルの篩で 1000 未満の因数は既にカバー → 1000 以上のみ試す
    //    ただし、n 自体がテーブル外から来たので、小さい因数も試す必要がある
    //    → SPF テーブルの素数リストは使えないが、3〜997 の試し割りを実行

    // 3〜997 の素数で試し割り (SPF テーブル構築時に 1000 未満を篩った範囲)
    // 素数は SPF テーブルから取得可能: data[p/2] == 0 なら素数
    {
        const auto& spf = getSpfTable();
        for (uint64_t p = 3; p < 1000 && p * p <= n; p += 2) {
            if (spf.data[p >> 1] != 0) continue;  // 合成数はスキップ
            if (n % p == 0) {
                int exp = 0;
                while (n % p == 0) { n /= p; ++exp; }
                addFactor64(factors, p, exp);
            }
        }
        if (n <= 1) return;
        // n がテーブル範囲に収まった場合
        if (n <= TABLE_LIMIT) { factorize_spf(n, factors); return; }
    }

    // 1009〜1999 の試し割り (SPF テーブルの √TABLE_LIMIT ≈ 1000 以上)
    for (uint64_t p = 1009; p < 2000 && p * p <= n; p += 2) {
        if (n % p == 0) {
            int exp = 0;
            while (n % p == 0) { n /= p; ++exp; }
            addFactor64(factors, p, exp);
        }
    }
    if (n <= 1) return;
    if (n <= TABLE_LIMIT) { factorize_spf(n, factors); return; }

    // 残りが素数なら追加
    if (n < static_cast<uint64_t>(1987) * 1987) {
        // 1987² ≈ 395 万未満 → 試し割り範囲の² 未満なので素数確定
        addFactor64(factors, n, 1);
        return;
    }
    if (isPrime64(n)) {
        addFactor64(factors, n, 1);
        return;
    }

    // 合成数: Fermat + pollardRho64 で分割し再帰
    std::vector<uint64_t> composites;
    composites.push_back(n);

    while (!composites.empty()) {
        uint64_t c = composites.back();
        composites.pop_back();

        if (c <= 1) continue;

        // テーブル範囲内なら SPF で分解
        if (c <= TABLE_LIMIT) {
            factorize_spf(c, factors);
            continue;
        }

        if (isPrime64(c)) {
            addFactor64(factors, c, 1);
            continue;
        }

        // Fermat 法 (均衡半素数を即座に発見)
        uint64_t f = fermat64(c, 1000);
        if (f > 1 && f < c) {
            composites.push_back(f);
            composites.push_back(c / f);
            continue;
        }

        // Pollard rho
        f = pollardRho64_local(c);
        if (f > 1 && f < c) {
            composites.push_back(f);
            composites.push_back(c / f);
            continue;
        }

        // 見つからない (理論上到達しない)
        addFactor64(factors, c, 1);
    }
}

} // anonymous namespace

// ============================================================================
// findFactor — 単一因数探索
// 複数アルゴリズムの段階的適用
// ============================================================================

Int IntFactorization::findFactor(const Int& n, int millerRabinK) {
    // 偶数
    if (n.isEven()) return Int(2);

    // 1以下
    if (n <= Int(1)) return n;

    // 素数チェック — 小さい n では決定的テスト (少ないラウンドで十分)
    int bits = static_cast<int>(n.bitLength());
    {
        // n < 2^64: k=2 で十分 (SPRP-2,3 は 2^64 未満で決定的)
        // n < 2^128: k=5 で十分 (実用的に十分な確率)
        int k = (bits <= 64) ? 2 : (bits <= 128) ? 5 : millerRabinK;
        if (IntPrime::isProbablePrime(n, k)) return n;
    }

    // ================================================================
    // Stage 1: Fermat 法 — 軽い初期試行 (< 1us)
    // p ≈ q の均衡半素数を即座に発見。コストが極めて低いので常に最初に試す。
    // 1000 回で |p-q| ≤ ~2000·√2 ≈ 2828 の範囲をカバー。
    // ================================================================
    {
        Int f = IntPrime::findFactor_Fermat(n, 1000);
        if (f > Int(1) && f < n) return f;
    }

    // ================================================================
    // Stage 2: Pollard rho / Brent (数us〜数ms)
    // 最も汎用的な手法。O(p^{1/2}) なので小さい因数ほど速い。
    // uint64 (≤63 bits) / 128-bit Montgomery (≤127 bits) 専用版で高速化済み。
    // ================================================================
    {
        Int f = IntPrime::findFactor_PollardRho(n);
        if (f > Int(1) && f < n) return f;
    }

    // ================================================================
    // Stage 3: 特殊手法群 — 桁数に応じて選択
    // ================================================================

    // 3a: SQUFOF (≤60 bits のみ、O(n^{1/4})、メモリ極小)
    // rho で見つからなかった小〜中規模のフォールバック。
    if (bits <= 60) {
        Int f = IntPrime::findFactor_SQUFOF(n);
        if (f > Int(1) && f < n) return f;
    }

    // 3b: Pollard p-1 — bits >= 50 のときのみ実行
    // 小さい数は rho で十分。B1 は bits に応じて調整。
    if (bits >= 50) {
        uint64_t B1 = (bits <= 80) ? 5000 : 50000;
        Int f = IntPrime::findFactor_PollardP1(n, B1);
        if (f > Int(1) && f < n) return f;
    }

    // 3c: Williams p+1 — bits >= 50 のときのみ実行
    if (bits >= 50) {
        uint64_t B1 = (bits <= 80) ? 5000 : 50000;
        Int f = IntPrime::findFactor_WilliamsP1(n, B1);
        if (f > Int(1) && f < n) return f;
    }

    // ================================================================
    // Stage 4: ECM + SIQS 並列実行
    // 64 bits 以上では ECM と SIQS を同時実行し、先に因数を見つけた方を採用。
    // ECM は小さい因数に強く、SIQS は均衡半素数に強い。
    // 小さい数では逐次実行（並列化オーバーヘッド回避）。
    // ================================================================
    if (bits >= 64) {
        std::atomic<bool> cancel{false};
        Int result_factor;
        std::mutex result_mutex;

        int curves = (bits <= 80) ? 50 : (bits <= 120) ? 100 : 200;

        auto fut_ecm = calx::threadPool().submit([&]() {
            Int f = IntPrime::findFactor_ECM(n, curves, 0, 0, &cancel);
            if (f > Int(1) && f < n) {
                std::lock_guard<std::mutex> lock(result_mutex);
                if (!cancel.load(std::memory_order_relaxed)) {
                    result_factor = std::move(f);
                    cancel.store(true, std::memory_order_release);
                }
            }
        });

        auto fut_siqs = calx::threadPool().submit([&]() {
            Int f = IntPrime::findFactor_SIQS(n, &cancel);
            if (f > Int(1) && f < n) {
                std::lock_guard<std::mutex> lock(result_mutex);
                if (!cancel.load(std::memory_order_relaxed)) {
                    result_factor = std::move(f);
                    cancel.store(true, std::memory_order_release);
                }
            }
        });

        fut_ecm.get();
        fut_siqs.get();

        if (cancel.load(std::memory_order_acquire))
            return result_factor;
    } else {
        // 小さい数: 逐次実行
        {
            int curves = (bits <= 80) ? 50 : 100;
            Int f = IntPrime::findFactor_ECM(n, curves, 0);
            if (f > Int(1) && f < n) return f;
        }
        {
            Int f = IntPrime::findFactor_SIQS(n);
            if (f > Int(1) && f < n) return f;
        }
    }

    // ================================================================
    // Stage 7: GNFS（一般数体篩法）
    // 60桁以上の大きな合成数向け。最終手段。
    // ================================================================
    {
        Int f = IntPrime::findFactor_GNFS(n);
        if (f > Int(1) && f < n) return f;
    }

    return n;  // 因数が見つからなかった
}

// ============================================================================
// factorize — 完全素因数分解
// ============================================================================

std::vector<std::pair<Int, int>> IntFactorization::factorize(const Int& n, int millerRabinK) {
    // 特殊ケース
    if (n.isNaN() || n.isInfinite()) return {};
    if (n.isZero()) return {{Int(0), 1}};

    Int remaining = n;
    if (remaining.isNegative()) remaining = -remaining;
    if (remaining == Int(1)) return {};

    std::vector<std::pair<Int, int>> factors;

    // Stage 1: 素因数 2 の除去（ビットシフトで高速）
    if (remaining.isEven()) {
        int exp2 = 0;
        while (remaining.isEven()) {
            remaining >>= 1;
            ++exp2;
        }
        factors.push_back({Int(2), exp2});
    }

    if (remaining == Int(1)) {
        std::sort(factors.begin(), factors.end());
        return factors;
    }

    // ================================================================
    // Stage 1.5: uint64 ファストパス — 1 limb なら Int を一切使わず分解
    // SPF テーブル (≤ 10^6) + 試し割り + pollardRho64 で完結。
    // 19 桁以下の全ての数がこのパスで処理される。
    // ================================================================
    if (remaining.bitLength() <= 63) {
        uint64_t n64 = remaining.toUInt64();
        factorize_u64(n64, factors);
        std::sort(factors.begin(), factors.end());
        return factors;
    }

    // Stage 2: 小素数バッチ GCD で 3〜1987 の因数を除去
    IntPrime::extractSmallPrimeFactors(remaining, factors);

    // remaining == 1 なら完了
    if (remaining == Int(1)) {
        std::sort(factors.begin(), factors.end());
        return factors;
    }

    // Stage 3: 残余が素数なら追加して完了
    if (IntPrime::isProbablePrime(remaining, millerRabinK)) {
        factors.push_back({remaining, 1});
        std::sort(factors.begin(), factors.end());
        return factors;
    }

    // Stage 4: 残余を再帰的に分解
    // 合成数をスタックで管理
    std::vector<Int> composites;
    composites.push_back(remaining);

    while (!composites.empty()) {
        Int c = composites.back();
        composites.pop_back();

        if (c == Int(1)) continue;

        if (IntPrime::isProbablePrime(c, millerRabinK)) {
            // 素数：因数リストに追加
            bool found = false;
            for (auto& [p, e] : factors) {
                if (p == c) {
                    ++e;
                    found = true;
                    break;
                }
            }
            if (!found) factors.push_back({c, 1});
            continue;
        }

        // 合成数：因数を1つ見つけて分割
        Int f = findFactor(c, millerRabinK);
        if (f == c || f <= Int(1)) {
            // 因数が見つからなかった（理論上は到達しないが安全策）
            // 残余をそのまま因数として追加
            bool found = false;
            for (auto& [p, e] : factors) {
                if (p == c) {
                    ++e;
                    found = true;
                    break;
                }
            }
            if (!found) factors.push_back({c, 1});
            continue;
        }

        Int cofactor = c / f;

        // f と cofactor をスタックに追加して再帰的に処理
        composites.push_back(f);
        composites.push_back(cofactor);
    }

    // 昇順ソート
    std::sort(factors.begin(), factors.end());
    return factors;
}

// ============================================================================
// toString — 素因数分解結果の文字列表現
// ============================================================================

std::string IntFactorization::toString(const std::vector<std::pair<Int, int>>& factors) {
    if (factors.empty()) return "1";

    std::ostringstream oss;
    bool first = true;
    for (auto& [p, e] : factors) {
        if (!first) oss << " * ";
        first = false;
        oss << p.toString();
        if (e > 1) oss << "^" << e;
    }
    return oss.str();
}

} // namespace calx