"""Quasi-bound states of a massive scalar field around a Kerr black hole.

Units G = c = M = 1, field mass mu (so mu equals the gravitational coupling alpha = G M mu / (hbar c)).
Time dependence exp(-i omega t): Im(omega) > 0 means the cloud grows.
Method: Leaver-type continued fraction with the coefficients of Dolan, Phys. Rev. D 76, 084001 (2007),
arXiv:0705.2880v2, Eqs. (33)-(48), with Re(q) < 0 for bound states.
The spheroidal separation constant uses c^2 = a^2 (omega^2 - mu^2) (Dolan Eq. 46), which is complex.
"""
import numpy as np
from spheroidal_eigenvalue import spheroidal_eigenvalue


def _recurrence_coefficients(omega, a, mu, l, m):
    b = np.sqrt(1.0 - a * a)
    q = -np.sqrt(mu * mu - omega * omega + 0j)
    if q.real > 0:
        q = -q
    separation = spheroidal_eigenvalue(l, m, a * a * (omega * omega - mu * mu))
    w = omega - a * m / 2.0
    s = omega - 1j * q
    c0 = 1 - 2j * omega - 2j / b * w
    c1 = -4 + 4j * (omega - 1j * q * (1 + b)) + 4j / b * w - 2 * (omega ** 2 + q ** 2) / q
    c2 = 3 - 2j * omega - 2 * (q ** 2 - omega ** 2) / q - 2j / b * w
    c3 = (2j * s ** 3 / q + 2 * s ** 2 * b + q * q * a * a + 2j * q * a * m - separation - 1
          - s ** 2 / q + 2 * q * b + 2j / b * (s ** 2 / q + 1) * w)
    c4 = s ** 4 / q ** 2 + 2j * omega * s ** 2 / q - 2j / b * s ** 2 / q * w
    return c0, c1, c2, c3, c4


def continued_fraction(omega, a, mu, l=1, m=1, depth=4000):
    """Left side of Dolan Eq. (48); zero at a quasi-bound frequency."""
    c0, c1, c2, c3, c4 = _recurrence_coefficients(omega, a, mu, l, m)
    n = np.arange(depth + 1, dtype=float)
    alpha = n * n + (c0 + 1) * n + c0
    beta = -2 * n * n + (c1 + 2) * n + c3
    gamma = n * n + (c2 - 3) * n + c4
    tail = 0j
    for k in range(depth, 0, -1):
        tail = alpha[k - 1] * gamma[k] / (beta[k] - tail)
    return (beta[0] - tail) / beta[0]


def hydrogenic_frequency(a, mu, l=1, m=1, overtone=0):
    """Brito, Cardoso and Pani (2015, v8), Eq. (5.27), with gamma_011 = 48 for l = m = 1, n = 0."""
    r_plus = 1 + np.sqrt(1 - a * a)
    real = mu - mu / 2 * (mu / (l + overtone + 1)) ** 2
    if (l, m, overtone) != (1, 1, 0):
        return real
    return real + 1j * (a * m - 2 * mu * r_plus) * mu ** (4 * l + 5) / 48.0


def find_bound_state(a, mu, l=1, m=1, guess=None, depth=4000, tolerance=1e-15, iterations=200):
    """Complex secant iteration on the continued fraction, started from the hydrogenic guess."""
    x0 = hydrogenic_frequency(a, mu, l, m) if guess is None else complex(guess)
    # Neighbouring levels are separated by about mu^3, so steps are scaled and clamped to that size.
    spacing = mu ** 3
    x1 = x0 + 1e-4 * spacing
    f0 = continued_fraction(x0, a, mu, l, m, depth)
    f1 = continued_fraction(x1, a, mu, l, m, depth)
    for _ in range(iterations):
        if f1 == f0:
            break
        step = -f1 * (x1 - x0) / (f1 - f0)
        if abs(step) > 0.05 * spacing:
            step *= 0.05 * spacing / abs(step)
        x2 = x1 + step
        x0, f0 = x1, f1
        x1, f1 = x2, continued_fraction(x2, a, mu, l, m, depth)
        if abs(x1 - x0) < tolerance * abs(x1):
            break
    return x1


if __name__ == '__main__':
    for a, mu in ((0.99, 0.421), (0.9, 0.293), (0.7, 0.187), (0.99, 0.1), (0.99, 0.05)):
        omega = find_bound_state(a, mu)
        deeper = find_bound_state(a, mu, guess=omega, depth=8000)
        print(a, mu, omega, deeper, hydrogenic_frequency(a, mu))