"""Offline continued-fraction grid of the l = m = 1, n = 0 cloud frequency for the Lab.

Usage: python3 build_growth_reference_grid.py OUTPUT.json
Rows are spins a, columns couplings alpha = M mu. Stored values: M Im(omega) (four significant figures,
signed; negative = the cloud is absorbed by the hole) and M Re(omega) (six figures).
Small couplings start from the hydrogenic guess; from alpha = 0.15 each column is continued from a
linear extrapolation of the two previous solutions, because the hydrogenic guess drifts at large alpha.
Every entry is checked against the superradiance condition: Im(omega) > 0 exactly when Re(omega) < m Omega_H.
"""
import json
import sys
import numpy as np
from kerr_bound_states import find_bound_state, hydrogenic_frequency

SPINS = [round(0.02 * k, 2) for k in range(50)] + [0.99]
COUPLINGS = [round(0.05 + 0.01 * k, 2) for k in range(66)]
CONTINUATION_START = 0.15


def depth_for(alpha):
    return 6000 if alpha < 0.1 else 3000


def build_row(a):
    horizon_rate = a / (2 * (1 + np.sqrt(1 - a * a)))
    rates, reals, solved = [], [], []
    for alpha in COUPLINGS:
        if alpha < CONTINUATION_START or len(solved) < 2:
            guess = hydrogenic_frequency(a, alpha)
        else:
            (alpha_1, omega_1), (alpha_2, omega_2) = solved[-2], solved[-1]
            guess = omega_2 + (omega_2 - omega_1) * (alpha - alpha_2) / (alpha_2 - alpha_1)
        omega = find_bound_state(a, alpha, guess=guess, depth=depth_for(alpha))
        hydrogenic_real = alpha * (1 - alpha * alpha / 8)
        if not np.isfinite(omega) or abs(omega.real / hydrogenic_real - 1) > 0.08:
            raise RuntimeError(f'Mode lost at a={a}, alpha={alpha}: {omega}')
        if np.sign(omega.imag) != np.sign(horizon_rate - omega.real):
            raise RuntimeError(f'Superradiance condition violated at a={a}, alpha={alpha}: {omega}')
        rates.append(float(f'{omega.imag:.4g}'))
        reals.append(float(f'{omega.real:.6g}'))
        solved.append((alpha, omega))
    return rates, reals


def main(output_path):
    solved = [build_row(a) for a in SPINS]
    depth_checks = []
    for a, alpha in ((0.99, 0.05), (0.99, 0.42), (0.7, 0.19), (0.5, 0.1), (0.2, 0.7)):
        depth = depth_for(alpha)
        shallow = find_bound_state(a, alpha, guess=solved[SPINS.index(a)][1][COUPLINGS.index(alpha)] + 1j * solved[SPINS.index(a)][0][COUPLINGS.index(alpha)], depth=depth)
        deep = find_bound_state(a, alpha, guess=shallow, depth=4 * depth)
        depth_checks.append({'a': a, 'alpha': alpha, 'depth': depth, 'imag': shallow.imag, 'imagAtFourTimesDepth': deep.imag,
                             'relativeChange': abs(deep.imag / shallow.imag - 1)})
    payload = {
        'quantity': 'M omega of the l = m = 1, n = 0 quasi-bound scalar mode; Im(omega) > 0 grows',
        'units': 'G = c = M = 1; alpha = G M mu / (hbar c)',
        'method': 'Leaver-type continued fraction, Dolan (2007) arXiv:0705.2880v2 Eqs. (33)-(48), depth 6000 below alpha = 0.1 and 3000 above',
        'provenance': 'Offline numerical reference computed by this site; not observational data',
        'spins': SPINS,
        'couplings': COUPLINGS,
        'growthRate': [rates for rates, _ in solved],
        'realFrequency': [reals for _, reals in solved],
        'depthChecks': depth_checks,
    }
    with open(output_path, 'w') as handle:
        json.dump(payload, handle, separators=(',', ':'))
        handle.write('\n')


if __name__ == '__main__':
    main(sys.argv[1])