#!/usr/bin/env python3
"""
benchmarks.py
=============
Original three benchmark families accompanying the revised technical edition of
"First-Principles Validation and Extension of the S-I-C-T Category-Theoretic
Reconstruction" (Roth Complexity Lab working paper, 21 October 2026).

Executable research source code. Run locally with:

    python -m pip install numpy
    python benchmarks.py

Requires Python 3 and NumPy (runtime approx. 15 s). Benchmarks:
  1. Shift vs Rule 90 (N = 6, 7, 8)  -- exact finite-state arithmetic
  2. Harmonic oscillator integrators (explicit/implicit Euler, Verlet, exact)
  3. Lorenz-63 RK4 divergence horizons and floating-point ulp merging

Benchmark 1 results are exact and architecture-independent. Floating-point
ulp merger times and rounding drift (Benchmarks 2-3) can depend on the runtime,
processor and NumPy version; the paper notes this explicitly.
"""
import json
import math

import numpy as np

RESULTS = {}


# ---------------------------------------------------------------------------
# Benchmark 1: shift conjugacy and Rule 90 erasure (Table 3)
# ---------------------------------------------------------------------------
def rule90(state):
    n = len(state)
    return tuple(state[(i - 1) % n] ^ state[(i + 1) % n] for i in range(n))


def shift(state):
    n = len(state)
    return tuple(state[(i + 1) % n] for i in range(n))


def entropy_bits(dist):
    h = 0.0
    for p in dist.values():
        if p > 0:
            h -= p * math.log2(p)
    return h


bench1 = {}
for N in (6, 7, 8):
    size = 1 << N
    states = [tuple(int(b) for b in format(i, f"0{N}b")) for i in range(size)]

    # Shift is a permutation (exhaustive conjugacy check).
    shift_images = [shift(s) for s in states]
    assert len(set(shift_images)) == size

    # Rule 90 multiplicities.
    r90_images = [rule90(s) for s in states]
    distinct = len(set(r90_images))
    multiplicity = size // distinct

    # Erasure under the uniform ensemble.
    mu = {s: 1.0 / size for s in states}
    series = [entropy_bits(mu)]
    for _ in range(12):
        new_mu = {}
        for s, p in mu.items():
            img = rule90(s)
            new_mu[img] = new_mu.get(img, 0.0) + p
        mu = new_mu
        series.append(entropy_bits(mu))
    total_erased = series[0] - series[-1]

    bench1[f"N={N}"] = {
        "shift_conjugate": True,
        "rule90_distinct_images": distinct,
        "multiplicity": f"{multiplicity}-to-1",
        "entropy_series_bits": series,
        "total_erased_bits": total_erased,
    }

# Table 3 expected totals: N=6 -> 2, N=7 -> 1, N=8 -> 8.
assert abs(bench1["N=6"]["total_erased_bits"] - 2) < 1e-12
assert abs(bench1["N=7"]["total_erased_bits"] - 1) < 1e-12
assert abs(bench1["N=8"]["total_erased_bits"] - 8) < 1e-12
RESULTS["benchmark1_shift_rule90"] = bench1


# ---------------------------------------------------------------------------
# Benchmark 2: harmonic oscillator (Table 4)
# ---------------------------------------------------------------------------
dt = 0.01
x0, p0, E0 = 1.0, 0.0, 0.5
horizons = [10**4, 10**5, 10**6]


def explicit_euler(x, p, n):
    for _ in range(n):
        x, p = x + dt * p, p - dt * x
    return x, p


def implicit_euler(x, p, n):
    c = 1.0 / (1.0 + dt * dt)
    for _ in range(n):
        x, p = c * (x + dt * p), c * (p - dt * x)
    return x, p


def verlet(x, p, n):
    for _ in range(n):
        p = p - 0.5 * dt * x
        x = x + dt * p
        p = p - 0.5 * dt * x
    return x, p


