"""Spin-0 spheroidal eigenvalue in a normalized associated-Legendre basis.

Angular equation (units with M = 1):
    (1/sin t) d/dt (sin t dS/dt) + (c2 cos^2 t - m^2/sin^2 t + A) S = 0
so A is an eigenvalue of L^2 - c2 x^2 with x = cos t. For c2 = 0, A = l(l+1).
c2 = (a omega)^2 for a massless wave and a^2 (omega^2 - mu^2) for a massive one;
complex c2 is allowed.
"""
import numpy as np
from scipy.special import lpmv, gammaln


def _normalized_legendre(degree, order, x):
    log_norm = 0.5 * (np.log((2 * degree + 1) / 2.0) + gammaln(degree - order + 1) - gammaln(degree + order + 1))
    return np.exp(log_norm) * lpmv(order, degree, x)


def spheroidal_eigenvalue(l, m, c2, basis_size=24):
    """Return A for the branch that tends to l(l+1) when c2 -> 0."""
    order = abs(m)
    if l < order:
        raise ValueError('l must be at least |m|')
    degrees = np.arange(order + (l - order) % 2, order + (l - order) % 2 + 2 * basis_size, 2)
    nodes, weights = np.polynomial.legendre.leggauss(4 * basis_size + 40)
    functions = np.array([_normalized_legendre(degree, order, nodes) for degree in degrees])
    # einsum avoids spurious floating-point warnings from some BLAS matmul builds.
    x_squared = np.einsum('ik,jk->ij', functions * weights * nodes ** 2, functions)
    matrix = np.diag(degrees * (degrees + 1.0)).astype(complex) - complex(c2) * x_squared
    values = np.linalg.eigvals(matrix)
    target = int(np.where(degrees == l)[0][0])
    # Follow the c2 = 0 branch: start from the diagonal and pick the closest eigenvalue.
    first_order = l * (l + 1.0) - complex(c2) * x_squared[target, target]
    value = values[np.argmin(np.abs(values - first_order))]
    return value if np.iscomplexobj(c2) or abs(value.imag) > 1e-14 else value.real


if __name__ == '__main__':
    # Small-c limit: A = l(l+1) - c2 <x^2>, with <x^2> = 1/5 for l = m = 1.
    for c2 in (0.0, 1e-4, 0.1, 0.5):
        print(c2, spheroidal_eigenvalue(1, 1, c2), 2 - c2 / 5)