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

// IntNumberTheory.hpp
// Number-theoretic functions (divisor enumeration, Möbius function, Euler's totient, Jacobi symbol)

#pragma once

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

#include <vector>
#include <algorithm>
#include <bit>

namespace sangi {

/**
 * @brief Number-theoretic functions
 *
 * Provides number-theoretic functions based on prime factorization.
 * Uses IntFactorTable to perform prime factorization, then computes
 * divisor enumeration, the Möbius function, and Euler's totient.
 */
class IntNumberTheory {
public:
    /**
     * @brief Returns all positive divisors of n in ascending order
     *
     * @param n Positive integer (returns an empty vector when n <= 0)
     * @param table Prime-factor table
     * @return Vector of divisors (sorted ascending)
     *
     * Algorithm:
     *   From the prime factorization n = p1^e1 * p2^e2 * ... * pk^ek,
     *   recursively generate all combinations of prime powers 0..ei.
     *
     * Example:
     *   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 <= 0) return {};

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

        // Compute the number of divisors
        size_t count = 1;
        for (auto& [p, e] : factors)
            count *= static_cast<size_t>(e + 1);

        // Generate divisors: every combination of prime-power exponents
        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 Möbius function mu(n)
     *
     * @param n Positive integer (n <= 0 -> 0)
     * @param table Prime-factor table
     * @return μ(n):
     *   - μ(1) = 1
     *   - mu(n) = 0  if n has a squared prime factor
     *   - mu(n) = (-1)^k  if n = p1*p2*...*pk (product of distinct primes)
     *
     * Reference: Serizawa Masami, "Introduction to Primes", p.206
     */
    static int moebiusMu(const Int& n, const IntFactorTable& table) {
        if (n <= 0) return 0;
        if (n.isOne()) return 1;

        auto factors = table.factorize(n);

        // Check whether n has a squared prime factor
        for (auto& [p, e] : factors) {
            if (e >= 2) return 0;
        }

        // +1 if the number of distinct prime factors is even, -1 if odd
        return (factors.size() & 1) ? -1 : 1;
    }