def exact_rotation(x, p, n):
    c, s = math.cos(dt), math.sin(dt)
    for _ in range(n):
        x, p = c * x + s * p, -s * x + c * p
    return x, p


bench2 = {}
for name, integrator in [
    ("explicit_euler", explicit_euler),
    ("implicit_euler", implicit_euler),
    ("stormer_verlet", verlet),
    ("exact_rotation_fp64", exact_rotation),
]:
    rec = {}
    for n in horizons:
        x, p = integrator(x0, p0, n)
        energy = 0.5 * (x * x + p * p)
        ref = (math.cos(n * dt), -math.sin(n * dt))
        traj_err = math.hypot(x - ref[0], p - ref[1])
        rec[f"n={n}"] = {
            "energy_error": energy - E0,
            "trajectory_error": traj_err,
        }
    bench2[name] = rec

# Explicit Euler energy drift matches the analytic (1 + dt^2)^n - 1 (Table 4).
for n in horizons:
    analytic = ((1.0 + dt * dt) ** n - 1.0) * E0
    measured = bench2["explicit_euler"][f"n={n}"]["energy_error"]
    assert abs(measured - analytic) / max(1.0, abs(analytic)) < 1e-6

RESULTS["benchmark2_oscillator"] = bench2


# ---------------------------------------------------------------------------
# Benchmark 3: Lorenz-63 RK4 (Table 5) + ulp merging (Finding 4.6)
# ---------------------------------------------------------------------------
SIGMA, RHO, BETA = 10.0, 28.0, 8.0 / 3.0


def lorenz_rhs(state):
    x, y, z = state
    return np.array([SIGMA * (y - x), x * (RHO - z) - y, x * y - BETA * z])


def rk4_step(state, h):
    k1 = lorenz_rhs(state)
    k2 = lorenz_rhs(state + 0.5 * h * k1)
    k3 = lorenz_rhs(state + 0.5 * h * k2)
    k4 = lorenz_rhs(state + h * k3)
    return state + (h / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)


h = 1e-3

# Transient to reference state uref (t = 50 from (1,1,1)).
state = np.array([1.0, 1.0, 1.0])
for _ in range(50000):
    state = rk4_step(state, h)
uref = state  # approximately (-4.40848, -3.32421, 24.10480)

bench3 = {
    "uref": uref.tolist(),
    "ulp_x": float(np.spacing(float(uref[0]))),
    "note": "ulp merger times are environment-dependent (Finding 4.6).",
}


def divergence_horizon(delta0):
    a = uref.copy()
    b = uref.copy()
    b[0] += delta0
    steps = 0
    max_steps = 200000  # t = 200
    while steps < max_steps:
        a = rk4_step(a, h)
        b = rk4_step(b, h)
        steps += 1
        if np.linalg.norm(a - b) >= 1.0:
            break
    return steps * h  # tcrit in time units


for delta0 in (1e-6, 1e-9, 1e-12):
    bench3[f"delta0={delta0}"] = {"measured_tcrit": divergence_horizon(delta0)}


def ulp_merge_steps(k):
    a = uref.copy()
    b = uref.copy()
    b[0] += k * float(np.spacing(float(uref[0])))
    for step in range(1, 10001):
        a = rk4_step(a, h)
        b = rk4_step(b, h)
        if np.array_equal(a, b):
            return step
    return None


for k in (1, 2, 3, 4, 8, 16):
    bench3[f"ulp_k={k}"] = {"first_coincidence_step": ulp_merge_steps(k)}

RESULTS["benchmark3_lorenz63"] = bench3


# ---------------------------------------------------------------------------
# Emit results
# ---------------------------------------------------------------------------
print(json.dumps(RESULTS, indent=2, default=str))
with open("original_results_clean.json", "w") as fh:
    json.dump(RESULTS, fh, indent=2, default=str)

print("\nAll original benchmark assertions passed.")