"""Amplification factor of a massless scalar wave scattered by a Kerr black hole.

Units G = c = M = 1. Radial Teukolsky equation for spin weight s = 0:
    d/dr (Delta dR/dr) + (K^2/Delta - lam) R = 0,
    Delta = r^2 - 2r + a^2,  K = omega (r^2 + a^2) - a m,
    lam = A + a^2 omega^2 - 2 a m omega.
Boundary condition: purely ingoing at the horizon, R ~ (r - r_+)^(-i sigma) with
sigma = 2 r_+ (omega - m Omega_H) / (r_+ - r_-).
At large r: R ~ I exp(-i omega r*)/r + R_out exp(+i omega r*)/r and
Z = |R_out/I|^2 - 1 (Brito, Cardoso and Pani, Eq. 4.93 with s = 0).
"""
import numpy as np
from scipy.integrate import solve_ivp
from spheroidal_eigenvalue import spheroidal_eigenvalue


def horizon_quantities(a):
    root = np.sqrt(1.0 - a * a)
    r_plus, r_minus = 1.0 + root, 1.0 - root
    return r_plus, r_minus, a / (2.0 * r_plus)


def _poly_mul(p, q):
    return np.convolve(p, q)


def _tortoise(r, a):
    r_plus, r_minus, _ = horizon_quantities(a)
    width = r_plus - r_minus
    tortoise = r + 2 * r_plus / width * np.log(r - r_plus)
    if width < 1e-14:
        raise ValueError('Extremal spin is outside this solver')
    return tortoise - 2 * r_minus / width * np.log(r - r_minus)


def _horizon_series(omega, a, m, lam, terms=8):
    """Frobenius coefficients b_n of R = x^nu sum b_n x^n, x = r - r_+ (ascending powers)."""
    r_plus, r_minus, omega_h = horizon_quantities(a)
    width = r_plus - r_minus
    nu = -1j * 2 * r_plus * (omega - m * omega_h) / width
    # Polynomials in x, lowest power first.
    delta = np.array([0.0, width, 1.0])
    delta_prime = np.array([width, 2.0])
    k_poly = omega * np.array([r_plus ** 2 + a * a, 2 * r_plus, 1.0], dtype=complex)
    k_poly[0] -= a * m
    p2 = _poly_mul(delta, delta)                 # multiplies R''
    p1 = _poly_mul(delta, delta_prime)           # multiplies R'
    p0 = _poly_mul(k_poly, k_poly) - lam * np.pad(delta, (0, 2)).astype(complex)
    def indicial(power):
        # Coefficient of x^(power) in p2 x^(power-2) power(power-1) + ... at the lowest order of each polynomial.
        return p2[2] * power * (power - 1) + p1[1] * power + p0[0]
    b = [1.0 + 0j]
    for n in range(1, terms):
        total = 0j
        for k in range(n):
            power = nu + k
            shift = n - k
            # x^(power + shift) gathers terms where polynomial power j satisfies j - 2 = shift (p2), j - 1 = shift (p1), j = shift (p0).
            coefficient = 0j
            if shift + 2 < len(p2):
                coefficient += p2[shift + 2] * power * (power - 1)
            if shift + 1 < len(p1):
                coefficient += p1[shift + 1] * power
            if shift < len(p0):
                coefficient += p0[shift]
            total += coefficient * b[k]
        b.append(-total / indicial(nu + n))
    return nu, np.array(b)


def _infinity_series(omega, a, m, lam, sign, terms=12):
    """Coefficients c_n of R = exp(i sign omega r*) sum c_n r^(-n-1)."""
    eps = sign
    # Polynomials in r, lowest power first.
    delta = np.array([a * a, -2.0, 1.0], dtype=complex)
    r2a2 = np.array([a * a, 0.0, 1.0], dtype=complex)
    q2 = _poly_mul(delta, delta)
    q1 = _poly_mul(delta, 2j * eps * omega * r2a2 + np.array([-2.0, 2.0, 0.0]))
    # q0 = Delta (2 i eps omega r - lam) + a^2 m^2 - 2 a m omega (r^2 + a^2)
    q0 = _poly_mul(delta, np.array([-lam, 2j * eps * omega], dtype=complex))
    q0[:3] += np.array([a * a * m * m - 2 * a * m * omega * a * a, 0.0, -2 * a * m * omega])
    c = [1.0 + 0j]
    # Term c_n r^(-k), k = n + 1: contributes q2_j k(k+1) r^(j-k-2), -q1_j k r^(j-k-1), q0_j r^(j-k).
    # The leading power r^(3-k) has coefficient 2 i eps omega (1 - k), which vanishes for k = 1.
    def contribution(k, power):
        value = 0j
        j = power + k + 2
        if 0 <= j < len(q2):
            value += q2[j] * k * (k + 1)
        j = power + k + 1
        if 0 <= j < len(q1):
            value += -q1[j] * k
        j = power + k
        if 0 <= j < len(q0):
            value += q0[j]
        return value
    for n in range(1, terms):
        k_new = n + 1
        power = 3 - k_new
        total = 0j
        for previous in range(n):
            total += contribution(previous + 1, power) * c[previous]
        c.append(-total / (2j * eps * omega * (1 - k_new)))
    return np.array(c)


