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

// legendre.hpp
// Legendre polynomial / associated Legendre function / spherical harmonic / Jacobi polynomial templates
//
// Functions provided:
//   legendreP(n, x)           — Legendre polynomial P_n(x)
//   assocLegendreP(n, m, x)   — associated Legendre function P_n^m(x) (without Condon-Shortley)
//   sphLegendre(n, m, theta)  — normalized spherical harmonic Y_n^m(θ)
//   jacobiP(n, alpha, beta, x) — Jacobi polynomial P_n^{(α,β)}(x)
//
// The Jacobi polynomial generalizes Legendre, Chebyshev, and Gegenbauer:
//   P_n^{(0,0)}(x) = P_n(x) (Legendre)
//   P_n^{(-1/2,-1/2)}(x) ∝ T_n(x) (Chebyshev, first kind)
//   P_n^{(1/2,1/2)}(x) ∝ U_n(x) (Chebyshev, second kind)
//   P_n^{(α,α)}(x) ∝ C_n^{α+1/2}(x) (Gegenbauer/ultraspherical)
//
// Supported types:
//   float, double, long double — Bonnet recurrence / three-term recurrence
//   Float                      — delegated to the sangi:: implementation (except sphLegendre)
//   Complex<double>            — Bonnet recurrence (complex argument)
//   Complex<Float>             — arbitrary-precision Bonnet recurrence

#ifndef SANGI_SPECIAL_LEGENDRE_HPP
#define SANGI_SPECIAL_LEGENDRE_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 <numbers>
#include <limits>

