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

// IntNumberTheory.hpp
// 数論的関数（約数列挙、メビウス関数、オイラーのファイ関数、Jacobi記号）
// 移植元: lib++20 整数.素数.cpp

#pragma once

#include "IntBase.hpp"
#include "IntFactorTable.hpp"
#include "IntGCD.hpp"

#include <vector>
#include <algorithm>

namespace calx {

/**
 * @brief 数論的関数
 *
 * 素因数分解に基づく数論的関数を提供します。
 * IntFactorTable を利用して素因数分解を行い、
 * 約数列挙・メビウス関数・オイラーのファイ関数を計算します。
 */
class IntNumberTheory {
public:
    /**
     * @brief n の全ての正の約数を昇順で返す
     *
     * @param n 正整数 (n <= 0 の場合は空ベクトル)
     * @param table 素因数テーブル
     * @return 約数のベクトル（昇順ソート済み）
     *
     * アルゴリズム:
     *   素因数分解 n = p1^e1 * p2^e2 * ... * pk^ek から
     *   各素因数の冪 0..ei の全組み合わせを再帰的に生成する。
     *
     * 例:
     *   divisors(12, table) → {1, 2, 3, 4, 6, 12}
     *   divisors(1, table)  → {1}
     *   divisors(0, table)  → {}
     */
    static std::vector<Int> divisors(const Int& n, const IntFactorTable& table) {
        if (n <= Int(0)) return {};

        auto factors = table.factorize(n);
        if (factors.empty()) return { Int(1) };  // n == 1

        // 約数の個数を計算
        size_t count = 1;
        for (auto& [p, e] : factors)
            count *= static_cast<size_t>(e + 1);

        // 約数を生成: 各素因数の冪の全組み合わせ
        std::vector<Int> result;
        result.reserve(count);
        result.push_back(Int(1));

        for (auto& [p, e] : factors) {
            size_t prev_size = result.size();
            Int pk(1);
            for (int k = 1; k <= e; ++k) {
                pk *= p;
                for (size_t j = 0; j < prev_size; ++j)
                    result.push_back(result[j] * pk);
            }
        }

        std::sort(result.begin(), result.end());
        return result;
    }

    /**
     * @brief メビウス関数 μ(n)
     *
     * @param n 正整数 (n <= 0 → 0)
     * @param table 素因数テーブル
     * @return μ(n):
     *   - μ(1) = 1
     *   - μ(n) = 0  if n が平方因子を持つ
     *   - μ(n) = (-1)^k  if n = p1*p2*...*pk (異なる素数の積)
     *
     * 参考: 芹沢正三「素数入門」p206
     */
    static int moebiusMu(const Int& n, const IntFactorTable& table) {
        if (n <= Int(0)) return 0;
        if (n == Int(1)) return 1;

        auto factors = table.factorize(n);

        // 平方因子を持つか確認
        for (auto& [p, e] : factors) {
            if (e >= 2) return 0;
        }

        // 異なる素因数の数が偶数個なら +1, 奇数個なら -1
        return (factors.size() & 1) ? -1 : 1;
    }

    /**
     * @brief オイラーのファイ関数 φ(n)
     *
     * n と互いに素な n 以下の正整数の個数を返す。
     *
     * @param n 正整数 (n <= 0 → 0)
     * @param table 素因数テーブル
     * @return φ(n) = n * Π_{p|n} (1 - 1/p)
     *
     * アルゴリズム:
     *   φ(n) = n * Π (1 - 1/p_i) = n - Σ n/p_i + ...
     *   整数演算で実装: result = n, 各素因数 p に対して result -= result / p
     *
     * 参考: 芹沢正三「素数入門」p200
     */
    static Int eulerPhi(const Int& n, const IntFactorTable& table) {
        if (n <= Int(0)) return Int(0);
        if (n == Int(1)) return Int(1);

        auto factors = table.factorize(n);
        Int result = n;
        for (auto& [p, e] : factors)
            result -= result / p;
        return result;
    }

    /**
     * @brief 約数の個数 σ_0(n)
     *
     * @param n 正整数
     * @param table 素因数テーブル
     * @return n の約数の個数
     */
    static Int divisorCount(const Int& n, const IntFactorTable& table) {
        if (n <= Int(0)) return Int(0);
        if (n == Int(1)) return Int(1);

        auto factors = table.factorize(n);
        Int count(1);
        for (auto& [p, e] : factors)
            count *= Int(e + 1);
        return count;
    }

    /**
     * @brief 約数の総和 σ_1(n)
     *
     * @param n 正整数
     * @param table 素因数テーブル
     * @return n の全約数の和
     *
     * アルゴリズム:
     *   σ(n) = Π_{p^e || n} (p^{e+1} - 1) / (p - 1)
     */
    static Int divisorSum(const Int& n, const IntFactorTable& table) {
        if (n <= Int(0)) return Int(0);
        if (n == Int(1)) return Int(1);

        auto factors = table.factorize(n);
        Int result(1);
        for (auto& [p, e] : factors) {
            // (p^{e+1} - 1) / (p - 1) = 1 + p + p^2 + ... + p^e
            Int sum(1), pk(1);
            for (int k = 1; k <= e; ++k) {
                pk *= p;
                sum += pk;
            }
            result *= sum;
        }
        return result;
    }
    // =====================================================================
    // Jacobi 記号 / Legendre 記号 / Kronecker 記号
    // =====================================================================

