"""Independent numerical Fresnel integral for an isolated point-mass lens.

It does not use the confluent hypergeometric closed form. The lens-plane integral
    F(w) = (w / 2 pi i) * integral d^2x exp[i w T(x, y)]
is rewritten through the area A(tau) of the region T < tau:
    I(tau) = (1 / 2 pi) dA/dtau,   F(w) = 1 - i w * integral_0^inf [I(tau) - 1] exp(i w tau) dtau.
For each tau the contour T = tau is found by bisection along every polar ray from the lens,
I(tau) is a Gauss-Legendre angle integral after a substitution that removes the endpoint
square-root singularity, the Fourier integral is exact for a piecewise-linear I(tau)
(Filon-type), and the tail beyond tau_max uses I - 1 = c1/tau + c2/tau^2 fitted on the grid.
"""
import numpy as np
from scipy.special import exp1

GAUSS_NODES, GAUSS_WEIGHTS = np.polynomial.legendre.leggauss(240)
GAUSS_NODES01, GAUSS_WEIGHTS01 = (GAUSS_NODES + 1) / 2, GAUSS_WEIGHTS / 2


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


def ray_minimum(theta, y):
    c = np.cos(theta)
    r = (y * c + np.sqrt(y * y * c * c + 4)) / 2
    return r, r * r / 2 - r * y * c + y * y / 2 - np.log(r) - arrival_offset(y)


def roots_on_rays(theta, tau, y):
    c = np.cos(theta)
    r_min, _ = ray_minimum(theta, y)
    offset = arrival_offset(y)

    def excess(log_r):
        r = np.exp(log_r)
        return r * r / 2 - r * y * c + y * y / 2 - log_r - offset - tau

    roots = []
    for low, high in [
        (np.full_like(theta, -(tau + offset + y * y + 8.0)), np.log(r_min)),
        (np.log(r_min), np.log(r_min + np.sqrt(2 * tau + 2) + 2 * y + 4)),
    ]:
        low, high = low.copy(), high.copy()
        sign_low = np.sign(excess(low))
        for _ in range(80):
            middle = (low + high) / 2
            same = np.sign(excess(middle)) == sign_low
            low = np.where(same, middle, low)
            high = np.where(same, high, middle)
        roots.append(np.exp((low + high) / 2))
    inner, outer = roots

    def radial_slope(r):
        return r - y * c - 1 / r

    return inner / np.abs(radial_slope(inner)) + outer / np.abs(radial_slope(outer))


def saddle_delay(y):
    root = np.sqrt(y * y + 4)
    return y * root / 2 + np.log((root + y) / (root - y))


def boundary_angle(tau, y):
    low, high = 0.0, np.pi
    for _ in range(80):
        middle = (low + high) / 2
        if ray_minimum(np.array([middle]), y)[1][0] < tau:
            low = middle
        else:
            high = middle
    return (low + high) / 2


def time_domain_amplification(tau, y):
    """I(tau): normalized to 1 for an unlensed wave."""
    if tau <= 0:
        return 0.0
    delay = saddle_delay(y)
    if tau < delay:
        edge = boundary_angle(tau, y)
        u = GAUSS_NODES01
        theta = edge * (1 - u * u)
        integrand = roots_on_rays(theta, tau, y) * 2 * edge * u
        return float(np.sum(GAUSS_WEIGHTS01 * integrand) / np.pi)
    # Above the saddle every ray crosses the contour; cluster nodes near theta = pi.
    excess = max(tau - delay, 1e-14)
    curvature = (delay - ray_minimum(np.array([np.pi - 1e-3]), y)[1][0]) / 1e-6
    width = np.sqrt(excess / max(curvature, 1e-12))
    v_max = np.arcsinh(np.pi / width)
    v = GAUSS_NODES01 * v_max
    phi = width * np.sinh(v)
    integrand = roots_on_rays(np.pi - phi, tau, y) * width * np.cosh(v) * v_max
    return float(np.sum(GAUSS_WEIGHTS01 * integrand) / np.pi)


def tau_grid(y, tau_max=400.0):
    delay = saddle_delay(y)
    near_zero = np.geomspace(1e-9, 0.05, 160)
    near_saddle = delay + np.concatenate([-np.geomspace(1e-9, min(0.5 * delay, 0.05), 160)[::-1], np.geomspace(1e-9, 0.05, 160)])
    body = np.linspace(0.0, 60.0, 6001)
    tail = np.geomspace(60.0, tau_max, 500)
    grid = np.unique(np.concatenate([near_zero, near_saddle, body[1:], tail]))
    return grid[grid > 0]


def filon_linear(tau, values, w):
    """Exact integral of a piecewise-linear function times exp(i w tau)."""
    t0, t1, f0, f1 = tau[:-1], tau[1:], values[:-1], values[1:]
    h = t1 - t0
    slope = (f1 - f0) / h
    e0, e1 = np.exp(1j * w * t0), np.exp(1j * w * t1)
    return np.sum((f1 * e1 - f0 * e0) / (1j * w) + slope * (e1 - e0) / w**2)


def tabulate(y, tau_max=400.0):
    tau = tau_grid(y, tau_max)
    values = np.array([time_domain_amplification(t, y) for t in tau])
    return tau, values


def fresnel_amplification(w, tau, values, y):
    """F(w) from a tabulated I(tau); the segment [0, tau_1] uses I(0+) = sqrt(mu+)."""
    minimum_amplitude = np.sqrt(0.5 + (y * y + 2) / (2 * y * np.sqrt(y * y + 4)))
    full_tau = np.concatenate([[0.0], tau])
    full_values = np.concatenate([[minimum_amplitude], values]) - 1.0
    body = filon_linear(full_tau, full_values, w)
    fit_mask = full_tau > 0.5 * full_tau[-1]
    design = np.vstack([1 / full_tau[fit_mask], 1 / full_tau[fit_mask] ** 2]).T
    c1, c2 = np.linalg.lstsq(design, full_values[fit_mask], rcond=None)[0]
    tau_max = full_tau[-1]
    # integral_{tau_max}^inf exp(i w t)/t dt = E1(-i w tau_max); 1/t^2 term by parts.
    e1_term = exp1(-1j * w * tau_max)
    tail = c1 * e1_term + c2 * (np.exp(1j * w * tau_max) / tau_max + 1j * w * e1_term)
    return 1 - 1j * w * (body + tail), (c1, c2)