#!/usr/bin/env python3
"""Benchmark of the first-order Ghosh quadrupolar scalar PDE.

Dimensionless R=r/mu, Omega=mu*omega, x=cos(theta), u=R*Phi, m=0.
On R in [2.5,10], impose homogeneous radial Dirichlet data, regular poles.
No horizon/outgoing conditions, QNM claim, time evolution, or all-order metric.
Source = 16 t^2(1-t)^2 chi_0(x), t=(R-Rin)/(Rout-Rin), chi_0=1/sqrt(2).
Norm^2 = integral |u|^2 dR dx/f (equivalently |Phi|^2 R^2 dR dx/f).
The direct angular collocation uses pointwise multiplication and derivatives,
not the Galerkin coupling matrix. Radial collocation is shared and converged.

Original implementation for blackhole.org, 2026-09-05. Primary equation:
Ghosh, Franchini, Voelkel, Barausse, Physical Review D 108 (2023) 024038,
arXiv:2303.00088v2, printed pp. 4-5, Eqs. (28)-(34).
Rerun with --output PATH; --verify-saved PATH compares a complete fresh run
with every saved field, excluding Python-version provenance, wall-clock time,
and unresolved last digits. Pinned numerical-library versions still match.
"""

import argparse
import hashlib
import json
import os
import platform
from pathlib import Path
import time

# Bound BLAS parallelism before NumPy import for predictable small-matrix work.
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
os.environ.setdefault("OMP_NUM_THREADS", "1")
import numpy as np
import scipy
from scipy.special import eval_legendre, roots_legendre
from scipy.sparse import csc_matrix
from scipy.sparse.linalg import splu, LinearOperator, onenormest

RIN, ROUT = 2.5, 10.0
FREQUENCIES = (0.35-0.025j, 0.55-0.01j)
EPSILONS = (0.0, 0.25, 0.5, 1.0, 2.0)
LMAX = (0, 2, 4, 6, 8)


def cheb(n, lo=-1., hi=1.):
    x = np.cos(np.pi*np.arange(n+1)/n)
    nodes = (lo+hi)/2+(hi-lo)*x/2
    weights = (-1.)**np.arange(n+1)
    weights[[0, -1]] *= 0.5
    differences = nodes[:, None]-nodes[None, :]
    np.fill_diagonal(differences, 1.)
    derivative = weights[None, :]/weights[:, None]/differences
    np.fill_diagonal(derivative, 0.)
    derivative[np.diag_indices_from(derivative)] = -derivative.sum(axis=1)
    assert np.max(np.abs(derivative@nodes-1)) < 1e-10
    return nodes, weights, derivative


def interpolation(nodes, weights, evaluation):
    difference = evaluation[:, None]-nodes[None, :]
    exact = abs(difference) < 5e-14
    difference[exact] = 1.
    matrix = weights[None, :]/difference
    matrix /= matrix.sum(axis=1)[:, None]
    for i in np.flatnonzero(exact.any(axis=1)):
        matrix[i] = exact[i].astype(float)
    return matrix


def f(R):
    return 1-2/R


def f1(R):
    return 2*f(R)/R**3


def forcing(R):
    tau = (R-RIN)/(ROUT-RIN)
    return 16*tau**2*(1-tau)**2


def chi(ell, x):
    return np.sqrt((2*ell+1)/2)*eval_legendre(ell, x)


def quadrupole(x):
    return (3*x*x-1)/2


def coupling(ells):
    matrix = np.zeros((len(ells), len(ells)))
    q = lambda ell: ell/np.sqrt(4*ell*ell-1) if ell > 0 else 0.
    for i, ell in enumerate(ells):
        for j, source_ell in enumerate(ells):
            if ell == source_ell:
                matrix[i,j] = 1.5*(q(ell+1)**2+q(ell)**2)-0.5
            elif ell == source_ell+2:
                matrix[i,j] = 1.5*q(source_ell+1)*q(source_ell+2)
            elif ell == source_ell-2:
                matrix[i,j] = 1.5*q(source_ell)*q(source_ell-1)
    return matrix


