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

// airy.hpp
// Templates for the Airy functions and their derivatives
//
// Provided functions:
//   airyAi(x), airyBi(x)           — Airy functions of the first/second kind
//   airyAiPrime(x), airyBiPrime(x) — derivatives of the Airy functions
//
// Supported types:
//   float, double, long double — computed directly from the Maclaurin series
//   Float                      — delegated to the sangi:: implementation
//   Complex<double>            — Maclaurin series (entire function)
//   Complex<Float>             — arbitrary-precision Maclaurin series

#ifndef SANGI_SPECIAL_AIRY_HPP
#define SANGI_SPECIAL_AIRY_HPP

#include <math/special/special_common.hpp>
#include <math/special/gamma.hpp>
#include <math/core/mp/Float.hpp>
#include <math/core/Complex.hpp>
#include <cmath>
#include <limits>

namespace sangi {
namespace special {

// ================================================================
// Airy constants (double precision)
// ================================================================

namespace detail {
    // Ai(0) = 1 / (3^{2/3} · Γ(2/3))
    inline constexpr double kAi0  =  0.35502805388781723926;
    // Ai'(0) = -1 / (3^{1/3} · Γ(1/3))
    inline constexpr double kAip0 = -0.25881940379280679841;
    // Bi(0) = 1 / (3^{1/6} · Γ(2/3))
    inline constexpr double kBi0  =  0.61492662744600073515;
    // Bi'(0) = 3^{1/6} / Γ(1/3)
    inline constexpr double kBip0 =  0.44828835735382635789;
} // namespace detail

// ================================================================
// Shared Maclaurin-series computation
// ================================================================
//
// Two fundamental solutions of the Airy equation y'' = xy:
//   f(x) = 1 + x³/(2·3) + x⁶/(2·3·5·6) + ...
//   g(x) = x + x⁴/(3·4) + x⁷/(3·4·6·7) + ...
//
// Recurrence:
//   f_term_{k+1} = f_term_k · x³ / ((3k+2)(3k+3))
//   g_term_{k+1} = g_term_k · x³ / ((3k+3)(3k+4))
//
// Ai(x) = Ai(0)·f(x) + Ai'(0)·g(x)
// Bi(x) = Bi(0)·f(x) + Bi'(0)·g(x)
//
// Derivatives:
//   f'(x) = Σ 3(k+1)·f_term_{k+1}/x  (k=0,1,2,...)
//   g'(x) = 1 + Σ (3k+4)·g_term_{k+1}/x  (k=0,1,2,...)
//
// Ai'(x) = Ai(0)·f'(x) + Ai'(0)·g'(x)
// Bi'(x) = Bi(0)·f'(x) + Bi'(0)·g'(x)

// ================================================================
// Airy Ai(x)
// ================================================================

template<IsNativeFloat T>
[[nodiscard]] T airyAi(T x) {
    if (std::isnan(x)) return std::numeric_limits<T>::quiet_NaN();
    if (std::isinf(x)) return T(0);  // Ai(±∞) → 0

    const T ai0  = static_cast<T>(detail::kAi0);
    const T aip0 = static_cast<T>(detail::kAip0);

    if (x == T(0)) return ai0;

    const T x3 = x * x * x;
    T f_sum = T(1);
    T g_sum = x;
    T f_term = T(1);
    T g_term = x;

    for (int k = 0; k < 500; k++) {
        f_term *= x3 / (T(3*k+2) * T(3*k+3));
        g_term *= x3 / (T(3*k+3) * T(3*k+4));
        f_sum += f_term;
        g_sum += g_term;

        if (k >= 3 && std::abs(f_term) + std::abs(g_term) <
            std::numeric_limits<T>::epsilon() * (std::abs(f_sum) + std::abs(g_sum)))
            break;
    }

    return ai0 * f_sum + aip0 * g_sum;
}

// Float type: use sangi::airyAi(x, precision) directly.

// ================================================================
// Airy Bi(x)
// ================================================================

template<IsNativeFloat T>
[[nodiscard]] T airyBi(T x) {
    if (std::isnan(x)) return std::numeric_limits<T>::quiet_NaN();
    if (std::isinf(x)) {
        if (x < T(0)) return T(0);   // Bi(-∞) → 0
        return std::numeric_limits<T>::infinity();  // Bi(+∞) → +∞
    }

    const T bi0  = static_cast<T>(detail::kBi0);
    const T bip0 = static_cast<T>(detail::kBip0);

    if (x == T(0)) return bi0;

    const T x3 = x * x * x;
    T f_sum = T(1);
    T g_sum = x;
    T f_term = T(1);
    T g_term = x;

    for (int k = 0; k < 500; k++) {
        f_term *= x3 / (T(3*k+2) * T(3*k+3));
        g_term *= x3 / (T(3*k+3) * T(3*k+4));
        f_sum += f_term;
        g_sum += g_term;

        if (k >= 3 && std::abs(f_term) + std::abs(g_term) <
            std::numeric_limits<T>::epsilon() * (std::abs(f_sum) + std::abs(g_sum)))
            break;
    }

    return bi0 * f_sum + bip0 * g_sum;
}

// Float type: use sangi::airyBi(x, precision) directly.

// ================================================================
// Airy Ai'(x)
// ================================================================

template<IsNativeFloat T>
[[nodiscard]] T airyAiPrime(T x) {
    if (std::isnan(x)) return std::numeric_limits<T>::quiet_NaN();
    if (std::isinf(x)) return T(0);  // Ai'(±∞) → 0

    const T ai0  = static_cast<T>(detail::kAi0);
    const T aip0 = static_cast<T>(detail::kAip0);

    // Ai'(0) = aip0  (f'(0)=0, g'(0)=1)
    if (x == T(0)) return aip0;

    const T x3 = x * x * x;
    const T inv_x = T(1) / x;
    T f_term = T(1);
    T g_term = x;
    T fp_sum = T(0);  // f'(0) = 0
    T gp_sum = T(1);  // g'(0) = 1

    for (int k = 0; k < 500; k++) {
        f_term *= x3 / (T(3*k+2) * T(3*k+3));
        g_term *= x3 / (T(3*k+3) * T(3*k+4));

        T fp_contrib = T(3*(k+1)) * f_term * inv_x;
        T gp_contrib = T(3*k+4) * g_term * inv_x;
        fp_sum += fp_contrib;
        gp_sum += gp_contrib;

        if (k >= 3 && std::abs(fp_contrib) + std::abs(gp_contrib) <
            std::numeric_limits<T>::epsilon() * (std::abs(fp_sum) + std::abs(gp_sum)))
            break;
    }

    return ai0 * fp_sum + aip0 * gp_sum;
}

// Float type: use sangi::airyAiPrime(x, precision) directly.

// ================================================================
// Airy Bi'(x)
// ================================================================

template<IsNativeFloat T>
[[nodiscard]] T airyBiPrime(T x) {
    if (std::isnan(x)) return std::numeric_limits<T>::quiet_NaN();
    if (std::isinf(x)) {
        if (x < T(0)) return T(0);   // Bi'(-∞) → 0
        return std::numeric_limits<T>::infinity();  // Bi'(+∞) → +∞
    }

    const T bi0  = static_cast<T>(detail::kBi0);
    const T bip0 = static_cast<T>(detail::kBip0);

    if (x == T(0)) return bip0;

    const T x3 = x * x * x;
    const T inv_x = T(1) / x;
    T f_term = T(1);
    T g_term = x;
    T fp_sum = T(0);
    T gp_sum = T(1);

    for (int k = 0; k < 500; k++) {
        f_term *= x3 / (T(3*k+2) * T(3*k+3));
        g_term *= x3 / (T(3*k+3) * T(3*k+4));

        T fp_contrib = T(3*(k+1)) * f_term * inv_x;
        T gp_contrib = T(3*k+4) * g_term * inv_x;
        fp_sum += fp_contrib;
        gp_sum += gp_contrib;

        if (k >= 3 && std::abs(fp_contrib) + std::abs(gp_contrib) <
            std::numeric_limits<T>::epsilon() * (std::abs(fp_sum) + std::abs(gp_sum)))
            break;
    }

    return bi0 * fp_sum + bip0 * gp_sum;
}

// Float type: use sangi::airyBiPrime(x, precision) directly.

// ================================================================
// Complex Airy — shared Maclaurin-series template
// ================================================================
//
// f(z), g(z) are entire functions, so the same recurrence as the real case can be used as is.
// The derivatives are likewise computed in the complex domain as f'(z), g'(z).

namespace detail {

/// Shared computation of Ai(z), Bi(z): returns f_sum, g_sum
template<typename R>
void airy_fg_complex(const Complex<R>& z, R eps, int max_iter,
                     Complex<R>& f_sum, Complex<R>& g_sum) {
    using C = Complex<R>;
    C z3 = z * z * z;
    f_sum = C(R(1));
    g_sum = z;
    C f_term(R(1));
    C g_term = z;

    for (int k = 0; k < max_iter; k++) {
        f_term = f_term * z3 / R(static_cast<int64_t>(3*k+2) * (3*k+3));
        g_term = g_term * z3 / R(static_cast<int64_t>(3*k+3) * (3*k+4));
        f_sum = f_sum + f_term;
        g_sum = g_sum + g_term;

        if (k >= 3 && sangi::abs(f_term) + sangi::abs(g_term) <
            eps * (sangi::abs(f_sum) + sangi::abs(g_sum)))
            break;
    }
}

/// Shared computation of Ai'(z), Bi'(z): returns fp_sum, gp_sum
template<typename R>
void airy_fg_prime_complex(const Complex<R>& z, R eps, int max_iter,
                           Complex<R>& fp_sum, Complex<R>& gp_sum) {
    using C = Complex<R>;
    C z3 = z * z * z;
    C inv_z = C(R(1)) / z;
    C f_term(R(1));
    C g_term = z;
    fp_sum = C(R(0));
    gp_sum = C(R(1));

    for (int k = 0; k < max_iter; k++) {
        f_term = f_term * z3 / R(static_cast<int64_t>(3*k+2) * (3*k+3));
        g_term = g_term * z3 / R(static_cast<int64_t>(3*k+3) * (3*k+4));

        C fp_contrib = C(R(3*(k+1))) * f_term * inv_z;
        C gp_contrib = C(R(3*k+4)) * g_term * inv_z;
        fp_sum = fp_sum + fp_contrib;
        gp_sum = gp_sum + gp_contrib;

        if (k >= 3 && sangi::abs(fp_contrib) + sangi::abs(gp_contrib) <
            eps * (sangi::abs(fp_sum) + sangi::abs(gp_sum)))
            break;
    }
}

} // namespace detail

// ================================================================
// Complex<double> overloads
// ================================================================

[[nodiscard]] inline Complex<double> airyAi(const Complex<double>& z) {
    using C = Complex<double>;
    C f_sum, g_sum;
    detail::airy_fg_complex<double>(z, std::numeric_limits<double>::epsilon(), 500,
                                    f_sum, g_sum);
    return C(detail::kAi0) * f_sum + C(detail::kAip0) * g_sum;
}

[[nodiscard]] inline Complex<double> airyBi(const Complex<double>& z) {
    using C = Complex<double>;
    C f_sum, g_sum;
    detail::airy_fg_complex<double>(z, std::numeric_limits<double>::epsilon(), 500,
                                    f_sum, g_sum);
    return C(detail::kBi0) * f_sum + C(detail::kBip0) * g_sum;
}

[[nodiscard]] inline Complex<double> airyAiPrime(const Complex<double>& z) {
    using C = Complex<double>;
    if (sangi::abs(z) == 0.0) return C(detail::kAip0);
    C fp_sum, gp_sum;
    detail::airy_fg_prime_complex<double>(z, std::numeric_limits<double>::epsilon(), 500,
                                          fp_sum, gp_sum);
    return C(detail::kAi0) * fp_sum + C(detail::kAip0) * gp_sum;
}

[[nodiscard]] inline Complex<double> airyBiPrime(const Complex<double>& z) {
    using C = Complex<double>;
    if (sangi::abs(z) == 0.0) return C(detail::kBip0);
    C fp_sum, gp_sum;
    detail::airy_fg_prime_complex<double>(z, std::numeric_limits<double>::epsilon(), 500,
                                          fp_sum, gp_sum);
    return C(detail::kBi0) * fp_sum + C(detail::kBip0) * gp_sum;
}

// ================================================================
// Complex<Float> overloads
// ================================================================

[[nodiscard]] inline Complex<Float> airyAi(const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    int wp = precision + 30;
    Float::PrecisionScope _ps(wp);   // compute the in-function exact÷exact / pow(exact) at wp digits
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    // Ai(0) = 1 / (3^{2/3} · Γ(2/3)),  Ai'(0) = -1 / (3^{1/3} · Γ(1/3))
    Float three(3);   // no setResultPrecision needed: PrecisionScope promotes Float(1)/three to wp
    Float cbrt3 = sangi::pow(three, Float(1) / three);  // 3^{1/3}
    Float cbrt3_sq = cbrt3 * cbrt3;                     // 3^{2/3}
    Float g23 = sangi::gamma(Float(2) / three, wp);   // Γ(2/3)
    Float g13 = sangi::gamma(Float(1) / three, wp);   // Γ(1/3)
    Float ai0  = Float(1) / (cbrt3_sq * g23);
    Float aip0 = -Float(1) / (cbrt3 * g13);

    C f_sum, g_sum;
    detail::airy_fg_complex<Float>(zw, Float::epsilon(wp), 10 * wp,
                                   f_sum, g_sum);

    C result = C(ai0) * f_sum + C(aip0) * g_sum;
    result.re.setPrecision(precision);
    result.im.setPrecision(precision);
    return result;
}

[[nodiscard]] inline Complex<Float> airyBi(const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    int wp = precision + 30;
    Float::PrecisionScope _ps(wp);   // compute the in-function exact÷exact / pow(exact) at wp digits
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    Float three(3);   // no setResultPrecision needed: PrecisionScope promotes Float(1)/three and Float(1)/Float(6) to wp
    Float sixth = sangi::pow(three, Float(1) / Float(6));   // 3^{1/6} (PrecisionScope prevents the precision from being capped)
    Float g23 = sangi::gamma(Float(2) / three, wp);
    Float g13 = sangi::gamma(Float(1) / three, wp);
    Float bi0  = Float(1) / (sixth * g23);
    Float bip0 = sixth / g13;

    C f_sum, g_sum;
    detail::airy_fg_complex<Float>(zw, Float::epsilon(wp), 10 * wp,
                                   f_sum, g_sum);

    C result = C(bi0) * f_sum + C(bip0) * g_sum;
    result.re.setPrecision(precision);
    result.im.setPrecision(precision);
    return result;
}

[[nodiscard]] inline Complex<Float> airyAiPrime(const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    int wp = precision + 30;
    Float::PrecisionScope _ps(wp);   // compute the in-function exact÷exact / pow(exact) at wp digits
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    Float three(3);   // no setResultPrecision needed: PrecisionScope promotes Float(1)/three to wp
    Float cbrt3 = sangi::pow(three, Float(1) / three);
    Float cbrt3_sq = cbrt3 * cbrt3;
    Float g23 = sangi::gamma(Float(2) / three, wp);
    Float g13 = sangi::gamma(Float(1) / three, wp);
    Float ai0  = Float(1) / (cbrt3_sq * g23);
    Float aip0 = -Float(1) / (cbrt3 * g13);

    if (sangi::abs(zw).isZero()) {
        C result(aip0);
        result.re.setPrecision(precision);
        result.im.setPrecision(precision);
        return result;
    }

    C fp_sum, gp_sum;
    detail::airy_fg_prime_complex<Float>(zw, Float::epsilon(wp), 10 * wp,
                                         fp_sum, gp_sum);

    C result = C(ai0) * fp_sum + C(aip0) * gp_sum;
    result.re.setPrecision(precision);
    result.im.setPrecision(precision);
    return result;
}

[[nodiscard]] inline Complex<Float> airyBiPrime(const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    int wp = precision + 30;
    Float::PrecisionScope _ps(wp);   // compute the in-function exact÷exact / pow(exact) at wp digits
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    Float three(3);   // no setResultPrecision needed: PrecisionScope promotes Float(1)/three and Float(1)/Float(6) to wp
    Float sixth = sangi::pow(three, Float(1) / Float(6));   // 3^{1/6} (PrecisionScope prevents the precision from being capped)
    Float g23 = sangi::gamma(Float(2) / three, wp);
    Float g13 = sangi::gamma(Float(1) / three, wp);
    Float bi0  = Float(1) / (sixth * g23);
    Float bip0 = sixth / g13;

    if (sangi::abs(zw).isZero()) {
        C result(bip0);
        result.re.setPrecision(precision);
        result.im.setPrecision(precision);
        return result;
    }

    C fp_sum, gp_sum;
    detail::airy_fg_prime_complex<Float>(zw, Float::epsilon(wp), 10 * wp,
                                         fp_sum, gp_sum);

    C result = C(bi0) * fp_sum + C(bip0) * gp_sum;
    result.re.setPrecision(precision);
    result.im.setPrecision(precision);
    return result;
}

} // namespace special
} // namespace sangi

#endif // SANGI_SPECIAL_AIRY_HPP