#!/usr/bin/env python3
"""Timelike Schwarzschild scattering, c=1 and all lengths in GM.

Run: python3 schwarzschild-scattering.py [--figure-data DIRECTORY]
Requires mpmath==1.3.0. No network, fitting, or numerical-relativity input.
Original implementation of the radial geodesic equation documented on the
Zoom–Whirl Dynamics Near Separatrices page. Results are finite numerical
checks, not bounds for finite-mass-ratio dynamics.
"""
import argparse
import json
import math
from pathlib import Path
import platform

import mpmath as mp


def roots(energy, j):
    """All three real roots of P(u)=2u^3-u^2+2u/j^2+(E^2-1)/j^2."""
    p = 1 / j**2 - mp.mpf(1) / 12
    q = (3 * energy**2 - 2) / (6 * j**2) - mp.mpf(1) / 108
    phase = mp.acos(3 * q / (2 * p) * mp.sqrt(-3 / p)) / 3
    return sorted(mp.mpf(1) / 6 + 2 * mp.sqrt(-p / 3)
                  * mp.cos(phase + 2 * k * mp.pi / 3) for k in range(3))


def angles(delta, dps):
    with mp.workdps(dps):
        energy, jc = 3 / mp.sqrt(7), 7 / mp.sqrt(2)
        j = jc * (1 + mp.mpf(str(delta)))
        a, b, c = roots(energy, j)
        assert a < 0 < b < mp.mpf(2) / 7 < c
        # Remove the turning-point square root with u=b(1-s^2).
        integrand = lambda s: mp.sqrt(2 * b) / mp.sqrt(
            (b * (1 - s*s) - a) * (c - b * (1 - s*s)))
        half_quad = mp.quad(integrand, [0, mp.mpf('.001'), mp.mpf('.01'),
                                        mp.mpf('.1'), 1])
        modulus = (b - a) / (c - a)
        theta0 = mp.asin(mp.sqrt(-a / (b - a)))
        half_elliptic = mp.sqrt(2 / (c - a)) * (
            mp.ellipk(modulus) - mp.ellipf(theta0, modulus))
        residual = max(abs(2*u**3-u**2+2*u/j**2+(energy**2-1)/j**2)
                       for u in (a, b, c))
        assert abs(half_quad-half_elliptic) < mp.mpf(10)**(-dps+12)
        assert residual < mp.mpf(10)**(-dps+8)
        return {'delta': str(delta), 'rho_min': str(1 / b),
                'chi_rad': str(2*half_quad-mp.pi),
                'N_total': str(half_quad/mp.pi),
                'quadrature_elliptic_difference': str(abs(half_quad-half_elliptic)),
                'cubic_residual': str(residual)}


def rk4_angle(delta, step):
    """Independent double-precision orbit ODE, with a linear endpoint crossing."""
    with mp.workdps(50):
        j = float(7 / mp.sqrt(2) * (1 + mp.mpf(str(delta))))
        u = float(roots(3 / mp.sqrt(7), mp.mpf(j))[1])
    v, phi = 0.0, 0.0
    accel = lambda x: -x + 1/j**2 + 3*x*x
    for _ in range(200000):
        a1 = accel(u)
        u2, v2 = u + step*v/2, v + step*a1/2
        a2 = accel(u2)
        u3, v3 = u + step*v2/2, v + step*a2/2
        a3 = accel(u3)
        u4, v4 = u + step*v3, v + step*a3
        new_u = u + step*(v + 2*v2 + 2*v3 + v4)/6
        new_v = v + step*(a1 + 2*a2 + 2*a3 + accel(u4))/6
        if new_u <= 0:
            return 2*(phi + step*u/(u-new_u))-math.pi
        u, v, phi = new_u, new_v, phi+step
    raise AssertionError('Orbit did not escape within integration budget')


def figure_data(directory):
    directory.mkdir(parents=True, exist_ok=True)
    with mp.workdps(40):
        delta = mp.mpf('1e-5')
        j = 7 / mp.sqrt(2) * (1 + delta)
        a, b, c = roots(3 / mp.sqrt(7), j)
        m = (b-a)/(c-a)
        k = mp.ellipk(m)
        theta_edge = mp.asin(mp.sqrt((mp.mpf(1)/12-a)/(b-a)))
        phi_edge = mp.sqrt(2/(c-a))*(k-mp.ellipf(theta_edge,m))
        for name, sign in [('incoming', -1), ('outgoing', 1)]:
            coords = []
            for i in range(1201):
                phi = sign * phi_edge * i / 1200
                u = a + (b-a)*mp.ellipfun('sn', k-abs(phi)*mp.sqrt((c-a)/2), m)**2
                coords.append(f'{float(mp.cos(phi)/u):.10f} {float(mp.sin(phi)/u):.10f}')
            (directory/f'zoom-whirl-{name}.dat').write_text('x y\n'+'\n'.join(coords)+'\n')
        rows = [angles(f'1e-{n}', 40) for n in range(2, 11)]
        (directory/'zoom-whirl-scaling.dat').write_text('logdelta N\n'+''.join(
            f'{n} {float(row["N_total"]):.12f}\n' for n,row in zip(range(2,11),rows)))
        slope = mp.sqrt(mp.mpf(7)/5)*mp.log(10)/(2*mp.pi)
        anchor = mp.mpf(rows[-1]['N_total'])
        (directory/'zoom-whirl-guide.dat').write_text('logdelta N\n'+''.join(
            f'{n} {float(anchor+slope*(n-10)):.12f}\n' for n in [2,10]))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--figure-data', type=Path)
    args = parser.parse_args()
    rows = []
    for n in [2, 4, 6, 8, 10]:
        delta = f'1e-{n}'
        low, high = angles(delta, 40), angles(delta, 65)
        # Compare stored decimal strings with precision sufficient for this check.
        with mp.workdps(70):
            difference = abs(mp.mpf(low['chi_rad'])-mp.mpf(high['chi_rad']))
            assert difference < mp.mpf('1e-27')
        high['precision_difference'] = str(difference)
        if n <= 6:
            exact = float(high['chi_rad'])
            high['rk4_errors'] = [abs(rk4_angle(delta,h)-exact) for h in [.004,.002,.001]]
            assert high['rk4_errors'][-1] < 2e-7
        rows.append(high)
    with mp.workdps(40):
        measured = (mp.mpf(rows[-1]['N_total'])-mp.mpf(rows[-2]['N_total'])) / (2*mp.log(10))
        predicted = mp.sqrt(mp.mpf(7)/5)/(2*mp.pi)
        assert abs(measured-predicted) < mp.mpf('1e-7')
    output = {'model':'timelike Schwarzschild geodesic; no radiation reaction',
        'units':'c=1, lengths divided by GM',
        'environment':{'python':platform.python_version(),'mpmath':mp.__version__},
        'energy':'3/sqrt(7)', 'rho_critical':'7/2', 'j_critical':'7/sqrt(2)',
        'control':'delta=j/j_critical-1 at fixed specific energy',
        'predicted_dN_d_minus_ln_delta':str(predicted),
        'measured_last_two_point_slope':str(measured), 'rows':rows}
    print(json.dumps(output, indent=2))
    if args.figure_data:
        figure_data(args.figure_data)


if __name__ == '__main__':
    main()