def radial(n, omega):
    nodes, bary, d = cheb(n, RIN, ROUT)
    R = nodes[1:-1]
    a = np.diag(f(R)**2)@(d@d)[1:-1, 1:-1]
    a += np.diag(f(R)*2/R**2)@d[1:-1, 1:-1]
    a = a.astype(complex)+np.diag(omega**2-2*f(R)/R**3)
    return nodes, bary, R, a


def solve(matrix, rhs):
    lu = splu(csc_matrix(matrix))
    result = lu.solve(rhs)
    backward = np.linalg.norm(matrix@result-rhs)/(
        np.linalg.norm(matrix, ord=np.inf)*np.linalg.norm(result)+np.linalg.norm(rhs))
    assert backward < 5e-13
    return result, backward, lu


def galerkin(nr, lmax, omega, epsilon, custom_source=None):
    nodes, bary, R, ar = radial(nr, omega)
    ells = np.arange(0, lmax+1, 2)
    c = coupling(ells)
    matrix = np.kron(np.eye(len(ells)), ar)
    matrix += np.kron(-np.diag(ells*(ells+1)), np.diag(f(R)/R**2))
    matrix += np.kron(c, np.diag(-2*epsilon*omega**2*f1(R)))
    rhs = np.zeros((len(ells), len(R)), dtype=complex)
    rhs[0] = forcing(R)
    if custom_source is not None:
        qx, qw = roots_legendre(48)
        basis = np.array([chi(ell,qx) for ell in ells])
        rhs = (basis*qw)@custom_source(R,qx)
    solution, backward, lu = solve(matrix, rhs.ravel())
    values = np.zeros((len(ells), nr+1), dtype=complex)
    values[:, 1:-1] = solution.reshape(len(ells), len(R))
    return dict(kind="galerkin", R=nodes, Rb=bary, ells=ells, values=values,
                backward=backward)


def direct(nr, nx, omega, epsilon, condition=False, custom_source=None):
    nodes, bary, R, ar = radial(nr, omega)
    x, xb, dx = cheb(nx)
    # At x=+-1 this is the regular degenerate PDE row, not Dirichlet data.
    angular = np.diag(1-x*x)@(dx@dx)-2*np.diag(x)@dx
    matrix = np.kron(np.eye(nx+1), ar)
    matrix += np.kron(angular, np.diag(f(R)/R**2))
    matrix += np.kron(np.diag(quadrupole(x)),
                      np.diag(-2*epsilon*omega**2*f1(R)))
    rhs = np.tile(forcing(R)/np.sqrt(2), nx+1).astype(complex)
    if custom_source is not None:
        rhs = custom_source(R,x).ravel()
    solution, backward, lu = solve(matrix, rhs)
    values = np.zeros((nx+1, nr+1), dtype=complex)
    values[:, 1:-1] = solution.reshape(nx+1, nr-1)
    info = dict(kind="direct", R=nodes, Rb=bary, x=x, xb=xb,
                values=values, backward=backward)
    if condition:
        inverse = LinearOperator(matrix.shape, dtype=complex,
            matvec=lambda v: lu.solve(v), rmatvec=lambda v: lu.solve(v, trans="H"),
            matmat=lambda v: lu.solve(v), rmatmat=lambda v: lu.solve(v, trans="H"))
        estimate = onenormest(inverse)
        info["inverse_1norm_estimate"] = float(estimate)
        info["condition_1norm_estimate"] = float(np.linalg.norm(matrix, 1)*estimate)
    return info


def evaluate(solution, r_eval, x_eval):
    pr = interpolation(solution["R"], solution["Rb"], r_eval)
    if solution["kind"] == "galerkin":
        px = np.array([chi(ell, x_eval) for ell in solution["ells"]]).T
    else:
        px = interpolation(solution["x"], solution["xb"], x_eval)
    return px@solution["values"]@pr.T


