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

// mittag_leffler.hpp
//
// Main implementation of the Mittag-Leffler function E_α(z) and the two-parameter E_{α,β}(z).
//
// Definitions:
//   E_α(z) = E_{α,1}(z) = Σ_{k=0}^∞ z^k / Γ(αk + 1)
//   E_{α,β}(z)         = Σ_{k=0}^∞ z^k / Γ(αk + β)
//
// Functions provided:
//   mittagLeffler(alpha, z)         — E_α(z)         (β=1)
//   mittagLeffler(alpha, beta, z)   — E_{α,β}(z)
//
// Supported types: float, double, long double (native float only)
//   Arbitrary-precision Float / Complex<Float> versions are planned for in a future release
//
// Algorithm (switches because Taylor loses precision / slows down for large |z|):
//   |z| ≤ |z|_cutoff(α): Taylor series (Σ z^k / Γ(αk+β))
//   |z| >  |z|_cutoff(α), z > 0:
//     Wiman asymptotics (DLMF 10.46.4 / Gorenflo et al. eq. 4.4.16):
//       E_{α,β}(z) ≈ (1/α) z^{(1-β)/α} exp(z^{1/α})
//                    - Σ_{k=1}^N z^{-k} / Γ(β - αk)
//   |z| >  |z|_cutoff(α), z < 0:
//     decaying asymptotics (heavy tail, removing Wiman's exponential term):
//       E_{α,β}(z) ≈ - Σ_{k=1}^N z^{-k} / Γ(β - αk)
//
// Since the asymptotic series diverges, truncate once the term magnitude starts increasing.
//
// Special cases:
//   α = 1, β = 1: E_{1,1}(z) = exp(z)
//   α = 1, β = 2: E_{1,2}(z) = (exp(z) - 1) / z
//   α = 2, β = 1, z ≥ 0: E_{2,1}(z) = cosh(√z)
//   α = 2, β = 2, z ≥ 0: E_{2,2}(z) = sinh(√z) / √z
//
// Constraints:
//   - α ∈ (0, 2] (α ≤ 0 throws)
//   - For now z is real only. For negative z, z^{1/α} becomes complex, but this
//     implementation returns only the real value of the decaying asymptotic system in that region
//   - For α extremely close to 0 (α < 0.05) convergence is extremely slow → truncated at max iterations

#ifndef SANGI_SPECIAL_MITTAG_LEFFLER_HPP
#define SANGI_SPECIAL_MITTAG_LEFFLER_HPP

#include <math/special/special_common.hpp>
#include <cmath>
#include <limits>
#include <stdexcept>

namespace sangi {
namespace special {

namespace detail {

    // Reciprocal gamma that treats the poles of Γ(x) as 0
    template<IsNativeFloat T>
    [[nodiscard]] T reciprocal_gamma_safe(T x) {
        // non-positive integer ⇒ pole of Γ, 1/Γ = 0
        if (x <= T(0) && x == std::floor(x)) return T(0);
        const T g = std::tgamma(x);
        if (!std::isfinite(g) || g == T(0)) return T(0);
        return T(1) / g;
    }

    // Taylor series: E_{α,β}(z) = Σ_{k=0}^∞ z^k / Γ(αk + β)
    // Recurrence: t_{k+1} = t_k * z * Γ(αk + β) / Γ(α(k+1) + β)
    //              = t_k * z / Γ(α(k+1) + β) * Γ(αk + β)
    // Evaluate t_k = z^k / Γ(αk + β) directly term by term (recurrence is complex since α is not an integer).
    template<IsNativeFloat T>
    [[nodiscard]] T mittag_leffler_taylor(T alpha, T beta, T z, int max_iter = 500) {
        // k = 0 term
        T sum = reciprocal_gamma_safe(beta);
        if (z == T(0)) return sum;

        // Stabilize the accumulated z^k by splitting into |z|^k and the sign ((-1)^k when z<0)
        const T abs_z = std::abs(z);
        const bool z_neg = z < T(0);

        const T eps = std::numeric_limits<T>::epsilon();
        T prev_sum = sum;
        int small_count = 0;
        T abs_zk = T(1);

        for (int k = 1; k < max_iter; ++k) {
            abs_zk *= abs_z;  // |z|^k
            const T inv_gamma = reciprocal_gamma_safe(alpha * T(k) + beta);
            if (inv_gamma == T(0)) {
                // term is 0 at a pole of Γ → skip
                continue;
            }
            T term = abs_zk * inv_gamma;
            if (z_neg && (k & 1)) term = -term;
            if (!std::isfinite(term)) break;

            sum += term;

            // Convergence test: term below relative eps of sum for 3 consecutive iterations
            if (std::abs(term) <= std::abs(sum) * eps * T(2)) {
                if (++small_count >= 3) break;
            } else {
                small_count = 0;
            }
            prev_sum = sum;
            (void)prev_sum;
        }
        return sum;
    }