def _asymptotic_solution(r, omega, a, m, lam, sign, coefficients):
    r_plus, r_minus, _ = horizon_quantities(a)
    delta = r * r - 2 * r + a * a
    phase = np.exp(1j * sign * omega * _tortoise(r, a))
    powers = np.arange(len(coefficients))
    value = np.sum(coefficients * r ** (-powers - 1.0))
    derivative = np.sum(coefficients * (-powers - 1.0) * r ** (-powers - 2.0))
    return phase * value, phase * (derivative + 1j * sign * omega * (r * r + a * a) / delta * value)


def amplification_factor(omega, a, l=1, m=1, outer_radius=None, rtol=1e-11, return_details=False):
    """Z_{0lm}(omega) for 0 < omega; spin 0 <= a < 1."""
    if omega <= 0:
        raise ValueError('Use positive frequency; Z(omega, m) = Z(-omega, -m)')
    r_plus, r_minus, omega_h = horizon_quantities(a)
    a_angular = spheroidal_eigenvalue(l, m, (a * omega) ** 2)
    lam = a_angular + a * a * omega * omega - 2 * a * m * omega
    nu, b = _horizon_series(omega, a, m, lam)
    x0 = 1e-3 * (r_plus - r_minus) if r_plus - r_minus > 1e-2 else 1e-5
    powers = np.arange(len(b))
    R0 = np.sum(b * x0 ** (nu + powers))
    dR0 = np.sum(b * (nu + powers) * x0 ** (nu + powers - 1))
    if outer_radius is None:
        outer_radius = max(250.0, 80.0 / omega)

    def rhs(r, y):
        delta = r * r - 2 * r + a * a
        k_value = omega * (r * r + a * a) - a * m
        second = -((2 * r - 2) * y[1] + (k_value * k_value / delta - lam) * y[0]) / delta
        return [y[1], second]

    solution = solve_ivp(rhs, (r_plus + x0, outer_radius), [R0, dR0], method='DOP853', rtol=rtol, atol=1e-14 * abs(R0))
    if not solution.success:
        raise RuntimeError(solution.message)
    R_end, dR_end = solution.y[0, -1], solution.y[1, -1]
    incoming = _asymptotic_solution(outer_radius, omega, a, m, lam, -1, _infinity_series(omega, a, m, lam, -1))
    outgoing = _asymptotic_solution(outer_radius, omega, a, m, lam, +1, _infinity_series(omega, a, m, lam, +1))
    matrix = np.array([[incoming[0], outgoing[0]], [incoming[1], outgoing[1]]])
    amplitude_in, amplitude_out = np.linalg.solve(matrix, [R_end, dR_end])
    z_value = abs(amplitude_out) ** 2 / abs(amplitude_in) ** 2 - 1.0
    if not return_details:
        return z_value
    # Conserved Wronskian W = Im(Delta conj(R) R'). With unit horizon amplitude, the horizon value is
    # -2 r_+ k_H and the value at infinity is omega (|R_out|^2 - |I|^2), so
    # omega (|I|^2 - |R_out|^2) = 2 r_+ k_H: energy leaves the hole (Z > 0) exactly when k_H < 0.
    k_h = omega - m * omega_h
    horizon_flux = 2 * r_plus * k_h
    infinity_flux = omega * (abs(amplitude_in) ** 2 - abs(amplitude_out) ** 2)
    return {'Z': z_value, 'lambda': lam, 'A': a_angular, 'I': amplitude_in, 'R': amplitude_out,
            'horizonWronskian': horizon_flux, 'infinityWronskian': infinity_flux, 'outerRadius': outer_radius,
            'steps': int(solution.t.size)}


def wronskian_along_path(omega, a, l=1, m=1, samples=6):
    """Im(Delta R* R') must be constant in r (the radial equation has real coefficients)."""
    r_plus, r_minus, omega_h = horizon_quantities(a)
    a_angular = spheroidal_eigenvalue(l, m, (a * omega) ** 2)
    lam = a_angular + a * a * omega * omega - 2 * a * m * omega
    nu, b = _horizon_series(omega, a, m, lam)
    x0 = 1e-3 * (r_plus - r_minus)
    powers = np.arange(len(b))
    R0 = np.sum(b * x0 ** (nu + powers)); dR0 = np.sum(b * (nu + powers) * x0 ** (nu + powers - 1))
    outer = max(250.0, 80.0 / omega)
    radii = np.geomspace(r_plus + x0, outer, samples)

    def rhs(r, y):
        delta = r * r - 2 * r + a * a
        k_value = omega * (r * r + a * a) - a * m
        return [y[1], -((2 * r - 2) * y[1] + (k_value * k_value / delta - lam) * y[0]) / delta]

    solution = solve_ivp(rhs, (radii[0], radii[-1]), [R0, dR0], method='DOP853', rtol=1e-11, atol=1e-16, t_eval=radii)
    delta = solution.t ** 2 - 2 * solution.t + a * a
    return solution.t, np.imag(delta * np.conj(solution.y[0]) * solution.y[1])