def offgrid_residual(solution, r_eval, x_eval, omega, epsilon):
    """Apply the literal 2D PDE away from all collocation nodes."""
    pr = interpolation(solution["R"],solution["Rb"],r_eval)
    _, _, dr = cheb(len(solution["R"])-1,RIN,ROUT)
    if solution["kind"] == "galerkin":
        ells = solution["ells"]
        px = np.array([chi(ell,x_eval) for ell in ells]).T
        angular_values = -ells[:,None]*(ells[:,None]+1)*solution["values"]
    else:
        px = interpolation(solution["x"],solution["xb"],x_eval)
        _, _, dx = cheb(len(solution["x"])-1)
        xx = solution["x"]
        da = np.diag(1-xx*xx)@(dx@dx)-2*np.diag(xx)@dx
        angular_values = da@solution["values"]
    values = px@solution["values"]@pr.T
    first = px@solution["values"]@dr.T@pr.T
    second = px@solution["values"]@(dr@dr).T@pr.T
    angular = px@angular_values@pr.T
    residual = second*f(r_eval)**2+first*(f(r_eval)*2/r_eval**2)
    residual += angular*(f(r_eval)/r_eval**2)
    residual += values*(omega**2-2*f(r_eval)/r_eval**3)
    residual -= 2*epsilon*omega**2*quadrupole(x_eval)[:,None]*values*f1(r_eval)
    residual -= forcing(r_eval)[None,:]/np.sqrt(2)
    return residual


def manufactured_check():
    """A separately differentiated polynomial checks all PDE signs and factors."""
    omega, epsilon = FREQUENCIES[1], 2.
    def exact(R,x):
        return (chi(0,x)+.2*chi(2,x)+.03*chi(4,x))[:,None]*forcing(R)
    def source(R,x):
        tau=(R-RIN)/(ROUT-RIN)
        first=(32*tau-96*tau**2+64*tau**3)/(ROUT-RIN)
        second=(32-192*tau+192*tau**2)/(ROUT-RIN)**2
        angular=chi(0,x)+.2*chi(2,x)+.03*chi(4,x)
        lap=-1.2*chi(2,x)-.6*chi(4,x)
        radial=f(R)**2*second+f(R)*2/R**2*first+(omega**2-2*f(R)/R**3)*forcing(R)
        return angular[:,None]*radial[None,:]+lap[:,None]*(f(R)/R**2*forcing(R))[None,:] \
            -2*epsilon*omega**2*(quadrupole(x)*angular)[:,None]*(f1(R)*forcing(R))[None,:]
    R,x,w=quadrature()
    expected=exact(R,x)
    errors={}
    for label, solution in (
        ("galerkin_R24_lmax6",galerkin(24,6,omega,epsilon,custom_source=source)),
        ("direct_R24_x12",direct(24,12,omega,epsilon,custom_source=source))):
        errors[label]=norm(evaluate(solution,R,x)-expected,w)/norm(expected,w)
        assert errors[label] < 3e-11
    return errors


def quadrature(nr=100, nx=48):
    r, rw = roots_legendre(nr)
    R = (RIN+ROUT)/2+(ROUT-RIN)*r/2
    rw = rw*(ROUT-RIN)/2/f(R)
    x, xw = roots_legendre(nx)
    return R, x, xw[:, None]*rw[None, :]


def norm(array, weights):
    return float(np.sqrt(np.sum(weights*abs(array)**2)))


def exact_checks():
    import sympy as s
    R, x, e, q, om = s.symbols("R x epsilon q Omega")
    density = s.sqrt((1+e*q)*(1-e*q)**3).series(e, 0, 2).removeO()
    assert s.expand(density-(1-e*q)) == 0
    assert s.series(density/(1+e*q),e,0,2).removeO() == 1-2*e*q
    assert s.series(density/(1-e*q),e,0,2).removeO() == 1
    u = s.Function("u")(R, x)
    ff = 1-2/R
    sphere = lambda v: s.diff((1-x*x)*s.diff(v,x),x)
    box = om**2/ff*(u/R)+s.diff(R**2*ff*s.diff(u/R,R),R)/R**2+sphere(u/R)/R**2
    radial_pde = ff**2*s.diff(u,R,2)+ff*s.diff(ff,R)*s.diff(u,R)
    radial_pde += ff/R**2*sphere(u)+(om**2-2*ff/R**3)*u
    assert s.simplify(R*ff*box-radial_pde) == 0
    angular_errors = []
    xq, wq = roots_legendre(80)
    for lmax in (0,2,4,8,12):
        ells = np.arange(0,lmax+1,2)
        basis = np.array([chi(ell,xq) for ell in ells])
        numeric = (basis*(wq*quadrupole(xq)))@basis.T
        angular_errors.append(float(np.max(abs(numeric-coupling(ells)))))
    assert max(angular_errors) < 3e-12
    return dict(metric_density_series="1-epsilon*q", time_density_series="1-2*epsilon*q",
        spatial_density_series="1", radial_reduction_residual="0",
        coupling_quadrature_max_errors=angular_errors,
        maximum_f1="27/1024 at R=8/3", maximum_metric_fraction_at_epsilon_2=27/512)