    // Asymptotic series: -Σ_{k=1}^N z^{-k} / Γ(β - αk)
    // The return value is this sum (separate from Wiman's exponential term).
    // Truncate once the term magnitude starts increasing (asymptotic series).
    template<IsNativeFloat T>
    [[nodiscard]] T mittag_leffler_asymp_decay_part(T alpha, T beta, T z,
                                                    int max_terms = 50) {
        const T inv_z = T(1) / z;
        T zk = inv_z;  // z^{-1}
        T sum = T(0);
        T prev_abs = std::numeric_limits<T>::infinity();

        for (int k = 1; k < max_terms; ++k) {
            const T inv_gamma = reciprocal_gamma_safe(beta - alpha * T(k));
            // The pole terms (β - αk a non-positive integer) are zeroed out, which
            // is a welcome case where the divergent terms vanish in the asymptotic series
            if (inv_gamma == T(0)) {
                zk *= inv_z;
                continue;
            }
            T term = zk * inv_gamma;
            T abs_term = std::abs(term);
            if (abs_term > prev_abs) break;  // divergence begins
            sum += term;
            prev_abs = abs_term;
            if (abs_term < std::abs(sum) * std::numeric_limits<T>::epsilon()) break;
            zk *= inv_z;
        }
        // It is -Σ, so return with the sign flipped
        return -sum;
    }

    // Wiman asymptotics (z > 0 only): (1/α) z^{(1-β)/α} exp(z^{1/α}) + decaying part
    template<IsNativeFloat T>
    [[nodiscard]] T mittag_leffler_asymp_positive(T alpha, T beta, T z,
                                                  int max_terms = 50) {
        const T inv_alpha = T(1) / alpha;
        const T z_pow_inv_alpha = std::pow(z, inv_alpha);
        const T pre = std::pow(z, (T(1) - beta) * inv_alpha);
        const T exp_term = inv_alpha * pre * std::exp(z_pow_inv_alpha);
        const T decay = mittag_leffler_asymp_decay_part(alpha, beta, z, max_terms);
        return exp_term + decay;
    }

    // |z| switching threshold — upper bound for using Taylor. Since |z|^{1/α} blows up for small α,
    // a threshold keeping |z|^{1/α} ≤ ~30 is adopted.
    // For negative z, alternating cancellation in intermediate terms loses double precision, so the threshold is shrunk to a bit over half.
    template<IsNativeFloat T>
    [[nodiscard]] T mittag_leffler_taylor_cutoff(T alpha, bool z_negative) {
        const T base = std::pow(T(30), alpha);
        return z_negative ? base * T(0.4) : base;
    }

} // namespace detail

// ================================================================
// Two-parameter Mittag-Leffler function E_{α,β}(z)
// ================================================================

/**
 * @brief E_{α,β}(z) = Σ_{k=0}^∞ z^k / Γ(αk + β)
 * @param alpha α (0 < α ≤ 2)
 * @param beta  β (real. At non-positive integers, treated as 1/Γ = 0 at the pole of Γ)
 * @param z     argument (real)
 */
template<IsNativeFloat T>
[[nodiscard]] T mittagLeffler(T alpha, T beta, T z) {
    if (std::isnan(alpha) || std::isnan(beta) || std::isnan(z)) {
        return std::numeric_limits<T>::quiet_NaN();
    }
    if (alpha <= T(0) || alpha > T(2)) {
        throw std::domain_error("mittagLeffler: alpha must be in (0, 2]");
    }

    // Special case: α = 1
    if (std::abs(alpha - T(1)) < T(1e-14)) {
        if (std::abs(beta - T(1)) < T(1e-14)) return std::exp(z);
        if (std::abs(beta - T(2)) < T(1e-14)) {
            if (z == T(0)) return T(1);
            return std::expm1(z) / z;
        }
        // General α=1 is left to Taylor or an incomplete-gamma-like representation
    }

    // Special case: α = 2, z ≥ 0
    if (std::abs(alpha - T(2)) < T(1e-14) && z >= T(0)) {
        const T r = std::sqrt(z);
        if (std::abs(beta - T(1)) < T(1e-14)) return std::cosh(r);
        if (std::abs(beta - T(2)) < T(1e-14)) {
            if (z == T(0)) return T(1);
            return std::sinh(r) / r;
        }
    }

    if (z == T(0)) return detail::reciprocal_gamma_safe(beta);

    const T abs_z = std::abs(z);
    const T cutoff = detail::mittag_leffler_taylor_cutoff(alpha, z < T(0));
    if (abs_z <= cutoff) {
        return detail::mittag_leffler_taylor(alpha, beta, z);
    }
    if (z > T(0)) {
        return detail::mittag_leffler_asymp_positive(alpha, beta, z);
    }
    return detail::mittag_leffler_asymp_decay_part(alpha, beta, z);
}

// ================================================================
// Single-parameter Mittag-Leffler function E_α(z) = E_{α,1}(z)
// ================================================================

template<IsNativeFloat T>
[[nodiscard]] T mittagLeffler(T alpha, T z) {
    return mittagLeffler(alpha, T(1), z);
}

} // namespace special
} // namespace sangi

#endif // SANGI_SPECIAL_MITTAG_LEFFLER_HPP
