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

// RationalFunction_partialFractions.hpp
// Partial-fraction expansion of rational functions (Heaviside's method)
//
// Usage:
//   #include <math/core/RationalFunction.hpp>
//   #include <math/core/RationalFunction_partialFractions.hpp>

#pragma once

#include <math/core/RationalFunction.hpp>
#include <math/core/Complex.hpp>
#include <math/roots/polynomial_roots.hpp>
#include <vector>

namespace sangi {

// ================================================================
// Struct holding the result of a partial-fraction expansion
// ================================================================

/// Each term of the partial-fraction expansion: coefficient / (x - pole)
template<typename T>
struct PartialFraction {
    Complex<T> coefficient;  // Residue
    Complex<T> pole;         // Pole
};

/// Result of the partial-fraction expansion
template<typename T>
struct PartialFractionExpansion {
    Polynomial<T> quotient;                     // Polynomial part (when deg(num) >= deg(den))
    std::vector<PartialFraction<T>> terms;      // Each partial-fraction term
    bool success = false;                       // Whether the expansion succeeded
};

// ================================================================
// Partial-fraction expansion (Heaviside's method)
// ================================================================

/// Compute the partial-fraction expansion of the rational function f = P(x)/Q(x).
/// Precondition: all roots of Q(x) are simple (multiple roots are unsupported)
///
/// Result: f(x) = quotient(x) + Sum A_k / (x - r_k)
/// where A_k = P(r_k) / Q'(r_k) (residue formula)
///
/// T is assumed to be a floating-point type (double, float)
template<typename T>
[[nodiscard]] PartialFractionExpansion<T> partialFractions(
    const RationalFunction<T>& f,
    T eps = std::numeric_limits<T>::epsilon() * T(100),
    size_t maxIter = 1000)
{
    static_assert(std::is_floating_point_v<T>,
                  "partialFractions requires floating-point coefficient type");

    PartialFractionExpansion<T> result;

    const auto& P = f.numerator();
    const auto& Q = f.denominator();

    if (Q.isZero()) {
        result.success = false;
        return result;
    }

    // When deg(P) >= deg(Q), extract the polynomial part via polynomial division
    Polynomial<T> remainder = P;
    if (P.degree() >= Q.degree()) {
        auto dr = P.divmod(Q);
        result.quotient = std::move(dr.quotient);
        remainder = std::move(dr.remainder);
    }

    // If the denominator is constant, partial fractions are unnecessary
    if (Q.degree() <= 0) {
        result.success = true;
        return result;
    }

    // Find all roots of Q(x)
    auto roots = solvePolynomial(Q, eps, maxIter);

    // Compute Q'(x)
    Polynomial<T> Qprime = Q.derivative();

    // For each root r_k, compute the residue A_k = remainder(r_k) / Q'(r_k)
    result.terms.reserve(roots.size());
    for (const auto& r : roots) {
        Complex<T> numVal = remainder(r);
        Complex<T> denVal = Qprime(r);

        if (abs(denVal) < eps) {
            // Q'(r_k) ~ 0 -> possible multiple root -> unsupported
            result.success = false;
            return result;
        }

        PartialFraction<T> term;
        term.pole = r;
        term.coefficient = numVal / denVal;
        result.terms.push_back(term);
    }

    result.success = true;
    return result;
}

/// Convert the partial-fraction expansion to a string
template<typename T>
[[nodiscard]] std::string toString(const PartialFractionExpansion<T>& pfe,
                                    const std::string& var = "x") {
    if (!pfe.success) return "(expansion failed)";

    std::ostringstream oss;
    bool first = true;

    // Polynomial part
    if (!pfe.quotient.isZero()) {
        oss << pfe.quotient.toString(var);
        first = false;
    }

    // Each partial-fraction term
    for (const auto& t : pfe.terms) {
        if (!first) oss << " + ";

        oss << "(" << t.coefficient << ")/(" << var << " - (" << t.pole << "))";
        first = false;
    }

    if (first) oss << "0";
    return oss.str();
}

// ================================================================
// Enumerate poles (roots of the denominator)
// ================================================================

/// Return the poles of the rational function f = P(x)/Q(x) (roots of Q(x) = 0)
template<typename T>
[[nodiscard]] std::vector<Complex<T>> poles(
    const RationalFunction<T>& f,
    T eps = std::numeric_limits<T>::epsilon() * T(100),
    size_t maxIter = 1000)
{
    static_assert(std::is_floating_point_v<T>,
                  "poles requires floating-point coefficient type");
    const auto& Q = f.denominator();
    if (Q.degree() <= 0) return {};
    return solvePolynomial(Q, eps, maxIter);
}

// ================================================================
// Enumerate zeros (roots of the numerator)
// ================================================================

/// Return the zeros of the rational function f = P(x)/Q(x) (roots of P(x) = 0)
template<typename T>
[[nodiscard]] std::vector<Complex<T>> zeros(
    const RationalFunction<T>& f,
    T eps = std::numeric_limits<T>::epsilon() * T(100),
    size_t maxIter = 1000)
{
    static_assert(std::is_floating_point_v<T>,
                  "zeros requires floating-point coefficient type");
    const auto& P = f.numerator();
    if (P.degree() <= 0) return {};
    return solvePolynomial(P, eps, maxIter);
}

// ================================================================
// Residue (compute the residue at a specified pole)
// ================================================================

/// Residue at the simple pole `pole` of the rational function f = P(x)/Q(x)
/// Res(f, pole) = P(pole) / Q'(pole)
template<typename T>
[[nodiscard]] Complex<T> residue(
    const RationalFunction<T>& f,
    const Complex<T>& pole)
{
    static_assert(std::is_floating_point_v<T>,
                  "residue requires floating-point coefficient type");
    Complex<T> numVal = f.numerator()(pole);
    Complex<T> denDerVal = f.denominator().derivative()(pole);
    return numVal / denDerVal;
}

} // namespace sangi