"""Representation-Invariant Authority Demand: small numerical verification.

Toy system: xdot = alpha*x + u, x in R^2, ||u||_inf <= rho.
Safe set: h(x)=1-||x||^2 >= 0.  On the boundary, the exact IAD is alpha,
because r=2 alpha and a=2 rho (|x_1|+|x_2|), whose worst ratio is alpha/rho.
The exact boundary feasibility threshold is therefore rho=alpha.
"""
import json
import numpy as np

SEED = 3140
rng = np.random.default_rng(SEED)
alpha = 0.73
N = 200_000
# Uniform angular boundary sample, plus axes so the supremum is represented.
theta = rng.uniform(0, 2*np.pi, N)
theta = np.concatenate([theta, [0, np.pi/2, np.pi, 3*np.pi/2]])
x = np.column_stack([np.cos(theta), np.sin(theta)])

def base_h(z):
    return 1.0 - np.sum(z*z, axis=-1)

def grad_base_h(z):
    return -2.0*z

def transformed(z, name):
    h = base_h(z)
    if name == 'h':
        q, qp = h, np.ones_like(h)
    elif name == '2h':
        q, qp = 2*h, 2*np.ones_like(h)
    elif name == 'h+h^3':
        q, qp = h + h**3, 1 + 3*h*h
    elif name == 'log(1+h)':
        q, qp = np.log1p(h), 1/(1+h)
    elif name == 'h^3 (invalid)':
        q, qp = h**3, 3*h*h
    else:
        raise ValueError(name)
    # On/near the boundary grad(q(h)) = q'(h) grad(h).
    grad = qp[:, None] * grad_base_h(z)
    return q, grad

def iad_for(rho, name):
    _, grad = transformed(x, name)
    f = alpha*x
    r = -np.sum(grad*f, axis=1)
    # support of [-rho,rho]^2: rho * L1 norm of grad
    a = rho*np.sum(np.abs(grad), axis=1)
    d = np.zeros_like(r)
    valid = a > 1e-14
    d[valid] = np.maximum(0, r[valid]) / a[valid]
    d[~valid & (r > 0)] = np.inf
    return float(np.max(d)), float(np.min(-r+a)), float(np.mean(r))

def softsup(values, tau=0.02):
    m = np.max(values)
    return float(m + tau*np.log(np.mean(np.exp((values-m)/tau))))

def run():
    names = ['h', '2h', 'h+h^3', 'log(1+h)', 'h^3 (invalid)']
    rows = {}
    for name in names:
        D, _, mean_r = iad_for(1.0, name)
        # raw boundary outward penalty changes with barrier representation;
        # IAD is a ratio and should not (for valid q'(0)>0).
        rows[name] = {'D_at_rho_1': D, 'raw_mean_r': mean_r}

    # Transition: exact sampled min controlled derivative at each actuator scale.
    scales = np.array([0.50, 0.65, 0.72, 0.73, 0.74, 0.85, 1.00])
    transition = []
    for rho in scales:
        D, min_deriv, _ = iad_for(rho, 'h')
        transition.append({'rho': float(rho), 'D': D,
                           'min_boundary_derivative': min_deriv,
                           'feasible_sampled': bool(min_deriv >= -1e-10)})

    # Empirical soft supremum converges to the exact supremum as sampling grows.
    # Compute pointwise demand for the base representation for a stable diagnostic.
    _, grad = transformed(x, 'h')
    r = -np.sum(grad*(alpha*x), axis=1)
    a = np.sum(np.abs(grad), axis=1)  # rho=1
    d = np.maximum(0, r)/a
    soft = softsup(d)
    # Finite-difference check of the analytic gradient on several boundary points.
    eps = 1e-6
    check = x[:32]
    numeric = np.empty_like(check)
    for j in range(2):
        step = np.zeros_like(check)
        step[:, j] = eps
        numeric[:, j] = (base_h(check + step) - base_h(check - step)) / (2 * eps)
    analytic = grad_base_h(check)
    grad_err = float(np.max(np.abs(numeric - analytic)))

    result = {
        'seed': SEED, 'alpha_exact_threshold': alpha,
        'representation_checks': rows,
        'transition': transition,
        'softsup_tau_0.02_rho_1': soft,
        'max_boundary_gradient_finite_difference_error': grad_err,
        'valid_invariance_relative_spread': float((max(rows[n]['D_at_rho_1'] for n in names[:4]) - min(rows[n]['D_at_rho_1'] for n in names[:4]))/alpha),
        'invalid_h3_D': rows['h^3 (invalid)']['D_at_rho_1'],
        'notes': 'For h^3, q_prime(0)=0, so the defining-function regularity assumption fails and the boundary gradient vanishes.'
    }
    with open('results.json', 'w') as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    run()