def run():
    import sympy
    np.random.seed(417)
    start = time.time()
    report = dict(scope=__doc__, python=platform.python_version(), numpy=np.__version__,
                  scipy=scipy.__version__, sympy=sympy.__version__,
                  source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                  annulus=[RIN,ROUT], frequencies=[[w.real,w.imag] for w in FREQUENCIES],
                  epsilons=list(EPSILONS), exact_checks=exact_checks())
    report["manufactured_solution_relative_errors"] = manufactured_check()
    report["conventions"] = dict(radial_field="u=R Phi", polar_basis="chi_l=sqrt((2l+1)/2) P_l",
        full_sphere_conversion="Y_l0=chi_l/sqrt(2 pi)", radial_boundaries="homogeneous Dirichlet",
        poles="regular; no additional Dirichlet data", forcing="16 t^2(1-t)^2 chi_0; t=(R-2.5)/7.5",
        empirical_comparison_floor=1e-12)
    report["norm_definition"] = "squared norm integral |u|^2 dR dx/f, evaluated by Gauss quadrature (100 radial,48 angular); not a physical energy norm"
    report["conditioning_note"] = "Inverse/condition 1-norm estimates refer only to the unweighted finite nodal matrix and are not continuum resolvent bounds. Response norm ratio tests this one forcing only."
    report["resolution_floor_note"] = "Differences below the observed independent-reference convergence floor must not be claimed as resolved continuum errors."
    re, xe, weights = quadrature()
    source = forcing(re)[None, :]*np.ones((len(xe),1))/np.sqrt(2)
    source_norm = norm(source, weights)
    report["source_norm"] = source_norm
    main_rows, convergence, refinements = [], [], []
    for wi, omega in enumerate(FREQUENCIES):
        for epsilon in EPSILONS:
            print(f"Solving omega={omega}, epsilon={epsilon}", flush=True)
            reference = direct(64, 20, omega, epsilon, condition=True)
            reference_values = evaluate(reference, re, xe)
            reference_norm = norm(reference_values, weights)
            row_base = dict(omega_index=wi, epsilon=epsilon, reference_radial_degree=64,
                reference_angular_degree=20, reference_field_norm=reference_norm,
                source_norm=source_norm, response_norm_ratio=reference_norm/source_norm,
                inverse_1norm_estimate=reference["inverse_1norm_estimate"],
                condition_1norm_estimate=reference["condition_1norm_estimate"],
                reference_backward_error=reference["backward"],
                reference_offgrid_residual_relative=norm(offgrid_residual(reference,re,xe,omega,epsilon),weights)/source_norm)
            for lmax in LMAX:
                coupled = galerkin(64, lmax, omega, epsilon)
                values = evaluate(coupled, re, xe)
                coefficients = coupled["values"]@interpolation(coupled["R"],coupled["Rb"],re).T
                # P2 carries only the last retained even harmonic into lmax+2.
                qcoef = coupling(np.array([lmax,lmax+2]))[1,0]
                omitted = -2*epsilon*omega**2*f1(re)*qcoef*coefficients[-1]
                omitted_values = chi(lmax+2,xe)[:,None]*omitted[None,:]
                residual = offgrid_residual(coupled,re,xe,omega,epsilon)
                main_rows.append(dict(row_base, lmax=lmax,
                    relative_field_error=norm(values-reference_values,weights)/reference_norm,
                    omitted_forcing_norm=norm(omitted_values,weights),
                    relative_omitted_forcing=norm(omitted_values,weights)/source_norm,
                    full_offgrid_residual_relative=norm(residual,weights)/source_norm,
                    full_residual_minus_omitted_relative=norm(residual-omitted_values,weights)/source_norm,
                    galerkin_backward_error=coupled["backward"]))
            # Full resolution records at every frequency/epsilon, not just easy case.
            for nr,nx in ((40,20),(48,20),(56,20),(64,10),(64,14),(64,18)):
                other = direct(nr,nx,omega,epsilon)
                convergence.append(dict(omega_index=wi,epsilon=epsilon,radial_degree=nr,
                    angular_degree=nx,reference_radial_degree=64,reference_angular_degree=20,
                    relative_field_difference=norm(evaluate(other,re,xe)-reference_values,weights)/reference_norm,
                    backward_error=other["backward"]))
            if epsilon in (0.,2.):
                fine = direct(80,24,omega,epsilon)
                refinements.append(dict(omega_index=wi,epsilon=epsilon,
                    radial_degree=80,angular_degree=24,
                    comparison_radial_degree=64,comparison_angular_degree=20,
                    relative_field_difference=norm(evaluate(fine,re,xe)-reference_values,weights)/reference_norm))
    report["comparison"] = main_rows
    report["reference_convergence"] = convergence
    report["additional_reference_refinements"] = refinements
    accepted_convergence = [v["relative_field_difference"] for v in convergence
                           if v["radial_degree"] >= 48]
    assert max(accepted_convergence) < 1e-12
    assert max(v["relative_field_difference"] for v in refinements) < 1e-12
    assert max(v["full_residual_minus_omitted_relative"] for v in main_rows) < 2e-10
    assert max(v["relative_field_error"] for v in main_rows if v["lmax"] == 8) < 1e-8
    for wi in range(2):
        for epsilon in EPSILONS[1:]:
            sequence = [v["relative_field_error"] for v in main_rows
                        if v["omega_index"] == wi and v["epsilon"] == epsilon]
            assert all(b < a for a,b in zip(sequence,sequence[1:]) if a > 1e-10)
    report["acceptance_thresholds"] = dict(reference_relative_difference=1e-12,
        full_residual_minus_omitted_relative=2e-10, lmax8_relative_error=1e-8,
        coupling_quadrature_absolute=3e-12, manufactured_relative=3e-11)
    report["wall_clock_seconds"] = time.time()-start
    report["checks_passed"] = True
    return report