    /**
     * @brief Jacobi 記号 (a/n)
     *
     * n が奇数正整数のとき、Jacobi 記号 (a/n) を計算します。
     * n が素数のとき Legendre 記号に一致します。
     *
     * アルゴリズム: 二次相互法則に基づく反復法
     *   1. a = a mod n に簡約
     *   2. a が偶数なら 2 の因子を取り出し (2/n) を計算
     *   3. 二次相互法則で (a/n) = ±(n/a) と反転
     *   4. a と n を入れ替えて繰り返し
     *
     * @param a 任意の整数
     * @param n 奇数正整数 (n > 0, n は奇数)
     * @return -1, 0, 1 のいずれか
     *
     * 特殊ケース:
     *   jacobi(0, n) = 0 (n > 1), jacobi(0, 1) = 1
     *   jacobi(1, n) = 1
     *   n が偶数または n <= 0 → 0 (未定義)
     */
    static int jacobiSymbol(const Int& a, const Int& n) {
        // n は奇数正整数でなければならない
        if (n <= Int(0) || n.isEven()) return 0;
        if (n == Int(1)) return 1;

        // a を n で剰余 (0 <= a_mod < n)
        Int a_mod = a % n;
        if (a_mod.isNegative()) a_mod += n;

        if (a_mod.isZero()) return 0;
        if (a_mod == Int(1)) return 1;

        Int aa = a_mod;
        Int nn = n;
        int result = 1;

        while (true) {
            if (aa.isZero()) return 0;
            if (aa == Int(1)) return result;

            // 2 の因子を取り出す
            // (2/n) = (-1)^{(n²-1)/8}
            // n mod 8 が 3 or 5 のとき (2/n) = -1
            while (aa.isEven()) {
                aa >>= 1;
                // n mod 8
                int n8 = static_cast<int>(nn.toUInt64() & 7);
                if (n8 == 3 || n8 == 5) result = -result;
            }

            if (aa == Int(1)) return result;

            // 二次相互法則: (a/n) = (-1)^{(a-1)(n-1)/4} * (n/a)
            // a ≡ 3 (mod 4) かつ n ≡ 3 (mod 4) のとき符号反転
            int a4 = static_cast<int>(aa.toUInt64() & 3);
            int n4 = static_cast<int>(nn.toUInt64() & 3);
            if (a4 == 3 && n4 == 3) result = -result;

            // swap(aa, nn), aa = nn mod aa
            Int temp = aa;
            aa = nn % temp;
            nn = temp;
        }
    }

    /**
     * @brief Legendre 記号 (a/p)
     *
     * p が奇素数のとき、a が p を法とする二次剰余かどうかを判定します。
     * Jacobi 記号のラッパーです。
     *
     * @param a 任意の整数
     * @param p 奇素数
     * @return -1: 非剰余, 0: p|a, 1: 剰余
     */
    static int legendreSymbol(const Int& a, const Int& p) {
        return jacobiSymbol(a, p);
    }

    /**
     * @brief Kronecker 記号 (a/n) の拡張
     *
     * n = 0, 1, 2, 負数に対しても定義される拡張 Jacobi 記号。
     *
     * @param a 任意の整数
     * @param n 任意の整数
     * @return Kronecker 記号の値
     */
    static int kroneckerSymbol(const Int& a, const Int& n) {
        if (n.isZero()) {
            // (a/0) = 1 if |a|=1, 0 otherwise
            Int abs_a = a.isNegative() ? -a : a;
            return (abs_a == Int(1)) ? 1 : 0;
        }

        Int nn = n;
        int result = 1;

        // 負の n: (a/-1) = -1 if a < 0, 1 otherwise
        if (nn.isNegative()) {
            nn = -nn;
            if (a.isNegative()) result = -result;
        }

        // n の因子 2 を取り出す: (a/2)
        int v = 0;
        while (nn.isEven()) {
            nn >>= 1;
            ++v;
        }

        if (v > 0) {
            // (a/2^v)
            // (a/2) = 0 if a is even, (-1)^{(a²-1)/8} if a is odd
            Int abs_a = a.isNegative() ? -a : a;
            if (abs_a.isEven()) return 0;
            if (v & 1) {
                int a8 = static_cast<int>(abs_a.toUInt64() & 7);
                if (a8 == 3 || a8 == 5) result = -result;
            }
        }

        // nn が 1 なら完了
        if (nn == Int(1)) return result;

        // 残りは奇数部分に対する Jacobi 記号
        return result * jacobiSymbol(a, nn);
    }
};

} // namespace calx