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

// spheroidal.hpp
// spheroidal wave functions
//
// The angular functions S_{mn}(c, η) and eigenvalues λ_{mn}(c) obtained by separating
// the Helmholtz equation in prolate / oblate spheroidal coordinates. This drop is the
// first version in the stocking phase and provides only the native-double prolate /
// oblate angular functions (first kind) and eigenvalues. The radial functions
// R_{mn}(c, ξ) and the Float / Complex versions are future work.
//
// Functions provided (native double / float / long double):
//   spheroidalEigenvalueProlate(m, n, c)   λ_{mn}^{P}(c)
//   spheroidalEigenvalueOblate(m, n, c)    λ_{mn}^{O}(c)
//   spheroidalPS1(m, n, c, eta)            S_{mn}^{P,1}(c, η)  (prolate)
//   spheroidalOS1(m, n, c, eta)            S_{mn}^{O,1}(c, η)  (oblate)
//
// Differential equation (DLMF 30.2.1, prolate; sign of γ²=c² > 0):
//   d/dη [(1-η²) dS/dη] + [λ - c²η² - m²/(1-η²)] S = 0,  |η| ≤ 1
//
// Bouwkamp / Flammer expansion (DLMF 30.6, A&S 21.7):
//   S_{mn}(c, η) = Σ_{r=0,2,...}' d_r^{mn}(c) · P^m_{m+r}(η)
// (' sums over r of the same parity as (n-m))
//
// The coefficients d_r satisfy the 3-term recurrence
//   A_r d_{r+2} + (B_r - λ) d_r + C_r d_{r-2} = 0
// and λ_{mn} is the eigenvalue chosen so that this recurrence has a decaying solution as r→∞.
//
// Recurrence coefficients (Flammer 1957, A&S 21.7.4):
//   A_r = (2m+r+1)(2m+r+2) c² / [(2m+2r+3)(2m+2r+5)]
//   B_r = (m+r)(m+r+1) + [2(m+r)(m+r+1) - 2m² - 1] c² / [(2m+2r-1)(2m+2r+3)]
//   C_r = r(r-1) c² / [(2m+2r-3)(2m+2r-1)]
//
// Oblate reuses the same recurrence via the substitution c² → -c² (DLMF 30.2.1 note).
//
// Numerical method:
//   1. truncate r at size N=40, symmetric tridiagonalization
//      (replace the off-diagonal with √(A_{r}·C_{r+2}), branching real / imaginary by the sign of c²)
//   2. compute eigenvalues / eigenvectors via implicit-shift QL
//   3. the (n-m)/2-th eigenvalue in ascending order is λ_{mn}
//   4. Flammer-normalize the corresponding eigenvector (Σ d_r = 1) and extract d_r^{mn}
//
// Limitations (stocking phase initial):
//   - native float only (Float arbitrary-precision version is future work)
//   - targets 12-digit precision for |c| ≤ ~10 (large c requires increasing N)
//   - second and third kinds not implemented
//   - radial R_{mn}^{(j)}(c, ξ) not implemented
//   - resonance (region where n-m is large and λ becomes degenerate) not handled
//
// Uses: spheroidal scattering, generalization of Mathieu, 3D Helmholtz analytic solutions
//
// References:
//   - DLMF §30
//   - C. Flammer, "Spheroidal Wave Functions" (Stanford 1957)
//   - Abramowitz & Stegun §21
//   - SciPy: scipy.special.{pro_ang1, obl_ang1}

#ifndef SANGI_SPECIAL_SPHEROIDAL_HPP
#define SANGI_SPECIAL_SPHEROIDAL_HPP

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

namespace sangi {
namespace special {

namespace detail_spheroidal {

    // compute the coefficients of the 3-term recurrence
    // r is parity, parity+2, parity+4, ...
    template<IsNativeFloat T>
    struct RecCoeffs {
        T A_r;  // d_{r+2} coefficient
        T B_r;  // diagonal (main)
        T C_r;  // d_{r-2} coefficient
    };

