"""Analytic point-lens amplification, Takahashi & Nakamura (2003) Eq. (17), and its limits.

F(w, y) = exp[pi w/4 + i (w/2)(ln(w/2) - 2 phi_m(y))] Gamma(1 - i w/2) 1F1(i w/2, 1; i w y^2/2)
with phi_m(y) = (x_m - y)^2/2 - ln x_m and x_m = (y + sqrt(y^2 + 4))/2.

Units: w = 8 pi G M_Lz f / c^3 is dimensionless; y is in Einstein radii.
Convention: F multiplies h(f) = integral h(t) exp(+2 pi i f t) dt, and the arrival-time
function T(x, y) = |x - y|^2/2 - ln|x| - phi_m(y) has its minimum at zero.
"""
import mpmath as mp
import numpy as np

mp.mp.dps = 60


def arrival_potential_minimum(y):
    x_minimum = (y + mp.sqrt(y * y + 4)) / 2
    return (x_minimum - y) ** 2 / 2 - mp.log(x_minimum)


def analytic_amplification(w, y):
    w, y = mp.mpf(w), mp.mpf(y)
    phase = mp.pi * w / 4 + 1j * (w / 2) * (mp.log(w / 2) - 2 * arrival_potential_minimum(y))
    return complex(mp.exp(phase) * mp.gamma(1 - 1j * w / 2) * mp.hyp1f1(1j * w / 2, 1, 1j * w * y * y / 2))


def image_properties(y):
    root = np.sqrt(y * y + 4)
    ratio = (y * y + 2) / (2 * y * root)
    return {
        "minimum_magnification": 0.5 + ratio,
        "saddle_magnification_magnitude": ratio - 0.5,
        "dimensionless_delay": y * root / 2 + np.log((root + y) / (root - y)),
    }


def geometric_optics_amplification(w, y):
    image = image_properties(y)
    return np.sqrt(image["minimum_magnification"]) - 1j * np.sqrt(image["saddle_magnification_magnitude"]) * np.exp(
        1j * w * image["dimensionless_delay"]
    )