def compare_saved(fresh, saved, location="report"):
    """Compare complete records; do not demand numerical digits below the floor."""
    if isinstance(fresh, dict):
        if not isinstance(saved, dict) or fresh.keys() != saved.keys():
            raise AssertionError("Saved record keys differ at "+location)
        for key in fresh:
            if key not in ("wall_clock_seconds", "python"):
                compare_saved(fresh[key], saved[key], location+"."+key)
    elif isinstance(fresh, list):
        if not isinstance(saved, list) or len(fresh) != len(saved):
            raise AssertionError("Saved array length differs at "+location)
        for i, (a, b) in enumerate(zip(fresh, saved)):
            compare_saved(a, b, location+"["+str(i)+"]")
    elif isinstance(fresh, float):
        if not isinstance(saved, (int, float)) or not np.isfinite(saved):
            raise AssertionError("Saved nonfinite or nonnumeric value at "+location)
        if not np.isclose(fresh, saved, rtol=1e-6, atol=2e-12):
            raise AssertionError("Fresh numerical result differs at "+location)
    elif fresh != saved:
        raise AssertionError("Saved value differs at "+location)


if __name__ == "__main__":
    parser=argparse.ArgumentParser()
    parser.add_argument("--output",type=Path)
    parser.add_argument("--verify-saved",type=Path)
    args=parser.parse_args()
    result=run()
    if args.verify_saved:
        compare_saved(result, json.loads(args.verify_saved.read_text()))
        print("Complete saved-report comparison passed.")
    payload=json.dumps(result,indent=2)
    if args.output:
        args.output.write_text(payload+"\n")
        print(f"Saved {args.output}")
    else:
        print(payload)