namespace sangi {
namespace special {

// ================================================================
// Legendre polynomial P_n(x)
// ================================================================
// Bonnet recurrence: (n+1)P_{n+1} = (2n+1)·x·P_n - n·P_{n-1}
// P_0 = 1, P_1 = x

template<IsNativeFloat T>
[[nodiscard]] T legendreP(int n, T x) {
    if (std::isnan(x)) return std::numeric_limits<T>::quiet_NaN();
    if (n < 0) return std::numeric_limits<T>::quiet_NaN();
    if (n == 0) return T(1);
    if (n == 1) return x;

    T p_prev = T(1);   // P_0
    T p_curr = x;       // P_1

    for (int k = 1; k < n; k++) {
        T p_next = (T(2 * k + 1) * x * p_curr - T(k) * p_prev) / T(k + 1);
        p_prev = p_curr;
        p_curr = p_next;
    }

    return p_curr;
}

// ================================================================
// Associated Legendre function P_n^m(x)
// ================================================================
// Without the Condon-Shortley phase factor (matches C++17 std::assoc_legendre)
//
// 1. P_m^m = (2m-1)!! · (1-x²)^{m/2}
// 2. P_{m+1}^m = x·(2m+1)·P_m^m
// 3. (n-m+1)P_{n+1}^m = (2n+1)·x·P_n^m - (n+m)·P_{n-1}^m

template<IsNativeFloat T>
[[nodiscard]] T assocLegendreP(int n, int m, T x) {
    if (std::isnan(x)) return std::numeric_limits<T>::quiet_NaN();
    if (n < 0 || m < 0) return std::numeric_limits<T>::quiet_NaN();
    if (m > n) return T(0);

    // m = 0 → P_n(x)
    if (m == 0) return legendreP(n, x);

    // P_m^m = (2m-1)!! · (1-x²)^{m/2}
    T sin2 = T(1) - x * x;
    if (sin2 < T(0)) sin2 = T(0);  // guard against numerical error
    T sin_factor = std::sqrt(sin2);

    T pmm = T(1);
    for (int i = 1; i <= m; i++) {
        pmm *= T(2 * i - 1) * sin_factor;
    }

    if (n == m) return pmm;

    // P_{m+1}^m = x·(2m+1)·P_m^m
    T pm1m = x * T(2 * m + 1) * pmm;
    if (n == m + 1) return pm1m;

    // Recurrence: (k-m+1)P_{k+1}^m = (2k+1)·x·P_k^m - (k+m)·P_{k-1}^m
    T p_prev = pmm;
    T p_curr = pm1m;

    for (int k = m + 1; k < n; k++) {
        T p_next = (T(2 * k + 1) * x * p_curr - T(k + m) * p_prev) / T(k - m + 1);
        p_prev = p_curr;
        p_curr = p_next;
    }

    return p_curr;
}

// ================================================================
// Normalized spherical harmonic Y_n^m(θ) (real part)
// ================================================================
// Y_n^m(θ) = √((2n+1)/(4π) · (n-m)!/(n+m)!) · P_n^m(cos θ)
// Only m ≥ 0 supported (real part only)

template<IsNativeFloat T>
[[nodiscard]] T sphLegendre(int n, int m, T theta) {
    if (std::isnan(theta)) return std::numeric_limits<T>::quiet_NaN();
    if (n < 0 || m < 0 || m > n) return std::numeric_limits<T>::quiet_NaN();

    // Normalization coefficient: √((2n+1)/(4π) · (n-m)!/(n+m)!)
    // (n-m)!/(n+m)! = 1/((n-m+1)(n-m+2)···(n+m)) computed in log space
    T log_coeff = T(0);
    for (int i = n - m + 1; i <= n + m; i++) {
        log_coeff -= std::log(T(i));
    }
    log_coeff += std::log(T(2 * n + 1) / (T(4) * std::numbers::pi_v<T>));

    T coeff = std::exp(T(0.5) * log_coeff);

    T cos_theta = std::cos(theta);
    T plm = assocLegendreP(n, m, cos_theta);

    return coeff * plm;
}

// Float type: use sangi::legendreP, sangi::assocLegendreP directly.

// ================================================================
// Complex<double> Legendre polynomial P_n(z)
// ================================================================
// The Bonnet recurrence holds as-is for complex arguments.

[[nodiscard]] inline Complex<double> legendreP(int n, const Complex<double>& z) {
    using C = Complex<double>;
    if (n < 0) return C(std::numeric_limits<double>::quiet_NaN());
    if (n == 0) return C(1.0);
    if (n == 1) return z;

    C p_prev(1.0);
    C p_curr = z;

    for (int k = 1; k < n; k++) {
        C p_next = (C(2.0 * k + 1.0) * z * p_curr - C(double(k)) * p_prev) / C(double(k + 1));
        p_prev = p_curr;
        p_curr = p_next;
    }

    return p_curr;
}

// ================================================================
// Complex<double> associated Legendre function P_n^m(z)
// ================================================================
// In the complex plane (1-z²) may become negative, so the real-valued
// version's sin2 < 0 guard is unnecessary. sqrt(1-z²) uses Complex sqrt.

[[nodiscard]] inline Complex<double> assocLegendreP(int n, int m, const Complex<double>& z) {
    using C = Complex<double>;
    if (n < 0 || m < 0) return C(std::numeric_limits<double>::quiet_NaN());
    if (m > n) return C(0.0);
    if (m == 0) return legendreP(n, z);

    // sqrt(1 - z²)
    C sin_factor = sangi::sqrt(C(1.0) - z * z);

    // P_m^m = (2m-1)!! · (1-z²)^{m/2}
    C pmm(1.0);
    for (int i = 1; i <= m; i++) {
        pmm = pmm * C(2.0 * i - 1.0) * sin_factor;
    }

    if (n == m) return pmm;

    // P_{m+1}^m = z·(2m+1)·P_m^m
    C pm1m = z * C(2.0 * m + 1.0) * pmm;
    if (n == m + 1) return pm1m;

    // Recurrence
    C p_prev = pmm;
    C p_curr = pm1m;

    for (int k = m + 1; k < n; k++) {
        C p_next = (C(2.0 * k + 1.0) * z * p_curr - C(double(k + m)) * p_prev) / C(double(k - m + 1));
        p_prev = p_curr;
        p_curr = p_next;
    }

    return p_curr;
}

// ================================================================
// Complex<Float> Legendre polynomial P_n(z)
// ================================================================

[[nodiscard]] inline Complex<Float> legendreP(int n, const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    if (n < 0) return C(Float::nan());
    if (n == 0) { Float one(1); one.setPrecision(precision); return C(one); }
    if (n == 1) {
        C result = z;
        result.re.setPrecision(precision);
        result.im.setPrecision(precision);
        return result;
    }

    int wp = precision + 15;
    // PrecisionGuard removed: zw's setResultPrecision(wp) propagates req=wp.
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    C p_prev(Float(1));
    C p_curr = zw;

    for (int k = 1; k < n; k++) {
        C p_next = (C(Float(2 * k + 1)) * zw * p_curr - C(Float(k)) * p_prev) / C(Float(k + 1));
        p_prev = p_curr;
        p_curr = p_next;
    }

    p_curr.re.setPrecision(precision);
    p_curr.im.setPrecision(precision);
    return p_curr;
}

// ================================================================
// Complex<Float> associated Legendre function P_n^m(z)
// ================================================================

[[nodiscard]] inline Complex<Float> assocLegendreP(int n, int m, const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    if (n < 0 || m < 0) return C(Float::nan());
    if (m > n) { Float zero(0); zero.setPrecision(precision); return C(zero); }
    if (m == 0) return legendreP(n, z, precision);

    int wp = precision + 20;
    // PrecisionGuard removed: zw's setResultPrecision(wp) propagates req=wp.
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    // sqrt(1 - z²)
    C sin_factor = sangi::sqrt(C(Float(1)) - zw * zw);

    // P_m^m = (2m-1)!! · (1-z²)^{m/2}
    C pmm(Float(1));
    for (int i = 1; i <= m; i++) {
        pmm = pmm * C(Float(2 * i - 1)) * sin_factor;
    }

    if (n == m) {
        pmm.re.setPrecision(precision);
        pmm.im.setPrecision(precision);
        return pmm;
    }

    // P_{m+1}^m = z·(2m+1)·P_m^m
    C pm1m = zw * C(Float(2 * m + 1)) * pmm;
    if (n == m + 1) {
        pm1m.re.setPrecision(precision);
        pm1m.im.setPrecision(precision);
        return pm1m;
    }

    // Recurrence
    C p_prev = pmm;
    C p_curr = pm1m;

    for (int k = m + 1; k < n; k++) {
        C p_next = (C(Float(2 * k + 1)) * zw * p_curr - C(Float(k + m)) * p_prev) / C(Float(k - m + 1));
        p_prev = p_curr;
        p_curr = p_next;
    }

    p_curr.re.setPrecision(precision);
    p_curr.im.setPrecision(precision);
    return p_curr;
}

// ================================================================
// Jacobi polynomial P_n^{(α,β)}(x)
// ================================================================
// DLMF §18.9.2 three-term recurrence:
//   P_0^{(α,β)}(x) = 1
//   P_1^{(α,β)}(x) = (α - β)/2 + (α + β + 2)·x/2
//   2n(n+α+β)(2n+α+β-2) P_n = (2n+α+β-1)[(2n+α+β)(2n+α+β-2)x + α²-β²] P_{n-1}
//                               - 2(n+α-1)(n+β-1)(2n+α+β) P_{n-2}
//
// Requires α, β > -1 (orthogonality condition)

template<IsNativeFloat T>
[[nodiscard]] T jacobiP(int n, T alpha, T beta, T x) {
    if (std::isnan(x) || std::isnan(alpha) || std::isnan(beta))
        return std::numeric_limits<T>::quiet_NaN();
    if (n < 0) return std::numeric_limits<T>::quiet_NaN();
    if (alpha <= T(-1) || beta <= T(-1)) return std::numeric_limits<T>::quiet_NaN();

    if (n == 0) return T(1);

    T ab = alpha + beta;
    T p_prev = T(1);  // P_0
    T p_curr = (alpha - beta) / T(2) + (ab + T(2)) * x / T(2);  // P_1

    if (n == 1) return p_curr;

    for (int k = 2; k <= n; k++) {
        T kf = T(k);
        T c = T(2) * kf + ab;  // 2k + α + β

        // Coefficients (DLMF 18.9.2)
        T a1 = T(2) * kf * (kf + ab) * (c - T(2));
        T a2 = (c - T(1)) * (alpha * alpha - beta * beta);
        T a3 = (c - T(2)) * (c - T(1)) * c;
        T a4 = T(2) * (kf + alpha - T(1)) * (kf + beta - T(1)) * c;

        T p_next = ((a2 + a3 * x) * p_curr - a4 * p_prev) / a1;
        p_prev = p_curr;
        p_curr = p_next;
    }

    return p_curr;
}

// ================================================================
// Complex<double> Jacobi polynomial P_n^{(α,β)}(z)
// ================================================================

[[nodiscard]] inline Complex<double> jacobiP(int n, double alpha, double beta,
                                              const Complex<double>& z) {
    using C = Complex<double>;
    if (n < 0 || alpha <= -1.0 || beta <= -1.0)
        return C(std::numeric_limits<double>::quiet_NaN());
    if (n == 0) return C(1.0);

    double ab = alpha + beta;
    C p_prev(1.0);
    C p_curr = C((alpha - beta) / 2.0) + C((ab + 2.0) / 2.0) * z;

    if (n == 1) return p_curr;

    for (int k = 2; k <= n; k++) {
        double kf = double(k);
        double c = 2.0 * kf + ab;

        double a1 = 2.0 * kf * (kf + ab) * (c - 2.0);
        double a2 = (c - 1.0) * (alpha * alpha - beta * beta);
        double a3 = (c - 2.0) * (c - 1.0) * c;
        double a4 = 2.0 * (kf + alpha - 1.0) * (kf + beta - 1.0) * c;

        C p_next = (C(a2) + C(a3) * z) * p_curr - C(a4) * p_prev;
        p_next = p_next / C(a1);
        p_prev = p_curr;
        p_curr = p_next;
    }

    return p_curr;
}

// ================================================================
// Complex<Float> Jacobi polynomial P_n^{(α,β)}(z)
// ================================================================

[[nodiscard]] inline Complex<Float> jacobiP(int n, const Float& alpha, const Float& beta,
                                             const Complex<Float>& z, int precision) {
    using C = Complex<Float>;
    if (n < 0 || alpha <= Float(-1) || beta <= Float(-1))
        return C(Float::nan());
    if (n == 0) {
        Float one(1); one.setPrecision(precision);
        return C(one);
    }

    int wp = precision + 20;
    // PrecisionGuard removed: aw/bw/zw's setResultPrecision(wp) propagates req=wp.

    Float aw = alpha; aw.setResultPrecision(wp);
    Float bw = beta;  bw.setResultPrecision(wp);
    C zw = z;
    zw.re.setResultPrecision(wp);
    zw.im.setResultPrecision(wp);

    Float ab = aw + bw;
    C p_prev(Float(1));
    C p_curr = C(ldexp(aw - bw, -1)) + C(ldexp(ab + Float(2), -1)) * zw;

    if (n == 1) {
        p_curr.re.setPrecision(precision);
        p_curr.im.setPrecision(precision);
        return p_curr;
    }

    for (int k = 2; k <= n; k++) {
        Float kf(k);
        Float c = mulScalarF(kf, uint64_t(2)) + ab;

        Float a1 = mulScalarF(kf, uint64_t(2)) * (kf + ab) * (c - Float(2));
        Float a2 = (c - Float(1)) * (aw * aw - bw * bw);
        Float a3 = (c - Float(2)) * (c - Float(1)) * c;
        Float a4 = mulScalarF(kf + aw - Float(1), uint64_t(2)) * (kf + bw - Float(1)) * c;

        C p_next = (C(a2) + C(a3) * zw) * p_curr - C(a4) * p_prev;
        p_next = p_next / C(a1);
        p_prev = p_curr;
        p_curr = p_next;
    }

    p_curr.re.setPrecision(precision);
    p_curr.im.setPrecision(precision);
    return p_curr;
}

[[nodiscard]] inline Complex<Float> jacobiP(int n, const Float& alpha, const Float& beta,
                                             const Complex<Float>& z) {
    return jacobiP(n, alpha, beta, z, Float::defaultPrecision());
}

// ================================================================
// Gegenbauer (ultraspherical) polynomial C_n^{λ}(x)
// ================================================================
// C_n^{λ}(x) is a specialization of the Jacobi polynomial:
//   C_n^{λ}(x) = Γ(2λ+n)Γ(λ+1/2) / (Γ(2λ)Γ(λ+n+1/2)) · P_n^{(λ-1/2, λ-1/2)}(x)
//
// Three-term recurrence (DLMF 18.9.1):
//   C_0^{λ}(x) = 1
//   C_1^{λ}(x) = 2λx
//   n·C_n^{λ}(x) = 2(n+λ-1)·x·C_{n-1}^{λ}(x) - (n+2λ-2)·C_{n-2}^{λ}(x)

template<IsNativeFloat T>
[[nodiscard]] T gegenbauerC(int n, T lambda, T x) {
    if (std::isnan(x) || std::isnan(lambda))
        return std::numeric_limits<T>::quiet_NaN();
    if (n < 0) return std::numeric_limits<T>::quiet_NaN();
    if (n == 0) return T(1);

    T p_prev = T(1);
    T p_curr = T(2) * lambda * x;

    if (n == 1) return p_curr;

    for (int k = 2; k <= n; k++) {
        T p_next = (T(2) * (T(k) + lambda - T(1)) * x * p_curr
                    - (T(k) + T(2) * lambda - T(2)) * p_prev) / T(k);
        p_prev = p_curr;
        p_curr = p_next;
    }

    return p_curr;
}

} // namespace special
} // namespace sangi

#endif // SANGI_SPECIAL_LEGENDRE_HPP