#!/usr/bin/env python3
"""Light bending near a Schwarzschild black hole with nothing but an ODE solver (standard library only).

Units G = c = 1 and M = 1.  For a light ray in the equatorial plane put u = 1/r and use the polar angle psi as the
independent variable.  The orbit equation is

        u'' + u = 3 u^2        (' = d/dpsi),        first integral:  u'^2 + u^2 - 2 u^3 = 1/b^2 ,

where b is the impact parameter.  The horizon is u = 1/2.  A ray with b < 3*sqrt(3) = 5.196... falls in, a ray with
b > 3*sqrt(3) comes back out, bent by the angle alpha.  Run:   python3 black_hole_rays.py
"""
import math

BC = 3 * math.sqrt(3)          # critical impact parameter
R0 = 1000.0                    # start (and finish) distance for the deflection angle


def f(y):
    """right-hand side of the first-order system y = (u, v), v = du/dpsi"""
    return (y[1], 3 * y[0] ** 2 - y[0])


def step(method, y, h):
    k1 = f(y)
    if method == "euler":
        return (y[0] + h * k1[0], y[1] + h * k1[1])
    if method == "heun":
        k2 = f((y[0] + h * k1[0], y[1] + h * k1[1]))
        return (y[0] + h / 2 * (k1[0] + k2[0]), y[1] + h / 2 * (k1[1] + k2[1]))
    k2 = f((y[0] + h / 2 * k1[0], y[1] + h / 2 * k1[1]))
    k3 = f((y[0] + h / 2 * k2[0], y[1] + h / 2 * k2[1]))
    k4 = f((y[0] + h * k3[0], y[1] + h * k3[1]))
    return tuple(y[i] + h / 6 * (k1[i] + 2 * k2[i] + 2 * k3[i] + k4[i]) for i in range(2))


def locate(method, y, h, target, frac):
    """fraction of the step h at which the method reaches u = target (Newton iterations)"""
    for _ in range(6):
        z = step(method, y, frac * h)
        frac = min(1.0, max(0.0, frac - (z[0] - target) / (z[1] * h)))
    return frac


def trace(b, method="rk4", h=0.02, max_psi=14 * math.pi):
    """integrate one ray; returns ('captured' | 'escaped' | 'orbit', deflection angle in radians or None)"""
    u0 = 1 / R0
    y = (u0, math.sqrt(1 / b ** 2 - u0 ** 2 * (1 - 2 * u0)))
    psi0 = math.asin(b / R0)
    psi = psi0
    while psi - psi0 < max_psi:
        z = step(method, y, h)
        if z[0] >= 0.5:                                     # reached the horizon
            return "captured", None
        if z[1] < 0 and z[0] <= u0:                         # back at the starting distance, moving outward
            fr = locate(method, y, h, u0, (y[0] - u0) / (y[0] - z[0]))
            return "escaped", (psi + fr * h - psi0) - math.pi + 2 * psi0
        y, psi = z, psi + h
    return "orbit", None


def critical(method, h):
    """critical impact parameter by bisection on 'is the ray captured?'"""
    lo, hi = 4.0, 6.5
    while hi - lo > 1e-9:
        mid = (lo + hi) / 2
        if trace(mid, method, h)[0] == "captured":
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2


if __name__ == "__main__":
    print("critical impact parameter  3*sqrt(3) = %.9f" % BC)
    print("\n  b      fate        alpha (deg)   4/b (deg)   4/b + 15 pi/(4 b^2) (deg)")
    for b in (14, 10, 8, 6, 5.3, 5.2, 5.0, 3.0):
        fate, al = trace(b, "rk4", 0.005)
        w1, w2 = 4 / b, 4 / b + 15 * math.pi / (4 * b * b)
        print("%5.2f  %-9s  %12s  %10.3f  %14.3f" % (b, fate, "%.3f" % math.degrees(al) if al is not None else "-", math.degrees(w1), math.degrees(w2)))
    print("\nbisection for b_c, error versus 3*sqrt(3):")
    for h in (0.1, 0.02):
        for m in ("euler", "heun", "rk4"):
            bc = critical(m, h)
            print("  h = %-5s %-6s b_c = %.6f   error %+.2e" % (h, m, bc, bc - BC))
    print("\norder of convergence of the deflection angle at b = 8:")
    ref = trace(8, "rk4", 0.0005)[1]
    for m in ("euler", "heun", "rk4"):
        e1 = abs(trace(8, m, 0.05)[1] - ref)
        e2 = abs(trace(8, m, 0.025)[1] - ref)
        print("  %-6s error(h=0.05) = %.2e   error(h=0.025) = %.2e   observed order %.2f" % (m, e1, e2, math.log(e1 / e2) / math.log(2)))