"""Where does the printed hydrogenic growth rate agree with a continued-fraction solution?

Usage: python3 compare_hydrogenic_growth.py hydrogenic-comparison-outputs.json

Hydrogenic form: Brito, Cardoso and Pani (2015, v8) Eq. (5.27), l = m = 1, n = 0, gamma_011 = 48.
Numerical form: Dolan (2007) continued fraction (kerr_bound_states.py).
Check against Dolan (2007) Table I: largest M Im(omega) and the coupling where it occurs.
Units G = c = M = 1, alpha = M mu.
"""
import json
import platform
import sys
import numpy as np
import scipy
from kerr_bound_states import find_bound_state, hydrogenic_frequency

SPINS = (0.5, 0.7, 0.9, 0.99)
COUPLINGS = [round(0.05 + 0.01 * k, 2) for k in range(46)]
AGREEMENT_BAND = 0.1
# Dolan (2007), arXiv:0705.2880v2, Table I: spin, coupling of the fastest growth, its M Im(omega).
DOLAN_TABLE_ONE = ((0.7, 0.187, 3.33e-10), (0.8, 0.231, 2.16e-9), (0.9, 0.293, 1.55e-8),
                   (0.95, 0.343, 4.88e-8), (0.98, 0.393, 1.11e-7), (0.99, 0.421, 1.50e-7))


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


def fastest_growth(a, coarse_coupling):
    """Golden-section search on M Im(omega) around a coarse maximum."""
    low, high = coarse_coupling - 0.02, coarse_coupling + 0.02
    ratio = (np.sqrt(5) - 1) / 2
    def rate(alpha):
        return find_bound_state(a, alpha, depth=3000).imag
    for _ in range(40):
        left, right = high - ratio * (high - low), low + ratio * (high - low)
        if rate(left) > rate(right):
            high = right
        else:
            low = left
    alpha = 0.5 * (low + high)
    return alpha, rate(alpha)


def main(output_path):
    comparisons = []
    for a in SPINS:
        rows = []
        for alpha in COUPLINGS:
            numerical = find_bound_state(a, alpha, depth=depth_for(alpha))
            hydrogenic = hydrogenic_frequency(a, alpha)
            if numerical.imag <= 0:
                rows.append({'alpha': alpha, 'numericalRate': numerical.imag, 'hydrogenicRate': hydrogenic.imag,
                             'note': 'not growing in the numerical solution'})
                break
            rows.append({'alpha': alpha, 'numericalRate': numerical.imag, 'hydrogenicRate': hydrogenic.imag,
                         'rateRatioHydrogenicOverNumerical': hydrogenic.imag / numerical.imag,
                         'realFrequencyRatio': hydrogenic.real / numerical.real})
        within = [r['alpha'] for r in rows if abs(r.get('rateRatioHydrogenicOverNumerical', np.inf) - 1) <= AGREEMENT_BAND]
        comparisons.append({'spin': a, 'rows': rows, 'couplingsWithinBand': within})
    table_checks = []
    for a, alpha_published, rate_published in DOLAN_TABLE_ONE:
        alpha, rate = fastest_growth(a, alpha_published)
        table_checks.append({'spin': a, 'publishedCoupling': alpha_published, 'publishedRate': rate_published,
                             'computedCoupling': alpha, 'computedRate': rate,
                             'rateRelativeDifference': abs(rate / rate_published - 1)})
    output = {'software': {'python': platform.python_version(), 'numpy': np.__version__, 'scipy': scipy.__version__},
              'agreementBand': AGREEMENT_BAND, 'comparisons': comparisons, 'dolanTableOneChecks': table_checks}
    with open(output_path, 'w') as handle:
        json.dump(output, handle, indent=1)
        handle.write('\n')
    for item in comparisons:
        print(item['spin'], 'within band:', item['couplingsWithinBand'])
    for check in table_checks:
        print(check)
    return 0 if all(c['rateRelativeDifference'] < 0.01 for c in table_checks) else 1


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