"""Einstein-Hilbert/Litim FRG benchmark used by UTFANSWF Row 19.

Implements the beta functions printed as Eqs. (124)-(125) in
Frank Saueressig, arXiv:2302.14152.  The numerical fixed point below is
obtained by solving the printed equations directly; it is not copied from
the numerical values printed later in that source.
"""
from __future__ import annotations

import numpy as np


def eta_N(g, lam):
    u = 1.0 - 2.0 * lam
    numerator = g * (5.0 / u - 9.0 / (u * u) - 7.0)
    denominator = 3.0 * np.pi * (
        1.0 + g / (12.0 * np.pi) * (5.0 / u - 6.0 / (u * u))
    )
    return numerator / denominator


def beta(x):
    x = np.asarray(x)
    g, lam = x[0], x[1]
    u = 1.0 - 2.0 * lam
    eta = eta_N(g, lam)
    beta_g = (2.0 + eta) * g
    beta_lam = -(2.0 - eta) * lam + g / (8.0 * np.pi) * (
        20.0 / u - 16.0 - (5.0 / 3.0) * eta / u
    )
    return np.asarray([beta_g, beta_lam])


def jacobian_complex_step(x, eps: float = 1.0e-30) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    J = np.empty((2, 2), dtype=float)
    for j in range(2):
        xc = x.astype(complex)
        xc[j] += 1j * eps
        J[:, j] = np.imag(beta(xc)) / eps
    return J


def solve_fixed_point(
    x0=(0.8759, 0.1566), tol: float = 1.0e-13, max_iter: int = 30
):
    x = np.asarray(x0, dtype=float)
    converged = False
    iterations = 0
    for iterations in range(1, max_iter + 1):
        b = np.real(beta(x)).astype(float)
        if float(np.max(np.abs(b))) <= tol:
            converged = True
            break
        J = jacobian_complex_step(x)
        step = np.linalg.solve(J, -b)
        x = x + step
        if not np.all(np.isfinite(x)) or abs(1.0 - 2.0 * x[1]) < 1.0e-8:
            break
    b = np.real(beta(x)).astype(float)
    if float(np.max(np.abs(b))) <= tol:
        converged = True
    return converged, x, b, iterations


def benchmark_result() -> dict:
    converged, x_star, b_star, iterations = solve_fixed_point()
    J = jacobian_complex_step(x_star)
    eig = np.linalg.eigvals(J)
    theta = -eig
    order = np.argsort(np.imag(theta))
    theta = theta[order]

    expected = {
        "gstar_direct_printed_equations": 0.875888555049176,
        "lambdastar_direct_printed_equations": 0.156604266763380,
        "theta_real_direct_printed_equations": 1.86654707333725,
        "theta_imag_abs_direct_printed_equations": 2.35256720924354,
    }

    metrics = {
        "converged": bool(converged),
        "iterations": int(iterations),
        "gstar": float(x_star[0]),
        "lambdastar": float(x_star[1]),
        "beta_at_star": [float(v) for v in b_star],
        "max_abs_beta": float(np.max(np.abs(b_star))),
        "stability_matrix": [[float(v) for v in row] for row in J],
        "critical_exponents": [
            {"re": float(np.real(t)), "im": float(np.imag(t))} for t in theta
        ],
        "expected_direct_equation_solution": expected,
        "regression_abs_error": {
            "gstar": abs(float(x_star[0]) - expected["gstar_direct_printed_equations"]),
            "lambdastar": abs(float(x_star[1]) - expected["lambdastar_direct_printed_equations"]),
            "theta_real": max(
                abs(float(np.real(t)) - expected["theta_real_direct_printed_equations"])
                for t in theta
            ),
            "theta_imag_abs": max(
                abs(abs(float(np.imag(t))) - expected["theta_imag_abs_direct_printed_equations"])
                for t in theta
            ),
        },
    }

    checks = {
        "non_gaussian_fixed_point": bool(
            converged and x_star[0] > 0.0 and abs(x_star[1]) > 0.0
        ),
        "beta_residual_below_1e_12": bool(np.max(np.abs(b_star)) <= 1.0e-12),
        "direct_equation_regression_matches": bool(
            max(metrics["regression_abs_error"].values()) <= 5.0e-10
        ),
        "two_uv_relevant_directions": bool(
            len(theta) == 2 and all(float(np.real(t)) > 0.0 for t in theta)
        ),
        "complex_conjugate_pair": bool(
            abs(float(np.real(theta[0]) - np.real(theta[1]))) <= 1.0e-10
            and abs(float(np.imag(theta[0]) + np.imag(theta[1]))) <= 1.0e-10
        ),
    }
    return {"metrics": metrics, "checks": checks, "pass": all(checks.values())}