    // c_squared = c² (prolate: > 0, oblate: < 0)
    template<IsNativeFloat T>
    [[nodiscard]] RecCoeffs<T> rec_coeffs(int m, int r, T c_squared) {
        T A = (T(2*m + r + 1) * T(2*m + r + 2) * c_squared)
            / (T(2*m + 2*r + 3) * T(2*m + 2*r + 5));
        T B = T(m + r) * T(m + r + 1)
            + (T(2 * (m+r) * (m+r+1) - 2*m*m - 1) * c_squared)
              / (T(2*m + 2*r - 1) * T(2*m + 2*r + 3));
        T C = T(r * (r - 1)) * c_squared
            / (T(2*m + 2*r - 3) * T(2*m + 2*r - 1));
        return {A, B, C};
    }

    // implicit-shift QL algorithm for symmetric tridiagonal matrices (Numerical Recipes "tqli" style)
    // d: diagonal (size N), e: sub-diagonal (size N, e[0] unused)
    // z: identity matrix on input, eigenvector columns on output
    // Return value: iteration count. Throws on failure.
    template<IsNativeFloat T>
    inline int tqli(std::vector<T>& d, std::vector<T>& e, std::vector<std::vector<T>>& z) {
        const int n = static_cast<int>(d.size());
        if (n == 1) return 0;

        // handle e with a 1-shift (e[i] is the sub-diagonal between d[i] and d[i+1])
        for (int i = 1; i < n; ++i) e[i - 1] = e[i];
        e[n - 1] = T(0);

        int total_iter = 0;
        const int max_iter = 60;

        for (int l = 0; l < n; ++l) {
            int iter = 0;
            int mz;
            do {
                for (mz = l; mz < n - 1; ++mz) {
                    T dd = std::abs(d[mz]) + std::abs(d[mz + 1]);
                    if (std::abs(e[mz]) <= std::numeric_limits<T>::epsilon() * dd) break;
                }
                if (mz != l) {
                    if (iter++ == max_iter) throw std::runtime_error("spheroidal tqli: convergence failure");
                    T g = (d[l + 1] - d[l]) / (T(2) * e[l]);
                    T r = std::hypot(g, T(1));
                    g = d[mz] - d[l] + e[l] / (g + std::copysign(r, g));
                    T s = T(1), c = T(1), p = T(0);
                    for (int i = mz - 1; i >= l; --i) {
                        T f = s * e[i];
                        T b = c * e[i];
                        e[i + 1] = (r = std::hypot(f, g));
                        if (r == T(0)) {
                            d[i + 1] -= p;
                            e[mz] = T(0);
                            break;
                        }
                        s = f / r;
                        c = g / r;
                        g = d[i + 1] - p;
                        T t = (d[i] - g) * s + T(2) * c * b;
                        p = s * t;
                        d[i + 1] = g + p;
                        g = c * t - b;
                        for (int k = 0; k < n; ++k) {
                            T zi1 = z[k][i + 1];
                            z[k][i + 1] = s * z[k][i] + c * zi1;
                            z[k][i]     = c * z[k][i] - s * zi1;
                        }
                    }
                    if (r == T(0) && mz - 1 >= l) continue;
                    d[l] -= p;
                    e[l] = g;
                    e[mz] = T(0);
                }
            } while (mz != l);
            total_iter += iter;
        }
        return total_iter;
    }

    // extract the d_r sequence and eigenvalue λ from the recurrence
    // expanded over the r sequence of the same parity as (n-m), size N
    template<IsNativeFloat T>
    struct EigPair {
        T lambda;
        std::vector<T> d;  // d_r for r = parity, parity+2, ...
        int parity;        // 0 if (n-m) even, 1 if odd
    };

