import json
import numpy as np

EPS = 1e-6


def admissibility(U, eps=EPS):
    U = np.asarray(U)
    D = U[..., 0]
    m = U[..., 1:-1]
    E = U[..., -1]
    return (D >= eps) & (E - np.sqrt(D * D + np.sum(m * m, axis=-1)) >= eps)


def g(U, s):
    U = np.asarray(U)
    s = np.asarray(s)
    # Supports a shared s=(d,) or batched s=(N,d).
    if s.ndim == 1:
        return U[..., -1] - s[0] * U[..., 0] - np.sum(s[1:] * U[..., 1:-1], axis=-1)
    return U[..., -1] - s[..., 0] * U[..., 0] - np.sum(s[..., 1:] * U[..., 1:-1], axis=-1)


def limit_residual(U0, U1, eps=EPS, iterations=60):
    """Exact scalar residual limiter via bisection on the relativistic margin.

    The admissible set is convex, so with valid U0 and invalid U1 there is one
    boundary crossing on [0,1]. This is a robust reference implementation of
    the scalar limiter; theta is detached because this NumPy MVP has no grads.
    """
    U0 = np.asarray(U0, dtype=float)
    U1 = np.asarray(U1, dtype=float)
    delta = U1 - U0
    def margin(t):
        V = U0 + t[..., None] * delta
        return V[..., -1] - np.sqrt(np.sum(V[..., :-1] * V[..., :-1], axis=-1))
    theta = np.ones(U0.shape[:-1])
    bad = (U1[..., 0] < eps) | (margin(np.ones(U0.shape[:-1])) < eps)
    lo = np.zeros_like(theta)
    hi = np.ones_like(theta)
    # Only bad rows need a root; valid rows remain at theta=1.
    for _ in range(iterations):
        mid = 0.5 * (lo + hi)
        ok = margin(mid) >= eps
        lo = np.where(ok, mid, lo)
        hi = np.where(ok, hi, mid)
    theta = np.where(bad, lo, theta)
    # Density is linear and can be the active constraint independently.
    dbad = U1[..., 0] < eps
    td = np.clip((U0[..., 0] - eps) / np.maximum(U0[..., 0] - U1[..., 0], 1e-30), 0., 1.)
    theta = np.where(dbad, np.minimum(theta, td), theta)
    # Move infinitesimally inside both closed constraints to avoid roundoff.
    theta = np.where(theta < 1.0, np.maximum(0.0, theta * (1.0 - 1e-12)), theta)
    V = U0 + theta[..., None] * delta
    # A direct safety correction handles ulp-scale endpoint errors.
    V[..., 0] = np.maximum(V[..., 0], eps)
    return V, theta

def component_clip(U, eps=EPS):
    """Control: independently clip density and energy, leaving momentum unchanged."""
    V = np.asarray(U, dtype=float).copy()
    V[..., 0] = np.maximum(V[..., 0], eps)
    V[..., -1] = np.maximum(V[..., -1], eps)
    return V


def main(seed=1234):
    rng = np.random.default_rng(seed)

    # Verify min_{||s||<=1} g_s(U) = E - sqrt(D^2+||m||^2).
    x = rng.normal(size=(10000, 3))
    x /= np.linalg.norm(x, axis=1, keepdims=True)
    U = rng.uniform(0.1, 3.0, size=(10000, 4))
    U[:, 1:3] = rng.normal(0.0, 1.0, size=(10000, 2))
    exact = U[:, -1] - np.linalg.norm(U[:, :-1], axis=1)
    sampled = np.min(np.stack([g(U, x[i]) for i in range(len(x))], axis=1), axis=1)
    # The sampled minimum is only an approximation; test the analytic worst normal.
    w = U[:, :-1]
    sstar = w / np.linalg.norm(w, axis=1, keepdims=True)
    analytic = np.array([g(U[i], sstar[i]) for i in range(len(U))])
    identity_abs_err = float(np.max(np.abs(exact - analytic)))
    sampled_gap_mean = float(np.mean(sampled - exact))

    # Valid baselines, then deliberately corrupt learned candidates.
    n = 20000
    U0 = rng.uniform(0.2, 2.0, size=(n, 4))
    U0[:, 1:3] = rng.normal(0.0, 0.7, size=(n, 2))
    U0[:, -1] = np.linalg.norm(U0[:, :-1], axis=1) + rng.uniform(0.05, 1.5, size=n)
    # Errors are large enough to produce both density and coupled energy/momentum violations.
    U1 = U0 + rng.normal(0.0, 0.9, size=U0.shape)
    U1[:, 0] -= rng.uniform(0.0, 1.3, size=n)
    U1[:, -1] -= rng.uniform(0.0, 0.9, size=n)

    safe, theta = limit_residual(U0, U1)
    clipped = component_clip(U1)
    valid0 = admissibility(U0).mean()
    valid1 = admissibility(U1).mean()
    valid_safe = admissibility(safe).mean()
    valid_clip = admissibility(clipped).mean()
    # Retention is the fraction of the learned residual preserved.
    dnorm = np.linalg.norm(U1 - U0, axis=1)
    retained = np.linalg.norm(safe - U0, axis=1) / np.maximum(dnorm, 1e-12)
    clip_retained = np.linalg.norm(clipped - U0, axis=1) / np.maximum(dnorm, 1e-12)
    result = {
        "seed": seed,
        "identity_max_abs_error": identity_abs_err,
        "random_sphere_mean_gap_vs_exact": sampled_gap_mean,
        "baseline_valid_rate": float(valid0),
        "candidate_valid_rate": float(valid1),
        "gql_safe_valid_rate": float(valid_safe),
        "component_clip_valid_rate": float(valid_clip),
        "gql_mean_theta": float(theta.mean()),
        "gql_fraction_limited": float(np.mean(theta < 1.0 - 1e-12)),
        "gql_mean_residual_retention": float(retained.mean()),
        "component_clip_mean_residual_retention": float(clip_retained.mean()),
        "gql_min_margin": float(np.min(safe[:, -1] - np.linalg.norm(safe[:, :-1], axis=1))),
        "component_clip_min_margin": float(np.min(clipped[:, -1] - np.linalg.norm(clipped[:, :-1], axis=1))),
    }
    print(json.dumps(result, indent=2, sort_keys=True))
    with open("gql_results.json", "w") as f:
        json.dump(result, f, indent=2, sort_keys=True)


if __name__ == "__main__":
    main()