#!/usr/bin/env python3
"""Numerical checks for paper_agentic_racing/agentic_racing.tex."""

import math


def integrate(f, lo: float, hi: float, n: int = 200_000) -> float:
    if n % 2:
        n += 1
    h = (hi - lo) / n
    total = f(lo) + f(hi)
    for i in range(1, n):
        total += (4 if i % 2 else 2) * f(lo + i * h)
    return total * h / 3


def assert_close(name: str, got: float, expected: float, tol: float) -> None:
    if abs(got - expected) > tol:
        raise AssertionError(f"{name}: got {got}, expected {expected}, tol {tol}")
    print(f"PASS {name}: {got:.10g}")


def assert_condition(name: str, condition: bool, detail: str) -> None:
    if not condition:
        raise AssertionError(f"{name}: {detail}")
    print(f"PASS {name}: {detail}")


def gompertz_G(s: float) -> float:
    return math.exp(-math.exp(-s))


def gompertz_g(s: float) -> float:
    return math.exp(-s - math.exp(-s))


def task_density(h: float) -> float:
    return math.exp(-0.20 * h * h) + 0.35 * math.exp(-0.5 * (h - 2.5) ** 2)


def W(q_level: float) -> float:
    return integrate(lambda h: task_density(h) * gompertz_G(q_level - h), -16, 16, 160_000)


def B(q_level: float) -> float:
    return integrate(lambda h: task_density(h) * gompertz_g(q_level - h), -16, 16, 160_000)


def q(C: float, L_inf: float, chi: float, kappa: float, eta: float) -> float:
    return -math.log(L_inf + chi * (C + kappa) ** (-eta))


def q_prime(C: float, L_inf: float, chi: float, kappa: float, eta: float) -> float:
    y = C + kappa
    b = L_inf / chi
    return eta / (y * (1 + b * y**eta))


def q_second(C: float, L_inf: float, chi: float, kappa: float, eta: float) -> float:
    y = C + kappa
    b = L_inf / chi
    return -eta * (1 + (1 + eta) * b * y**eta) / (y**2 * (1 + b * y**eta) ** 2)


def bisect_root(f, lo: float, hi: float, tol: float = 1e-12) -> float:
    flo = f(lo)
    fhi = f(hi)
    if flo == 0:
        return lo
    if fhi == 0:
        return hi
    if flo * fhi > 0:
        raise AssertionError(f"root not bracketed: f({lo})={flo}, f({hi})={fhi}")
    for _ in range(200):
        mid = 0.5 * (lo + hi)
        fmid = f(mid)
        if abs(fmid) < tol or (hi - lo) < tol:
            return mid
        if flo * fmid <= 0:
            hi = mid
            fhi = fmid
        else:
            lo = mid
            flo = fmid
    return 0.5 * (lo + hi)


def main() -> None:
    # 1. Gompertz reliability density integrates to one.
    mass = integrate(gompertz_g, -30, 30, 300_000)
    assert_close("Gompertz boundary kernel integrates to one", mass, 1.0, 1e-8)

    # 2. Boundary statistic equals finite-difference derivative of W(q).
    q_level = 1.3
    h = 2e-4
    fd_W = (W(q_level + h) - W(q_level - h)) / (2 * h)
    assert_close("boundary statistic finite difference", B(q_level), fd_W, 2e-8)

    # 3. Bounded menus saturate in a simple compact-support example.
    bounded_B = integrate(lambda h0: gompertz_g(18.0 - h0), -2.0, 2.0, 80_000)
    assert_condition("bounded-menu saturation", bounded_B < 1e-5, f"B={bounded_B:.6g}")

    # 4. Moving frontier lower bound is positive.
    lower = 0.4 * integrate(gompertz_g, -1.0, 1.0, 100_000)
    assert_condition("moving-frontier lower bound", lower > 0.0, f"lower={lower:.10g}")

    # 5. Scaling derivatives against centered finite differences.
    params = dict(L_inf=0.2, chi=3.0, kappa=1.7, eta=0.35)
    C = 19.0
    step = 1e-4
    fd1 = (q(C + step, **params) - q(C - step, **params)) / (2 * step)
    fd2 = (q(C + step, **params) - 2 * q(C, **params) + q(C - step, **params)) / (step * step)
    assert_close("q prime finite difference", q_prime(C, **params), fd1, 1e-9)
    assert_close("q second finite difference", q_second(C, **params), fd2, 1e-6)

    # 6. Local capability-step and task-value approximation.
    e = 2.1
    eps = 1e-3
    delta_q = q(C + e * eps, **params) - q(C, **params)
    local_delta = e * q_prime(C, **params) * eps
    assert_close("local scaling step", delta_q, local_delta, 5e-8)
    q_current = q(C, **params)
    exact_S = W(q_current + delta_q) - W(q_current)
    local_S = delta_q * B(q_current)
    assert_close("local frontier value", exact_S, local_S, 2e-6)

    # 7. Temporary lead present value.
    r = 0.06
    mu = 0.44
    duration = integrate(lambda t: math.exp(-(r + mu) * t), 0.0, 80.0, 100_000)
    assert_close("temporary lead duration", duration, 1.0 / (r + mu), 1e-10)
    pi = 0.30
    beta = 0.20
    prize = (pi + beta) * exact_S / (r + mu)
    lead_L = pi * exact_S / (r + mu)
    lead_D = beta * exact_S / (r + mu)
    assert_close("private prize equals L plus D", prize, lead_L + lead_D, 1e-14)

    # 8. Symmetric contest FOC against finite-difference derivative.
    gamma = 0.55

    def Lambda(X: float) -> float:
        return 1.0 - math.exp(-gamma * X)

    def Lambda_prime(X: float) -> float:
        return gamma * math.exp(-gamma * X)

    r0 = 0.25
    r1 = 0.08

    def R(X: float) -> float:
        return r0 + r1 * X

    def R_prime(X: float) -> float:
        return r1

    L = 4.0
    D = 1.0

    def payoff(xi: float, xj: float) -> float:
        X = xi + xj
        if X <= 0:
            return 0.0
        sigma = xi / X
        return Lambda(X) * (sigma * L - (1.0 - sigma) * D) - xi * R(X)

    x = 1.4
    fd_payoff = (payoff(x + step, x) - payoff(x - step, x)) / (2 * step)
    foc_formula = (
        0.5 * Lambda_prime(2 * x) * (L - D)
        + Lambda(2 * x) * (L + D) / (4 * x)
        - R(2 * x)
        - x * R_prime(2 * x)
    )
    assert_close("symmetric contest FOC finite difference", foc_formula, fd_payoff, 1e-8)

    # 9. Joint-profit over-racing identity.
    def joint_foc(X: float) -> float:
        return Lambda_prime(X) * (L - D) - (R(X) + X * R_prime(X))

    XJ = bisect_root(joint_foc, 0.01, 20.0)
    private_margin_at_joint = (
        0.5 * Lambda_prime(XJ) * (L - D)
        + Lambda(XJ) * (L + D) / (2 * XJ)
        - R(XJ)
        - 0.5 * XJ * R_prime(XJ)
    )
    threshold_margin = 0.5 * (Lambda(XJ) * (L + D) / XJ - R(XJ))
    assert_close("joint over-racing identity", private_margin_at_joint, threshold_margin, 1e-10)
    assert_condition(
        "joint over-racing sign",
        (private_margin_at_joint > 0) == (Lambda(XJ) * (L + D) / XJ > R(XJ)),
        f"XJ={XJ:.6g}, private_margin={private_margin_at_joint:.6g}",
    )


if __name__ == "__main__":
    main()