    template<IsNativeFloat T>
    [[nodiscard]] EigPair<T> compute_eigen(int m, int n, T c_squared, int N = 40) {
        if (m < 0 || n < m) throw std::invalid_argument("spheroidal: requires 0 ≤ m ≤ n");
        const int parity = (n - m) % 2;
        // index k = 0..N-1 → r = parity + 2k
        std::vector<T> diag(N), off(N);
        for (int k = 0; k < N; ++k) {
            int r = parity + 2 * k;
            auto co = rec_coeffs<T>(m, r, c_squared);
            diag[k] = co.B_r;
            // the sub-diagonal e[k] is between d[k] and d[k+1]
            if (k + 1 < N) {
                int r_next = parity + 2 * (k + 1);
                auto co_next = rec_coeffs<T>(m, r_next, c_squared);
                // symmetrization: T_{k,k+1} = sqrt(A_{r(k)} · C_{r(k+1)})
                T prod = co.A_r * co_next.C_r;
                if (prod < T(0)) {
                    // for c² < 0 (oblate), A·C should be > 0 (both have a c² factor)
                    // take absolute value just in case
                    prod = std::abs(prod);
                }
                off[k] = std::sqrt(prod);
            }
        }
        // set the identity matrix as the initial z
        std::vector<std::vector<T>> z(N, std::vector<T>(N, T(0)));
        for (int k = 0; k < N; ++k) z[k][k] = T(1);

        std::vector<T> e = off;
        // since tqli treats e[1..n-1] as the sub-diagonal, pass a shifted copy
        std::vector<T> e_shift(N, T(0));
        for (int i = 0; i < N - 1; ++i) e_shift[i + 1] = off[i];

        tqli(diag, e_shift, z);

        // sort eigenvalues in ascending order (eigenvectors kept in sync)
        std::vector<int> idx(N);
        for (int k = 0; k < N; ++k) idx[k] = k;
        std::sort(idx.begin(), idx.end(),
                  [&](int a, int b) { return diag[a] < diag[b]; });

        const int target = (n - m) / 2;
        if (target >= N) throw std::out_of_range("spheroidal: N too small");

        EigPair<T> out;
        out.lambda = diag[idx[target]];
        out.parity = parity;
        out.d.resize(N);
        // eigenvector: z[k][col] corresponds to diag[col]
        int col = idx[target];
        for (int k = 0; k < N; ++k) out.d[k] = z[k][col];

        // undo the similarity transform applied for symmetrization:
        // transform from the original recurrence → symmetric form d_r → d̃_r = d_r · ψ_r
        // ψ_0 = 1, ψ_{k+1}/ψ_k = sqrt(A_{r(k)} / C_{r(k+1)}) · sign
        // restore d_r to its original scale
        std::vector<T> psi(N, T(1));
        for (int k = 0; k + 1 < N; ++k) {
            int r = parity + 2 * k;
            int r_next = parity + 2 * (k + 1);
            auto co = rec_coeffs<T>(m, r, c_squared);
            auto co_next = rec_coeffs<T>(m, r_next, c_squared);
            T denom = co_next.C_r;
            if (denom == T(0)) {
                // when C_2 = 0 (parity=0, k=0): no scaling needed
                psi[k + 1] = psi[k];
            } else {
                T ratio_sq = co.A_r / denom;
                if (ratio_sq < T(0)) ratio_sq = std::abs(ratio_sq);
                psi[k + 1] = psi[k] * std::sqrt(ratio_sq);
            }
        }
        for (int k = 0; k < N; ++k) {
            if (psi[k] != T(0)) out.d[k] /= psi[k];
        }

        // Flammer normalization: make S_{mn}(c, 0) (n-m even) or S_{mn}'(c, 0) (n-m odd)
        // equal to the corresponding value of P_n^m ─ in practice normalizing by Σ d_r is fine, but
        // here we align by "Σ d_r · P_{m+r}^m(0) = P_n^m(0)" (even)
        // for (n-m) odd, Σ d_r · (P_{m+r}^m)'(0) = (P_n^m)'(0)
        // for simplicity: so that at c=0 only d_0 (parity=0) or d_1 (parity=1) is 1 and the rest 0
        // → "d_target = 1" normalization: scale so that d=1 at the target index k=(n-m)/2
        T pivot = out.d[target];
        if (std::abs(pivot) < std::numeric_limits<T>::min()) {
            // numerically 0 → leave as is (eigenvector being 0 at the target position is rare)
            return out;
        }
        for (T& x : out.d) x /= pivot;

        return out;
    }

}  // namespace detail_spheroidal

// ================================================================
// Prolate spheroidal eigenvalue λ_{mn}(c)  (c² > 0)
// ================================================================
template<IsNativeFloat T>
[[nodiscard]] T spheroidalEigenvalueProlate(int m, int n, T c) {
    if (std::isnan(c)) return std::numeric_limits<T>::quiet_NaN();
    if (m < 0 || n < m) return std::numeric_limits<T>::quiet_NaN();
    if (c == T(0)) return T(n) * T(n + 1);  // λ_{mn}(0) = n(n+1)

    auto eig = detail_spheroidal::compute_eigen<T>(m, n, c * c);
    return eig.lambda;
}

// ================================================================
// Oblate spheroidal eigenvalue (c → ic, i.e. c² → -c²)
// ================================================================
template<IsNativeFloat T>
[[nodiscard]] T spheroidalEigenvalueOblate(int m, int n, T c) {
    if (std::isnan(c)) return std::numeric_limits<T>::quiet_NaN();
    if (m < 0 || n < m) return std::numeric_limits<T>::quiet_NaN();
    if (c == T(0)) return T(n) * T(n + 1);

    auto eig = detail_spheroidal::compute_eigen<T>(m, n, -c * c);
    return eig.lambda;
}

// ================================================================
// Prolate spheroidal angular function S_{mn}^{(1)}(c, η) (first kind)
// ================================================================
//
// S_{mn}(c, η) = Σ_{k} d_{2k+parity}^{mn}(c) · P^m_{m+2k+parity}(η)
//
// Normalization: internally aligned to "d=1 at the target index". This differs from both
// the Meixner-Schäfke convention and the Flammer convention, but is consistent because
// S = P^m_n at c=0.
//
template<IsNativeFloat T>
[[nodiscard]] T spheroidalPS1(int m, int n, T c, T eta) {
    if (std::isnan(c) || std::isnan(eta))
        return std::numeric_limits<T>::quiet_NaN();
    if (m < 0 || n < m) return std::numeric_limits<T>::quiet_NaN();
    if (eta < T(-1) || eta > T(1)) return std::numeric_limits<T>::quiet_NaN();

    if (c == T(0)) return assocLegendreP(n, m, eta);

    auto eig = detail_spheroidal::compute_eigen<T>(m, n, c * c);
    T sum = T(0);
    const int N = static_cast<int>(eig.d.size());
    for (int k = 0; k < N; ++k) {
        int r = eig.parity + 2 * k;
        T pmr = assocLegendreP(m + r, m, eta);
        sum += eig.d[k] * pmr;
    }
    return sum;
}

// ================================================================
// Oblate spheroidal angular function S_{mn}^{O,1}(c, η)
// ================================================================
template<IsNativeFloat T>
[[nodiscard]] T spheroidalOS1(int m, int n, T c, T eta) {
    if (std::isnan(c) || std::isnan(eta))
        return std::numeric_limits<T>::quiet_NaN();
    if (m < 0 || n < m) return std::numeric_limits<T>::quiet_NaN();
    if (eta < T(-1) || eta > T(1)) return std::numeric_limits<T>::quiet_NaN();

    if (c == T(0)) return assocLegendreP(n, m, eta);

    auto eig = detail_spheroidal::compute_eigen<T>(m, n, -c * c);
    T sum = T(0);
    const int N = static_cast<int>(eig.d.size());
    for (int k = 0; k < N; ++k) {
        int r = eig.parity + 2 * k;
        T pmr = assocLegendreP(m + r, m, eta);
        sum += eig.d[k] * pmr;
    }
    return sum;
}

} // namespace special
} // namespace sangi

#endif // SANGI_SPECIAL_SPHEROIDAL_HPP