"""Complex-step/central-difference cosmology derivative validation for Row 19.

This implements the explicit flat-w0waCDM reproducibility benchmark printed in
main.tex: p=(Omega_m,h,w0,wa), z={0.1,0.3,0.6,1.0,1.5,2.0}, observables
{H,D_A,D_L}.  It validates derivative plumbing.  It does *not* claim that the
larger curvature-aware compressed vector {R,l_A,D_M/r_d,D_H/r_d} is fully
implemented here.
"""
from __future__ import annotations

import numpy as np
from scipy.integrate import quad

C_KM_S = 299792.458
Z_GRID = np.asarray([0.1, 0.3, 0.6, 1.0, 1.5, 2.0], dtype=float)
PIVOT = np.asarray([0.315, 0.674, -1.0, 0.0], dtype=float)
PARAMETERS = ["Omega_m", "h", "w0", "wa"]
OBSERVABLES = ["H", "D_A", "D_L"]


def E(z, p):
    Om, _h, w0, wa = p
    zp1 = 1.0 + z
    f_de = zp1 ** (3.0 * (1.0 + w0 + wa)) * np.exp(-3.0 * wa * z / zp1)
    return np.sqrt(Om * zp1**3 + (1.0 - Om) * f_de)


def H(z, p):
    return 100.0 * p[1] * E(z, p)


def D_c(z, p):
    # scipy.quad with complex_func=True preserves the imaginary complex-step signal.
    integral, _ = quad(
        lambda zp: 1.0 / E(zp, p),
        0.0,
        float(z),
        epsabs=1.0e-12,
        epsrel=1.0e-12,
        limit=300,
        complex_func=True,
    )
    return C_KM_S / (100.0 * p[1]) * integral


def D_A(z, p):
    return D_c(z, p) / (1.0 + z)


def D_L(z, p):
    return (1.0 + z) * D_c(z, p)


FUNCTIONS = {"H": H, "D_A": D_A, "D_L": D_L}


def complex_step_jacobian(fn, eps: float = 1.0e-30):
    J = np.empty((len(Z_GRID), len(PIVOT)), dtype=float)
    for i, z in enumerate(Z_GRID):
        for j in range(len(PIVOT)):
            pc = PIVOT.astype(complex)
            pc[j] += 1j * eps
            J[i, j] = float(np.imag(fn(float(z), pc)) / eps)
    return J


def central_difference_jacobian(fn, delta: float = 1.0e-5):
    J = np.empty((len(Z_GRID), len(PIVOT)), dtype=float)
    for i, z in enumerate(Z_GRID):
        for j in range(len(PIVOT)):
            pp = PIVOT.copy()
            pm = PIVOT.copy()
            pp[j] += delta
            pm[j] -= delta
            fp = np.real(fn(float(z), pp))
            fm = np.real(fn(float(z), pm))
            J[i, j] = float((fp - fm) / (2.0 * delta))
    return J


def validate(tolerance: float = 1.0e-3) -> dict:
    results = {}
    overall = 0.0
    finite = True
    for name in OBSERVABLES:
        fn = FUNCTIONS[name]
        Jcs = complex_step_jacobian(fn)
        Jcd = central_difference_jacobian(fn)
        denom = np.maximum(1.0e-16, np.abs(Jcs))
        rel = np.abs(Jcd - Jcs) / denom
        max_overall = float(np.max(rel))
        max_by_param = {
            p: float(np.max(rel[:, j])) for j, p in enumerate(PARAMETERS)
        }
        finite = finite and bool(
            np.all(np.isfinite(Jcs)) and np.all(np.isfinite(Jcd)) and np.all(np.isfinite(rel))
        )
        overall = max(overall, max_overall)
        results[name] = {
            "max_relative_error": max_overall,
            "max_relative_error_by_parameter": max_by_param,
            "complex_step_jacobian": Jcs.tolist(),
            "central_difference_jacobian": Jcd.tolist(),
        }
    checks = {
        "all_derivatives_finite": bool(finite),
        "complex_step_vs_central_difference_below_tolerance": bool(overall <= tolerance),
    }
    return {
        "pivot": dict(zip(PARAMETERS, [float(x) for x in PIVOT])),
        "redshift_grid": [float(z) for z in Z_GRID],
        "observables": OBSERVABLES,
        "complex_step_epsilon": 1.0e-30,
        "central_difference_step": 1.0e-5,
        "acceptance_max_relative_error": tolerance,
        "overall_max_relative_error": overall,
        "results": results,
        "checks": checks,
        "pass": all(checks.values()),
    }
