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

// Polynomial_factorization.hpp
// Exact factorization of integer polynomials
//
// Algorithm:
//   1. Extract content (GCD of all coefficients)
//   2. Separate linear factors via the rational-root theorem
//   3. Yun's square-free decomposition (separate repeated factors via gcd(f, f'))
//   4. Try higher-degree factors via Kronecker's method
//
// Usage:
//   #include <math/core/Polynomial.hpp>
//   #include <math/core/Polynomial_factorization.hpp>
//   #include <math/core/FactoredInteger.hpp>  // result type

#pragma once

#include <math/core/Polynomial.hpp>
#include <math/core/FactoredInteger.hpp>
#include <math/core/matrix.hpp>
#include <numeric>
#include <cmath>
#include <algorithm>
#include <set>
#include <cstdint>
#include <functional>

namespace sangi {

// ================================================================
// Internal helpers
// ================================================================

namespace detail {

// Enumerate divisors of an integer (positive divisors only, sorted)
inline std::vector<int> divisors(int n) {
    if (n < 0) n = -n;
    if (n == 0) return {};
    std::vector<int> result;
    for (int i = 1; i * i <= n; ++i) {
        if (n % i == 0) {
            result.push_back(i);
            if (i != n / i)
                result.push_back(n / i);
        }
    }
    std::sort(result.begin(), result.end());
    return result;
}

// Enumerate divisors of an integer (int64_t version)
inline std::vector<int64_t> divisors(int64_t n) {
    if (n < 0) n = -n;
    if (n == 0) return {};
    std::vector<int64_t> result;
    for (int64_t i = 1; i * i <= n; ++i) {
        if (n % i == 0) {
            result.push_back(i);
            if (i != n / i)
                result.push_back(n / i);
        }
    }
    std::sort(result.begin(), result.end());
    return result;
}

// Recover polynomial coefficients via Lagrange interpolation
// xs, ys: evaluation points and values (n+1 of them -> polynomial of degree at most n)
template<typename T>
Polynomial<T> lagrangeInterpolationPoly(
    const std::vector<T>& xs,
    const std::vector<T>& ys)
{
    int n = static_cast<int>(xs.size());
    Polynomial<T> result(T(0));
    Polynomial<T> one(T(1));

    for (int i = 0; i < n; ++i) {
        if (ys[i] == T(0)) continue;

        // L_i(x) = Π_{j≠i} (x - x_j) / (x_i - x_j)
        Polynomial<T> basis = one;
        T denom = T(1);
        for (int j = 0; j < n; ++j) {
            if (j == i) continue;
            // (x - x_j)
            basis = basis * Polynomial<T>({T(0) - xs[j], T(1)});
            denom = denom * (xs[i] - xs[j]);
        }
        // ys[i] / denom * basis
        result = result + basis * (ys[i] / denom);
    }
    return result;
}

// Check whether the polynomial has integer coefficients (always true when T is an integer type)
template<typename T>
bool hasIntegerCoeffs(const Polynomial<T>& p) {
    if constexpr (std::is_integral_v<T>) {
        return true;
    } else {
        // Floating-point case: check whether each coefficient is sufficiently close to an integer
        for (int i = 0; i <= p.degree(); ++i) {
            double v = static_cast<double>(p[i]);
            if (std::abs(v - std::round(v)) > 1e-9) return false;
        }
        return true;
    }
}

// Convert the polynomial to integer coefficients and return (Lagrange rational result -> integer-coefficient check)
template<typename T>
bool tryRoundToInt(const Polynomial<double>& src, Polynomial<T>& dst) {
    std::vector<T> coeffs(src.degree() + 1);
    for (int i = 0; i <= src.degree(); ++i) {
        double v = src[i];
        double r = std::round(v);
        if (std::abs(v - r) > 1e-6) return false;
        coeffs[i] = static_cast<T>(static_cast<int64_t>(r));
    }
    dst = Polynomial<T>(std::move(coeffs));
    return true;
}

} // namespace detail

// ================================================================
// GCD of integer polynomials (pseudo-remainder based)
// ================================================================

/// Pseudo-remainder of integer polynomials
/// prem(f, g) = lc(g)^(deg(f)-deg(g)+1) * f mod g
/// The result always keeps integer coefficients
template<PolynomialCoeff T>
[[nodiscard]] Polynomial<T> pseudoRemainder(
    const Polynomial<T>& f, const Polynomial<T>& g)
{
    static_assert(std::is_integral_v<T>);
    if (g.isZero()) return f;
    if (f.degree() < g.degree()) return f;

    int delta = f.degree() - g.degree() + 1;
    T lc_g = g.leadingCoefficient();

    // lc(g)^delta * f
    T scale = T(1);
    for (int i = 0; i < delta; ++i) scale *= lc_g;

    Polynomial<T> r = f * scale;
    auto dr = r.divmod(g);
    return dr.remainder;
}

/// GCD of integer polynomials (subresultant PRS)
/// At each step take the pseudo-remainder and shrink the coefficients via primitivePart
template<PolynomialCoeff T>
[[nodiscard]] Polynomial<T> intPolyGcd(Polynomial<T> a, Polynomial<T> b) {
    static_assert(std::is_integral_v<T>);

    // Zero check
    if (a.isZero()) return b;
    if (b.isZero()) return a;

    // Make both primitive
    a = primitivePart(a);
    b = primitivePart(b);

    // Ensure deg(a) >= deg(b)
    if (a.degree() < b.degree()) std::swap(a, b);

    while (!b.isZero()) {
        Polynomial<T> r = pseudoRemainder(a, b);
        if (r.isZero()) break;
        r = primitivePart(r);
        a = std::move(b);
        b = std::move(r);
    }

    // If b is 0, a is the GCD; otherwise b is the GCD
    Polynomial<T>& result = b.isZero() ? a : b;

    // Make the leading coefficient positive
    if (result.leadingCoefficient() < T(0))
        result = -result;

    return primitivePart(result);
}

// ================================================================
// Square-free decomposition (Yun's algorithm)
// ================================================================

/// Decompose polynomial f into square-free factors via Yun's algorithm
/// Result: f = c * Product f_i^i  (f_i are square-free and mutually coprime)
/// Returns FactoredInteger<Polynomial<T>> holding each factor and its exponent
///
/// T is an integer type (int, int64_t, etc.)
template<PolynomialCoeff T>
[[nodiscard]] FactoredInteger<Polynomial<T>> squareFreeFactorization(
    const Polynomial<T>& f)
{
    static_assert(std::is_integral_v<T>,
                  "squareFreeFactorization requires integral coefficient type");

    FactoredInteger<Polynomial<T>> result;

    if (f.isZero() || f.degree() <= 0) {
        if (!f.isZero()) result.addFactor(f, 1);
        return result;
    }

    Polynomial<T> fd = f.derivative();
    if (fd.isZero()) {
        // f' = 0 -> f may be constant (or in characteristic p, all exponents are multiples of p)
        result.addFactor(f, 1);
        return result;
    }

    // Use pseudo-remainder-based integer-polynomial GCD
    Polynomial<T> pp = primitivePart(f);
    Polynomial<T> ppd = pp.derivative();

    Polynomial<T> g = intPolyGcd(pp, ppd);
    Polynomial<T> v = pp / g;
    Polynomial<T> w = ppd / g;

    for (int n = 1; !v.isZero() && v.degree() > 0; ++n) {
        Polynomial<T> vd = v.derivative();
        Polynomial<T> diff = w - vd;

        Polynomial<T> h = intPolyGcd(v, diff);

        if (h.degree() > 0) {
            // Make the leading coefficient positive
            if (h.leadingCoefficient() < T(0))
                h = -h;
            result.addFactor(h, n);
        }

        v = v / h;
        if (diff.isZero()) break;
        w = diff / h;
    }

    // Add the remaining v if it is not constant
    if (!v.isZero() && v.degree() > 0) {
        if (v.leadingCoefficient() < T(0)) v = -v;
        result.addFactor(v, 1);
    }

    return result;
}

// ================================================================
// Separation of linear factors via the rational-root theorem
// ================================================================

namespace detail {

/// Rational-root theorem: search for rationals b/a with f(b/a) = 0
/// and separate all (ax - b) factors
/// f must be primitive (content = 1)
template<PolynomialCoeff T>
void extractLinearFactors(
    Polynomial<T>& f,
    FactoredInteger<Polynomial<T>>& factors,
    int multiplicity = 1)
{
    static_assert(std::is_integral_v<T>);

    while (f.degree() >= 1) {
        T f0 = f[0];
        T fn = f.leadingCoefficient();

        if (f0 < T(0)) f0 = -f0;
        if (fn < T(0)) fn = -fn;

        auto divs_f0 = divisors(static_cast<int64_t>(f0));
        auto divs_fn = divisors(static_cast<int64_t>(fn));

        // If f0 = 0, then x = 0 is a root -> divisible by (x)
        if (f[0] == T(0)) {
            divs_f0 = {0};
        }

        bool found = false;
        // Candidate: x = +/- b/a (b | f0, a | fn)
        // Test in integer polynomials: f(b/a) = 0 iff the remainder of f divided by (ax - b) is 0
        for (auto a : divs_fn) {
            if (a == 0) continue;
            for (auto b : divs_f0) {
                for (int sign : {1, -1}) {
                    T bv = static_cast<T>(b * sign);
                    T av = static_cast<T>(a);
                    // Candidate factor: (av * x - bv) = av * x + (-bv)
                    Polynomial<T> candidate({T(0) - bv, av});

                    auto dr = f.divmod(candidate);
                    if (dr.remainder.isZero()) {
                        if (candidate.leadingCoefficient() < T(0))
                            candidate = -candidate;
                        factors.addFactor(candidate, multiplicity);
                        f = std::move(dr.quotient);
                        found = true;
                        goto next_round;
                    }
                }
            }
        }
    next_round:
        if (!found) break;
    }
}

} // namespace detail

// ================================================================
// Factorization via Kronecker's method
// ================================================================

namespace detail {

/// Kronecker's algorithm:
/// Factors of deg(f) = n have degree at most n/2 -> evaluate at n/2 + 1 points,
/// then try every combination of divisors of the values via Lagrange interpolation
template<PolynomialCoeff T>
void kroneckerFactor(
    Polynomial<T>& f,
    FactoredInteger<Polynomial<T>>& factors,
    int multiplicity = 1)
{
    static_assert(std::is_integral_v<T>);

    while (f.degree() >= 2) {
        int deg = f.degree();
        int half = deg / 2;
        int npts = half + 1;

        // Choose evaluation points (centered at 0)
        std::vector<int64_t> xs(npts);
        int offset = half / 2;
        for (int i = 0; i < npts; ++i)
            xs[i] = i - offset;

        // Compute the value at each point and enumerate divisors
        std::vector<std::vector<int64_t>> allDivs(npts);
        for (int i = 0; i < npts; ++i) {
            T val = f(static_cast<T>(xs[i]));
            int64_t v = static_cast<int64_t>(val);
            auto d = divisors(v);
            std::vector<int64_t> fullDivs;
            fullDivs.reserve(d.size() * 2);
            for (auto dv : d) {
                fullDivs.push_back(dv);
                fullDivs.push_back(-dv);
            }
            if (v == 0) {
                fullDivs.clear();
                for (int64_t k = -10; k <= 10; ++k) fullDivs.push_back(k);
            }
            allDivs[i] = std::move(fullDivs);
        }

        size_t totalCombs = 1;
        for (int i = 0; i < npts; ++i) {
            if (allDivs[i].empty()) { totalCombs = 0; break; }
            totalCombs *= allDivs[i].size();
            if (totalCombs > 100000) { totalCombs = 100000; break; }
        }

        bool found = false;
        for (size_t idx = 0; idx < totalCombs; ++idx) {
            std::vector<double> xsd(npts), ysd(npts);
            size_t rem = idx;
            for (int i = 0; i < npts; ++i) {
                size_t j = rem % allDivs[i].size();
                rem /= allDivs[i].size();
                xsd[i] = static_cast<double>(xs[i]);
                ysd[i] = static_cast<double>(allDivs[i][j]);
            }

            Polynomial<double> cand_d = lagrangeInterpolationPoly(xsd, ysd);
            Polynomial<T> candidate;
            if (!tryRoundToInt(cand_d, candidate)) continue;
            if (candidate.degree() <= 0) continue;
            if (candidate.degree() > half) continue;

            T candLC = candidate.leadingCoefficient();
            if (candLC < T(0)) { candidate = -candidate; candLC = -candLC; }
            T fLC = f.leadingCoefficient();
            if (fLC < T(0)) fLC = -fLC;
            if (candLC == T(0) || fLC % candLC != T(0)) continue;

            auto dr = f.divmod(candidate);
            if (dr.remainder.isZero()) {
                factors.addFactor(candidate, multiplicity);
                f = std::move(dr.quotient);
                found = true;
                break;
            }
        }

        if (!found) break;
    }
}

} // namespace detail

// ================================================================
// Factorization (main entry point)
// ================================================================

/// Factor the integer polynomial f into irreducible factors
/// Result: FactoredInteger<Polynomial<T>> with f = content * Product factor_i ^ exp_i
///
/// Algorithm:
///   1. Extract content (GCD of all coefficients) as a constant factor
///   2. Separate linear factors via the rational-root theorem
///   3. Separate repeated factors via Yun's square-free decomposition
///   4. For each square-free factor, separate higher-degree factors via Kronecker's method
template<PolynomialCoeff T>
[[nodiscard]] FactoredInteger<Polynomial<T>> factorize(const Polynomial<T>& f) {
    static_assert(std::is_integral_v<T>,
                  "factorize requires integral coefficient type (int or int64_t)");

    FactoredInteger<Polynomial<T>> result;

    if (f.isZero()) return result;

    // 1. Extract content
    T c = content(f);
    Polynomial<T> prim = f;
    if (c != T(1) && c != T(0)) {
        prim = primitivePart(f);
        if (c != T(1))
            result.addFactor(Polynomial<T>(c), 1);
    }

    // Make the leading coefficient positive
    if (!prim.isZero() && prim.leadingCoefficient() < T(0)) {
        prim = -prim;
        if (result.size() > 0 && result[0].degree() == 0)
            result[0] = Polynomial<T>(T(0) - result[0][0]);
        else
            result.addFactor(Polynomial<T>(T(-1)), 1);
    }

    if (prim.degree() <= 0) {
        if (!prim.isZero() && !(prim[0] == T(1)))
            result.addFactor(prim, 1);
        return result;
    }
    if (prim.degree() == 1) {
        result.addFactor(prim, 1);
        return result;
    }

    // 2. Separate linear factors first via the rational-root theorem
    detail::extractLinearFactors(prim, result);
    if (prim.degree() <= 0) return result;
    if (prim.degree() == 1) {
        if (prim.leadingCoefficient() < T(0)) prim = -prim;
        result.addFactor(prim, 1);
        return result;
    }

    // 3. Square-free decomposition
    auto sqfree = squareFreeFactorization(prim);
    if (sqfree.empty()) {
        result.addFactor(prim, 1);
        return result;
    }

    // 4. Factor each square-free factor via Kronecker's method
    for (size_t i = 0; i < sqfree.size(); ++i) {
        Polynomial<T> fi = sqfree[i];
        int mult = sqfree.exponent(i);
        if (fi.degree() <= 0) continue;

        detail::extractLinearFactors(fi, result, mult);
        if (fi.degree() <= 0) continue;
        if (fi.degree() == 1) {
            if (fi.leadingCoefficient() < T(0)) fi = -fi;
            result.addFactor(fi, mult);
            continue;
        }

        detail::kroneckerFactor(fi, result, mult);
        if (fi.degree() > 0) {
            if (fi.leadingCoefficient() < T(0)) fi = -fi;
            result.addFactor(fi, mult);
        }
    }

    return result;
}



// ========================================================================
// CAS-4: Resultant / Discriminant (Sylvester-matrix based)
// ========================================================================

/// Build the Sylvester matrix
/// f(x) of degree m, g(x) of degree n -> (m+n) x (m+n) matrix
template<PolynomialCoeff T>
Matrix<T> sylvesterMatrix(const Polynomial<T>& f, const Polynomial<T>& g) {
    int m = f.degree(), n = g.degree();
    int sz = m + n;
    Matrix<T> S(sz, sz, T(0));

    // Place coefficients of f shifted across n rows
    for (int i = 0; i < n; ++i)
        for (int j = 0; j <= m; ++j)
            S(i, i + j) = f[m - j];  // Descending order

    // Place coefficients of g shifted across m rows
    for (int i = 0; i < m; ++i)
        for (int j = 0; j <= n; ++j)
            S(n + i, i + j) = g[n - j];  // Descending order

    return S;
}

namespace detail {

// Ring power base^e for e >= 0 (repeated multiplication, no division).
template<typename T>
[[nodiscard]] T ringPow(const T& base, int e) {
    T r = T(1);
    for (int i = 0; i < e; ++i) r = r * base;
    return r;
}

// Pseudo-remainder prem(f, g) = lc(g)^(deg f - deg g + 1) * f mod g.
// Self-contained version (no std::is_integral_v constraint) so it also works
// for multiprecision integer types such as sangi::Int. Requires deg f >= deg g.
// All coefficient divisions inside divmod() are exact thanks to the pre-scaling.
template<typename T>
[[nodiscard]] Polynomial<T> pseudoRemainderRing(
    const Polynomial<T>& f, const Polynomial<T>& g)
{
    int delta = f.degree() - g.degree() + 1;
    T scale = ringPow(g.leadingCoefficient(), delta);
    Polynomial<T> r = f * scale;
    return r.divmod(g).remainder;
}

// Divide every coefficient of p by the scalar d. The caller guarantees the
// division is exact (subresultant theory), so integer division loses nothing.
template<typename T>
[[nodiscard]] Polynomial<T> exactScalarDivide(const Polynomial<T>& p, const T& d) {
    std::vector<T> c = p.coefficients();
    for (auto& ci : c) ci = ci / d;
    return Polynomial<T>(std::move(c));
}

// Resultant via the subresultant polynomial remainder sequence
// (Cohen, "A Course in Computational Algebraic Number Theory", Algorithm 3.3.7).
//
// Compared with the Sylvester-determinant route (O((n+m)^3) ring operations on
// coefficients that can blow up to the product of all minors), the subresultant
// PRS performs O(min(n,m)) pseudo-divisions while the fraction-free quotient
// g * h^delta keeps every intermediate coefficient bounded by a single
// subresultant (Hadamard bound). This is decisive for multiprecision coefficients.
//
// Requires deg A >= 1 and deg B >= 1 (constant / zero arguments are handled by
// the public resultant() wrapper).
template<typename T>
[[nodiscard]] T resultantSubresultantPRS(Polynomial<T> A, Polynomial<T> B) {
    int s = 1;
    if (A.degree() < B.degree()) {
        std::swap(A, B);
        // Res(A, B) = (-1)^(deg A * deg B) Res(B, A); the sign is -1 iff both degrees are odd.
        if ((A.degree() % 2 == 1) && (B.degree() % 2 == 1)) s = -s;
    }

    T g = T(1), h = T(1);
    while (true) {
        int delta = A.degree() - B.degree();
        if ((A.degree() % 2 == 1) && (B.degree() % 2 == 1)) s = -s;

        Polynomial<T> R = pseudoRemainderRing(A, B);
        A = std::move(B);

        // B = R / (g * h^delta)   (exact)
        T denom = g * ringPow(h, delta);
        B = exactScalarDivide(R, denom);

        g = A.leadingCoefficient();
        // h = h^(1-delta) * g^delta. For delta >= 1 this is g^delta / h^(delta-1)
        // (exact); for delta == 0 it leaves h unchanged.
        if (delta >= 1)
            h = ringPow(g, delta) / ringPow(h, delta - 1);

        if (B.isZero() || B.degree() == 0) break;
    }

    // A common factor of higher degree survived: the resultant vanishes.
    if (B.isZero()) return T(0);

    // B is a non-zero constant and deg A = d >= 1.
    int d = A.degree();
    T finalH = ringPow(B.leadingCoefficient(), d) / ringPow(h, d - 1);
    return T(s) * finalH;
}

} // namespace detail

/// Resultant Res(f, g).
///
/// For integer coefficient types the subresultant PRS route is taken
/// (fraction-free, coefficient growth bounded by the subresultants); for field
/// types it falls back to the Sylvester-matrix determinant.
template<PolynomialCoeff T>
[[nodiscard]] T resultant(const Polynomial<T>& f, const Polynomial<T>& g) {
    if (f.isZero() || g.isZero()) return T(0);

    const int m = f.degree();
    const int n = g.degree();

    // Constant arguments: Res(c, B) = c^deg(B), Res(A, c) = c^deg(A), Res(c, d) = 1.
    if (m == 0 && n == 0) return T(1);
    if (m == 0) return detail::ringPow(f[0], n);
    if (n == 0) return detail::ringPow(g[0], m);

    if constexpr (numeric_traits<T>::is_integer) {
        return detail::resultantSubresultantPRS(f, g);
    } else {
        // Field coefficients: the determinant route is exact and well-conditioned.
        return determinant(sylvesterMatrix(f, g));
    }
}

/// Discriminant Disc(f) = (-1)^(n(n-1)/2) * Res(f, f') / a_n
/// a_n = leading coefficient
template<PolynomialCoeff T>
[[nodiscard]] T discriminant(const Polynomial<T>& f) {
    int n = f.degree();
    if (n < 1) return T(0);
    auto fp = f.derivative();
    T res = resultant(f, fp);
    T an = f[n];
    T sign = ((n * (n - 1) / 2) % 2 == 0) ? T(1) : T(-1);
    if (an != T(0)) return sign * res / an;
    return sign * res;
}

// ========================================================================
// CAS-10: integer partitions
// ========================================================================

/// Enumerate all partitions of n (descending partitions)
[[nodiscard]] inline std::vector<std::vector<int>> integerPartitions(int n) {
    std::vector<std::vector<int>> result;
    if (n <= 0) { result.push_back({}); return result; }

    std::function<void(int, int, std::vector<int>&)> gen;
    gen = [&](int remaining, int maxPart, std::vector<int>& current) {
        if (remaining == 0) {
            result.push_back(current);
            return;
        }
        for (int k = std::min(remaining, maxPart); k >= 1; --k) {
            current.push_back(k);
            gen(remaining - k, k, current);
            current.pop_back();
        }
    };
    std::vector<int> buf;
    gen(n, n, buf);
    return result;
}

/// Partition number p(n) (dynamic programming)
[[nodiscard]] inline uint64_t partitionCount(int n) {
    if (n < 0) return 0;
    std::vector<uint64_t> dp(n + 1, 0);
    dp[0] = 1;
    for (int k = 1; k <= n; ++k)
        for (int j = k; j <= n; ++j)
            dp[j] += dp[j - k];
    return dp[n];
}

} // namespace sangi
