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

// DiscreteLog.hpp
//
// 離散対数 (Discrete Logarithm)
//
// g^x ≡ h (mod p) を解く。
//
// 主な API:
//   discreteLog(g, h, p)                 — Baby-step Giant-step (BSGS)
//   discreteLogPohligHellman(g, h, p)    — Pohlig-Hellman
//   discreteLogBruteForce(g, h, p)       — 総当たり (小さい p 用)
//   primitiveRoot(p)                     — 原始根
//   orderMod(g, p)                       — 位数 ord_p(g)

#ifndef CALX_CORE_DISCRETE_LOG_HPP
#define CALX_CORE_DISCRETE_LOG_HPP

#include <cstdint>
#include <cmath>
#include <stdexcept>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <optional>
#include <numeric>

namespace calx {

namespace detail_dlog {

/// overflow-safe modular multiply using double-width intermediate
inline int64_t mulmod(int64_t a, int64_t b, int64_t m) {
    // Use unsigned to avoid UB; works for m < 2^62
    auto ua = static_cast<uint64_t>(((a % m) + m) % m);
    auto ub = static_cast<uint64_t>(((b % m) + m) % m);
    auto um = static_cast<uint64_t>(m);
    // Russian peasant multiplication to avoid overflow
    uint64_t result = 0;
    ua %= um;
    while (ub > 0) {
        if (ub & 1) { result += ua; if (result >= um) result -= um; }
        ua <<= 1; if (ua >= um) ua -= um;
        ub >>= 1;
    }
    return static_cast<int64_t>(result);
}

inline int64_t modpow(int64_t base, int64_t exp, int64_t mod) {
    int64_t result = 1;
    base = ((base % mod) + mod) % mod;
    while (exp > 0) {
        if (exp & 1) result = mulmod(result, base, mod);
        base = mulmod(base, base, mod);
        exp >>= 1;
    }
    return result;
}

inline int64_t gcd(int64_t a, int64_t b) {
    a = std::abs(a); b = std::abs(b);
    while (b) { a %= b; std::swap(a, b); }
    return a;
}

/// 拡張ユークリッド
inline int64_t modInverse(int64_t a, int64_t m) {
    a = ((a % m) + m) % m;
    int64_t old_r = a, r = m;
    int64_t old_s = 1, s = 0;
    while (r != 0) {
        int64_t q = old_r / r;
        int64_t tmp = r; r = old_r - q * r; old_r = tmp;
        tmp = s; s = old_s - q * s; old_s = tmp;
    }
    if (old_r != 1) return -1;  // 逆元なし
    return ((old_s % m) + m) % m;
}

/// p-1 の素因数分解
inline std::vector<std::pair<int64_t, int>> factorize(int64_t n) {
    std::vector<std::pair<int64_t, int>> factors;
    for (int64_t d = 2; d * d <= n; ++d) {
        if (n % d == 0) {
            int exp = 0;
            while (n % d == 0) { n /= d; ++exp; }
            factors.push_back({d, exp});
        }
    }
    if (n > 1) factors.push_back({n, 1});
    return factors;
}

} // namespace detail_dlog

/**
 * @brief g の mod p での位数
 *
 * ord_p(g) = min {k > 0 : g^k ≡ 1 (mod p)}
 */
[[nodiscard]]
inline int64_t orderMod(int64_t g, int64_t p) {
    if (p <= 1) throw std::invalid_argument("orderMod: p must be > 1");
    g = ((g % p) + p) % p;
    if (g == 0 || detail_dlog::gcd(g, p) != 1)
        throw std::invalid_argument("orderMod: g must be coprime to p");

    // p-1 の約数を効率的にチェック
    int64_t n = p - 1;
    auto factors = detail_dlog::factorize(n);

    int64_t order = n;
    for (auto& [q, _] : factors) {
        while (order % q == 0 && detail_dlog::modpow(g, order / q, p) == 1)
            order /= q;
    }
    return order;
}

/**
 * @brief 原始根 (primitive root mod p)
 *
 * g^{p-1} ≡ 1 かつ g^{(p-1)/q} ≢ 1 for all prime q | (p-1)
 */
[[nodiscard]]
inline int64_t primitiveRoot(int64_t p) {
    if (p <= 1) throw std::invalid_argument("primitiveRoot: p must be > 1");
    if (p == 2) return 1;

    int64_t phi = p - 1;
    auto factors = detail_dlog::factorize(phi);

    for (int64_t g = 2; g < p; ++g) {
        bool isPrimRoot = true;
        for (auto& [q, _] : factors) {
            if (detail_dlog::modpow(g, phi / q, p) == 1) {
                isPrimRoot = false;
                break;
            }
        }
        if (isPrimRoot) return g;
    }
    throw std::runtime_error("primitiveRoot: no primitive root found");
}

/**
 * @brief 離散対数 (Baby-step Giant-step)
 *
 * g^x ≡ h (mod p) を解く。解がなければ nullopt。
 * O(√n) 時間・空間。n = ord(g) or p-1。
 */
[[nodiscard]]
inline std::optional<int64_t> discreteLog(int64_t g, int64_t h, int64_t p) {
    if (p <= 1) throw std::invalid_argument("discreteLog: p must be > 1");
    g = ((g % p) + p) % p;
    h = ((h % p) + p) % p;

    if (h == 1) return 0;

    int64_t n = static_cast<int64_t>(std::ceil(std::sqrt(static_cast<double>(p)))) + 1;

    // Baby step: {g^j mod p → j} for j = 0, ..., n-1
    std::unordered_map<int64_t, int64_t> table;
    int64_t gj = 1;
    for (int64_t j = 0; j < n; ++j) {
        table[gj] = j;
        gj = detail_dlog::mulmod(gj, g, p);
    }

    // Giant step: g^{-n} mod p
    int64_t gn = detail_dlog::modpow(g, n, p);
    int64_t gnInv = detail_dlog::modInverse(gn, p);
    if (gnInv < 0) return std::nullopt;

    // h · (g^{-n})^i for i = 0, ..., n-1
    int64_t gamma = h;
    for (int64_t i = 0; i <= n; ++i) {
        auto it = table.find(gamma);
        if (it != table.end()) {
            int64_t x = i * n + it->second;
            return x;
        }
        gamma = detail_dlog::mulmod(gamma, gnInv, p);
    }

    return std::nullopt;
}

/**
 * @brief 離散対数 (Pohlig-Hellman)
 *
 * p-1 の素因数分解を利用して CRT で合成。
 * p-1 が小さい素因数のみで構成される場合に高速。
 */
[[nodiscard]]
inline std::optional<int64_t> discreteLogPohligHellman(int64_t g, int64_t h, int64_t p) {
    if (p <= 1) throw std::invalid_argument("discreteLogPH: p must be > 1");
    g = ((g % p) + p) % p;
    h = ((h % p) + p) % p;

    int64_t n = p - 1;  // 位数 (p が素数の場合)
    auto factors = detail_dlog::factorize(n);

    // 各素因数べきについて部分離散対数を計算
    std::vector<int64_t> residues;
    std::vector<int64_t> moduli;

    for (auto& [q, e] : factors) {
        int64_t qe = 1;
        for (int i = 0; i < e; ++i) qe *= q;

        // g_q = g^{n/q^e}, h_q = h^{n/q^e}
        // x ≡ x_q (mod q^e) を求める
        // Silver-Pohlig-Hellman: 桁ごとに解く

        int64_t gamma = detail_dlog::modpow(g, n / q, p);  // 位数 q の元
        int64_t x_q = 0;
        int64_t qPow = 1;  // q^k
        int64_t hCurrent = h;

        for (int k = 0; k < e; ++k) {
            // h_k = hCurrent^{n/q^{k+1}}
            int64_t hk = detail_dlog::modpow(hCurrent, n / (qPow * q), p);
            // BSGS で gamma^d_k = h_k を解く
            auto dk = discreteLog(gamma, hk, p);
            if (!dk.has_value()) return std::nullopt;
            int64_t d = dk.value() % q;

            x_q += d * qPow;
            // hCurrent = h · g^{-x_q}
            hCurrent = detail_dlog::mulmod(h,
                detail_dlog::modpow(detail_dlog::modInverse(g, p), x_q, p), p);
            qPow *= q;
        }

        residues.push_back(x_q);
        moduli.push_back(qe);
    }

    // CRT で合成
    int64_t result = 0;
    int64_t M = 1;
    for (auto m : moduli) M *= m;

    for (std::size_t i = 0; i < moduli.size(); ++i) {
        int64_t Mi = M / moduli[i];
        int64_t yi = detail_dlog::modInverse(Mi % moduli[i], moduli[i]);
        if (yi < 0) return std::nullopt;
        // overflow 回避: mulmod を使用
        int64_t term = detail_dlog::mulmod(
            detail_dlog::mulmod(residues[i], Mi % M, M), yi, M);
        result = (result + term) % M;
    }

    return result;
}

/**
 * @brief 離散対数 (総当たり、小さい p 用)
 */
[[nodiscard]]
inline std::optional<int64_t> discreteLogBruteForce(int64_t g, int64_t h, int64_t p) {
    g = ((g % p) + p) % p;
    h = ((h % p) + p) % p;
    int64_t x = 1;
    for (int64_t i = 0; i < p; ++i) {
        if (x == h) return i;
        x = detail_dlog::mulmod(x, g, p);
    }
    return std::nullopt;
}

} // namespace calx

#endif // CALX_CORE_DISCRETE_LOG_HPP
