#!/usr/bin/env python3
"""Recompute the gravitational-wave frequency evolution of a circular, equatorial
extreme-mass-ratio inspiral and compare a quadrupole flux with a frozen Teukolsky table.

Standard library only. Usage:
    python3 recompute_frequency_evolution.py inputs.json output.json

Model (adiabatic order, geometric units G = c = M = 1):
    dE/dt = -F(r),  dt/dr = (dE/dr) / F,  N_GW = integral of Omega / pi dt
E(r) and Omega(r) are the exact circular Kerr geodesic values; F is either the
quadrupole flux (32/5) Omega^(10/3) or the tabulated Teukolsky flux to infinity,
interpolated through the relativistic factor F_table / F_quadrupole.
"""
import json
import math
import sys


def orbital_frequency(r, a):
    return 1.0 / (r ** 1.5 + a)


def specific_energy(r, a):
    return (1 - 2 / r + a * r ** -1.5) / math.sqrt(1 - 3 / r + 2 * a * r ** -1.5)


def energy_slope(r, a, step=1e-5):
    # Fourth-order central difference; independent of the closed-form slope used in the Lab.
    h = step * r
    return (-specific_energy(r + 2 * h, a) + 8 * specific_energy(r + h, a)
            - 8 * specific_energy(r - h, a) + specific_energy(r - 2 * h, a)) / (12 * h)


def quadrupole_flux(r, a):
    return 32 / 5 * orbital_frequency(r, a) ** (10 / 3)


def lagrange_cubic(xs, ys, x):
    """Cubic Lagrange interpolation through the four tabulated points nearest x."""
    if x < xs[0] - 1e-12 or x > xs[-1] + 1e-12:
        raise ValueError(f'r = {x} lies outside the frozen table [{xs[0]}, {xs[-1]}]; refusing to extrapolate')
    index = min(max(0, min(range(len(xs)), key=lambda i: abs(xs[i] - x)) - 1), len(xs) - 4)
    nodes = range(index, index + 4)
    total = 0.0
    for i in nodes:
        weight = 1.0
        for j in nodes:
            if j != i:
                weight *= (x - xs[j]) / (xs[i] - xs[j])
        total += weight * ys[i]
    return total


def linear(xs, ys, x):
    for i in range(len(xs) - 1):
        if xs[i] - 1e-12 <= x <= xs[i + 1] + 1e-12:
            t = (x - xs[i]) / (xs[i + 1] - xs[i])
            return ys[i] + t * (ys[i + 1] - ys[i])
    raise ValueError(f'r = {x} lies outside the frozen table')


def table_flux_function(table, a, interpolate=lagrange_cubic):
    column = table['spins'].index(a)
    radii = table['radii']
    factors = [row[column] / quadrupole_flux(r, a) for r, row in zip(radii, table['values'])]
    return lambda r: quadrupole_flux(r, a) * interpolate(radii, factors, r)


def integrate(case, flux, intervals, energy=specific_energy):
    """Simpson integration in ln r from the stop radius out to the start radius."""
    a, q = case['spin'], case['massRatio']
    lower, upper = math.log(case['stopRadius']), math.log(case['startRadius'])
    n = intervals + intervals % 2
    h = (upper - lower) / n

    def slope(r):
        if energy is specific_energy:
            return energy_slope(r, a)
        return (energy(r * (1 + 1e-5)) - energy(r * (1 - 1e-5))) / (2e-5 * r)

    def integrands(s):
        r = math.exp(s)
        dt_dlnr = slope(r) / (q * flux(r)) * r
        return dt_dlnr, dt_dlnr * orbital_frequency(r, a) / math.pi

    time = cycles = 0.0
    for k in range(n + 1):
        weight = 1 if k in (0, n) else (4 if k % 2 else 2)
        dt, dn = integrands(lower + k * h)
        time += weight * dt
        cycles += weight * dn
    return time * h / 3, cycles * h / 3


def frequency_track(case, flux, samples, constants):
    """Time left before the stop radius and f_GW in hertz, on a radius grid."""
    seconds_per_m = case['blackHoleMassSolar'] * constants['solarMassTimeSeconds']
    rows = []
    for k in range(samples + 1):
        r = case['stopRadius'] + (case['startRadius'] - case['stopRadius']) * k / samples
        sub = dict(case, startRadius=r)
        time, cycles = integrate(sub, flux, 400) if k else (0.0, 0.0)
        rows.append({'radius': r, 'gwFrequencyHz': orbital_frequency(r, case['spin']) / (math.pi * seconds_per_m),
                     'timeLeftYears': time * seconds_per_m / constants['julianYearSeconds'], 'cyclesLeft': cycles})
    return rows


def run(inputs):
    table, constants = inputs['fluxTable'], inputs['constants']
    intervals = inputs['integration']['intervals']
    results = []
    for case in inputs['cases']:
        a = case['spin']
        seconds_per_m = case['blackHoleMassSolar'] * constants['solarMassTimeSeconds']
        table_flux = table_flux_function(table, a)
        t_table, n_table = integrate(case, table_flux, intervals)
        t_quad, n_quad = integrate(case, lambda r: quadrupole_flux(r, a), intervals)
        t_lin, n_lin = integrate(case, table_flux_function(table, a, linear), intervals)
        t_coarse, n_coarse = integrate(case, table_flux, 4)
        t_newton, n_newton = integrate(case, table_flux, intervals, energy=lambda r: -1 / (2 * r))
        results.append({
            'id': case['id'], 'input': case,
            'teukolskyTable': {'timeM': t_table, 'years': t_table * seconds_per_m / constants['julianYearSeconds'], 'gwCycles': n_table},
            'quadrupoleFlux': {'timeM': t_quad, 'gwCycles': n_quad, 'cycleRatioToTable': n_quad / n_table},
            'checks': {
                'linearInterpolationCycleChange': n_lin / n_table - 1,
                'fourIntervalSimpsonCycleChange': n_coarse / n_table - 1,
                'newtonianEnergyCycleChange': n_newton / n_table - 1},
            'frequencyTrack': frequency_track(case, table_flux, 8, constants)})
    failure = {}
    try:
        table_flux_function(table, 0.9)(4.0)
    except ValueError as error:
        failure['extrapolateBelowTable'] = str(error)
    return {'schema': 'sandboxphysics.emri-frequency-evolution.output/1',
            'python': sys.version.split()[0], 'results': results, 'refusedCases': failure}


if __name__ == '__main__':
    if len(sys.argv) != 3:
        sys.exit('usage: recompute_frequency_evolution.py inputs.json output.json')
    with open(sys.argv[1], encoding='utf-8') as handle:
        output = run(json.load(handle))
    with open(sys.argv[2], 'w', encoding='utf-8') as handle:
        json.dump(output, handle, indent=2)
        handle.write('\n')
    for row in output['results']:
        print(f"{row['id']}: table N = {row['teukolskyTable']['gwCycles']:.6g}, quadrupole N = {row['quadrupoleFlux']['gwCycles']:.6g}, ratio {row['quadrupoleFlux']['cycleRatioToTable']:.4f}")