    /**
     * @brief Euler's totient phi(n)
     *
     * Returns the count of positive integers <= n that are coprime to n.
     *
     * @param n Positive integer (n <= 0 -> 0)
     * @param table Prime-factor table
     * @return φ(n) = n * Π_{p|n} (1 - 1/p)
     *
     * Algorithm:
     *   φ(n) = n * Π (1 - 1/p_i) = n - Σ n/p_i + ...
     *   Implemented in integer arithmetic: result = n; for each prime factor p, result -= result / p
     *
     * Reference: Serizawa Masami, "Introduction to Primes", p.200
     */
    static Int eulerPhi(const Int& n, const IntFactorTable& table) {
        if (n <= 0) return Int(0);
        if (n.isOne()) return Int(1);

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

    /**
     * @brief Number of divisors sigma_0(n)
     *
     * @param n Positive integer
     * @param table Prime-factor table
     * @return Number of divisors of n
     */
    static Int divisorCount(const Int& n, const IntFactorTable& table) {
        if (n <= 0) return Int(0);
        if (n.isOne()) return Int(1);

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

    /**
     * @brief Sum of divisors sigma_1(n)
     *
     * @param n Positive integer
     * @param table Prime-factor table
     * @return Sum of all divisors of n
     *
     * Algorithm:
     *   σ(n) = Π_{p^e || n} (p^{e+1} - 1) / (p - 1)
     */
    static Int divisorSum(const Int& n, const IntFactorTable& table) {
        if (n <= 0) return Int(0);
        if (n.isOne()) 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 symbol / Legendre symbol / Kronecker symbol
    // =====================================================================

    /**
     * @brief Jacobi symbol (a/n)
     *
     * When n is a positive odd integer, compute the Jacobi symbol (a/n).
     * Coincides with the Legendre symbol when n is prime.
     *
     * Algorithm: iteration based on quadratic reciprocity
     *   1. Reduce a to a mod n
     *   2. If a is even, factor out 2 and compute (2/n)
     *   3. Apply quadratic reciprocity to invert (a/n) = +/-(n/a)
     *   4. Swap a and n and repeat
     *
     * @param a Arbitrary integer
     * @param n Positive odd integer (n > 0 and n odd)
     * @return One of -1, 0, 1
     *
     * Special cases:
     *   jacobi(0, n) = 0 (n > 1), jacobi(0, 1) = 1
     *   jacobi(1, n) = 1
     *   n even or n <= 0 -> 0 (undefined)
     */
    // uint64-based Jacobi helper: standard binary Jacobi
    // Precondition: n is a positive odd integer (n & 1 == 1); result is the accumulated sign (+/-1)
    static int jacobiSymbol_u64(uint64_t a, uint64_t n, int result) {
        while (a != 0) {
            while ((a & 1) == 0) {
                a >>= 1;
                int n8 = static_cast<int>(n & 7);
                if (n8 == 3 || n8 == 5) result = -result;
            }
            if ((a & 3) == 3 && (n & 3) == 3) result = -result;
            uint64_t tmp = a;
            a = n % a;
            n = tmp;
            if (n == 1) return result;
        }
        return (n == 1) ? result : 0;
    }

    // Convert the bits state to a +/-1 accumulated sign and pass to the 1-limb u64 fast path
    // bits encoding (from gmp-impl.h):
    //   bit 0: e (sign, 0=+1, 1=-1)
    //   bits 1-4: state index -> decodes (a mod 4, b mod 4, d)
    //   When bits >= 16, d=0 (a is denominator): swap a, b to canonical (a/b) form
    static int jacobiSymbol_finish_u64(uint64_t a, uint64_t b, unsigned bits) {
        if (bits == mpn::BITS_FAIL) return 0;
        int sign = 1 - 2 * static_cast<int>(bits & 1);
        if (bits >= 16) std::swap(a, b);
        // At this point b should be odd (the denominator)
        if (b == 0) return (a == 1) ? sign : 0;
        if (b == 1) return sign;
        if (a == 0) return (b == 1) ? sign : 0;
        return jacobiSymbol_u64(a, b, sign);
    }

    // mpn-level Jacobi loop (modeled on GMP's mpn_jacobi_n)
    // ap, np are mutable buffers (allocated by caller, padded to a common width)
    // bits is the state initialized by jacobi_init
    // Return value: +/-1 / 0
    //
    // Convention: ap = "a" (numerator, GMP's a), np = "b" (denominator, GMP's b)
    // Invariant: update bits at each Euclid step (within subdiv or hgcd2_jacobi)
    // The state machine handles 2-stripping automatically (Schönhage's rule via table)
    static int jacobiSymbol_mpn_loop(uint64_t* ap, size_t an,
                                      uint64_t* np, size_t nnn,
                                      unsigned bits) {
        thread_local std::vector<uint64_t> scratch;
        auto ensure_scratch = [&](size_t n) {
            if (scratch.size() < n) scratch.resize(n);
        };

        while (true) {
            // Termination test
            if (an == 0) {
                if (nnn == 1 && np[0] == 1) return jacobi_finish_or_fail(bits);
                return 0;
            }
            if (nnn == 0) {
                if (an == 1 && ap[0] == 1) return jacobi_finish_or_fail(bits);
                return 0;
            }
            if (an == 1 && ap[0] == 1) return jacobi_finish_or_fail(bits);
            if (nnn == 1 && np[0] == 1) return jacobi_finish_or_fail(bits);

            // 1-limb range: u64 fast path
            if (an == 1 && nnn == 1) {
                return jacobiSymbol_finish_u64(ap[0], np[0], bits);
            }

            // padding: align ap, np to a common width
            size_t width = std::max(an, nnn);
            if (an < width) std::memset(ap + an, 0, (width - an) * sizeof(uint64_t));
            if (nnn < width) std::memset(np + nnn, 0, (width - nnn) * sizeof(uint64_t));

            // Recursive HGCD batch (large size) — Phase 3
            // Jacobi-specific threshold: unlike GCD, the bits state machine adds overhead, so
            // recursion starts at a higher threshold than plain hgcd2_jacobi
            // 384 = 24K-bit cutoff. Up to 512, plain hgcd2_jacobi is fastest;
            // beyond that, recursive HGCD wins (per benchmarks)
            constexpr size_t JACOBI_HGCD_THRESHOLD = 384;
            if (width >= JACOBI_HGCD_THRESHOLD) {
                auto& arena = getThreadArena();
                size_t mark = arena.mark();

                mpn::HgcdMatrix M;
                mpn::hgcd_matrix_init(&M, 2 * width + 4);
                uint64_t* hgcd_tp = arena.alloc_limbs(std::max(width, M.alloc) + 16);

                size_t new_n = mpn::hgcd_jacobi(ap, np, width, &M, hgcd_tp, &bits);
                arena.rewind(mark);

                if (new_n > 0) {
                    an = mpn::normalized_size(ap, new_n);
                    nnn = mpn::normalized_size(np, new_n);
                    if (bits == mpn::BITS_FAIL) return 0;
                    continue;
                }
                // No progress -> fall through to hgcd2_jacobi
            }

            // Try hgcd2_jacobi (batch over the top 2 limbs)
            if (width >= 2) {
                uint64_t mask = ap[width-1] | np[width-1];
                if (mask == 0) {
                    // Both top limbs are 0: reduce width and retry
                    width--;
                    an = std::min(an, width);
                    nnn = std::min(nnn, width);
                    continue;
                }

                bool try_hgcd2 = (width >= 2);
                uint64_t ah, al, bh, bl;
                if (try_hgcd2) {
                    if (mask >> 63) {
                        ah = ap[width-1]; al = (width >= 2) ? ap[width-2] : 0;
                        bh = np[width-1]; bl = (width >= 2) ? np[width-2] : 0;
                    } else {
                        int shift = std::countl_zero(mask);
                        if (width >= 3) {
                            ah = (ap[width-1] << shift) | (ap[width-2] >> (64-shift));
                            al = (ap[width-2] << shift) | (ap[width-3] >> (64-shift));
                            bh = (np[width-1] << shift) | (np[width-2] >> (64-shift));
                            bl = (np[width-2] << shift) | (np[width-3] >> (64-shift));
                        } else if (width == 2) {
                            ah = (ap[1] << shift) | (ap[0] >> (64-shift));
                            al = ap[0] << shift;
                            bh = (np[1] << shift) | (np[0] >> (64-shift));
                            bl = np[0] << shift;
                        } else {
                            try_hgcd2 = false;
                        }
                    }
                }

                if (try_hgcd2) {
                    mpn::HgcdMatrix1 M;
                    if (mpn::hgcd2_jacobi(ah, al, bh, bl, &M, &bits)) {
                        ensure_scratch(width + 1);
                        size_t new_w = mpn::hgcd_mul_matrix1_vector(&M, scratch.data(),
                                                                     ap, np, width);
                        std::memcpy(ap, scratch.data(), new_w * sizeof(uint64_t));
                        an = mpn::normalized_size(ap, new_w);
                        nnn = mpn::normalized_size(np, new_w);
                        if (bits == mpn::BITS_FAIL) return 0;
                        continue;
                    }
                }
            }

            // subdiv fallback: bigger -= q * smaller, jacobi_update with d, q&3
            // d=1: a (=ap) shrinks (b is denominator); d=0: b (=np) shrinks (a is denominator)
            {
                int cmp_result;
                if (an > nnn) cmp_result = 1;
                else if (an < nnn) cmp_result = -1;
                else cmp_result = mpn::cmp(ap, an, np, nnn);

                if (cmp_result == 0) {
                    // ap == np: gcd. If 1, jacobi_finish; otherwise 0
                    if (an == 1 && ap[0] == 1) return jacobi_finish_or_fail(bits);
                    return 0;
                }

                uint64_t* rp; uint64_t* lp;
                size_t rn, ln;
                int d;
                if (cmp_result > 0) { rp = ap; rn = an; lp = np; ln = nnn; d = 1; }
                else                { rp = np; rn = nnn; lp = ap; ln = an; d = 0; }

                size_t qn_cap = rn - ln + 1;
                size_t div_sz = mpn::divide_scratch_size(rn, ln);
                size_t need = ln + qn_cap + div_sz;
                ensure_scratch(need);

                uint64_t* tp = scratch.data();        // remainder buffer (size ln)
                uint64_t* qp = tp + ln;                // quotient (size qn_cap)
                uint64_t* dp = qp + qn_cap;            // div scratch

                mpn::divide(qp, tp, rp, rn, lp, ln, dp);
                size_t rem_n = mpn::normalized_size(tp, ln);
                size_t qn = mpn::normalized_size(qp, qn_cap);

                if (rem_n == 0) {
                    // gcd = lp. If 1, finish; otherwise 0
                    if (qn > 0) bits = mpn::jacobi_update(bits, d, qp[0] & 3);
                    if (ln == 1 && lp[0] == 1) return jacobi_finish_or_fail(bits);
                    return 0;
                }

                if (qn > 0) {
                    bits = mpn::jacobi_update(bits, d, qp[0] & 3);
                    if (bits == mpn::BITS_FAIL) return 0;
                }

                // rp ← remainder
                std::memcpy(rp, tp, rem_n * sizeof(uint64_t));
                if (rem_n < rn) std::memset(rp + rem_n, 0, (rn - rem_n) * sizeof(uint64_t));

                if (cmp_result > 0) an = rem_n;
                else nnn = rem_n;
            }
        }
    }

    // bits termination: 0 if BITS_FAIL, otherwise jacobi_finish
    static int jacobi_finish_or_fail(unsigned bits) {
        return bits == mpn::BITS_FAIL ? 0 : mpn::jacobi_finish(bits);
    }

    static int jacobiSymbol(const Int& a, const Int& n) {
        // n must be a positive odd integer
        if (n <= 0 || n.isEven()) return 0;
        if (n.isOne()) return 1;

        // Reduce a modulo 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.isOne()) return 1;

        // mpn path: for >= 3 limbs, go to the mpn-level loop (avoid Int construction/normalize overhead)
        if (a_mod.size() >= 3 || n.size() >= 3) {
            std::vector<uint64_t> aa(a_mod.data(), a_mod.data() + a_mod.size());
            std::vector<uint64_t> nn(n.data(), n.data() + n.size());
            // Expand aa's capacity to match nn (Lehmer-matrix application requires zero padding)
            if (aa.size() < nn.size()) aa.resize(nn.size(), 0);
            if (nn.size() < aa.size()) nn.resize(aa.size(), 0);
            // bits init: a = a_mod (numerator), b = n (denominator), s = 0
            unsigned bits = mpn::jacobi_init(static_cast<unsigned>(aa[0] & 3),
                                              static_cast<unsigned>(nn[0] & 3),
                                              0);
            return jacobiSymbol_mpn_loop(aa.data(), a_mod.size(),
                                          nn.data(), n.size(), bits);
        }

        // Small size (<= 2 limbs): the naive Int-level implementation is fast enough
        Int aa = a_mod;
        Int nn = n;
        int result = 1;

        while (true) {
            if (aa.isZero()) return 0;
            if (aa.isOne()) return result;

            if (aa.bitLength() <= 64 && nn.bitLength() <= 64) {
                return jacobiSymbol_u64(aa.word(0), nn.word(0), result);
            }

            while (aa.isEven()) {
                aa >>= 1;
                int n8 = static_cast<int>(nn.word(0) & 7);
                if (n8 == 3 || n8 == 5) result = -result;
            }
            if (aa.isOne()) return result;

            int a4 = static_cast<int>(aa.word(0) & 3);
            int n4 = static_cast<int>(nn.word(0) & 3);
            if (a4 == 3 && n4 == 3) result = -result;

            Int temp = aa;
            aa = nn % temp;
            nn = temp;
        }
    }

    /**
     * @brief Legendre symbol (a/p)
     *
     * When p is an odd prime, determines whether a is a quadratic residue modulo p.
     * A wrapper around the Jacobi symbol.
     *
     * @param a Arbitrary integer
     * @param p Odd prime
     * @return -1: non-residue, 0: p | a, 1: residue
     */
    static int legendreSymbol(const Int& a, const Int& p) {
        return jacobiSymbol(a, p);
    }

    /**
     * @brief Extended Kronecker symbol (a/n)
     *
     * The extended Jacobi symbol, defined also for n = 0, 1, 2, and negative integers.
     *
     * @param a Arbitrary integer
     * @param n Arbitrary integer
     * @return Value of the Kronecker symbol
     */
    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.isOne() ? 1 : 0;
        }

        Int nn = n;
        int result = 1;

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

        // Factor out 2 from n: (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.word(0) & 7);  // toUInt64 throws for 2+ limbs
                if (a8 == 3 || a8 == 5) result = -result;
            }
        }

        // Done if nn is 1
        if (nn.isOne()) return result;

        // What remains is the Jacobi symbol for the odd part
        return result * jacobiSymbol(a, nn);
    }
};

} // namespace sangi