import json
import itertools
import time
import numpy as np
from scipy.optimize import minimize


def constraints(r, R, B):
    r = np.asarray(r)
    return np.array([r[0]-B[0], -r[0]-B[0], r[1]-B[1], -r[1]-B[1], r@r-R*R])


def shield(z, R, B, tol=1e-10):
    """Exact Euclidean projection onto disk intersected with an axis-aligned box.
    Candidates are the KKT active-set candidates for all subsets in 2D.
    """
    z = np.asarray(z, dtype=float)
    candidates = [(z.copy(), ())]
    # One box face: orthogonal projection onto the face.
    faces = [(0, 1), (0, -1), (1, 1), (1, -1)]
    for i, sign in faces:
        rr = z.copy(); rr[i] = sign * B[i]
        candidates.append((rr, (i, sign)))
    # Disk alone (radial KKT candidate).
    nz = np.linalg.norm(z)
    if nz > R:
        candidates.append((z * (R / nz), ('disk',)))
    # Two box faces (corners).
    for i, s1 in faces:
        for j, s2 in faces:
            if i < j:
                rr = z.copy(); rr[i] = s1 * B[i]; rr[j] = s2 * B[j]
                candidates.append((rr, (i, s1, j, s2)))
    # Disk/box active intersections (the finite KKT roots).
    for i, sign in faces:
        c = sign * B[i]
        rem = R*R-c*c
        if rem >= -tol:
            rem = max(0., rem)
            for y in (np.sqrt(rem), -np.sqrt(rem)):
                rr = np.array([0., 0.]); rr[i] = c; rr[1-i] = y
                candidates.append((rr, (i, sign, 'disk')))
    feasible = [(0.5*np.sum((rr-z)**2), rr, a)
                for rr, a in candidates if np.max(constraints(rr, R, B)) <= tol]
    if not feasible:
        # Conservative fallback is the box center, which is feasible whenever the
        # modeled actuator set is nonempty.
        rr = np.zeros(2)
        return rr, 0.5*np.sum((rr-z)**2), True
    obj, rr, active = min(feasible, key=lambda t: t[0])
    return rr, obj, False


def pg_projection(z, R, B, steps=300, alpha=.08):
    """A deliberately standard alternating projected-gradient feasibility repair."""
    r = z.copy()
    for _ in range(steps):
        r = np.clip(r, -B, B)
        n = np.linalg.norm(r)
        if n > R: r *= R/n
    return r


def reference(z, R, B):
    # SLSQP is only used for numerical verification, not the shield.
    fun = lambda r: .5*np.sum((r-z)**2)
    cons = ({'type':'ineq','fun':lambda r: R*R-r@r},
            {'type':'ineq','fun':lambda r: B-r},
            {'type':'ineq','fun':lambda r: B+r})
    q = minimize(fun, np.clip(z, -B, B), constraints=cons, method='SLSQP',
                 options={'ftol':1e-12, 'maxiter':300})
    return q.x, q.fun, q.success


def run():
    rng = np.random.default_rng(2906)
    B = np.array([1.35, 1.10])
    direction = np.array([.8, .6])
    direction /= np.linalg.norm(direction)
    # Prediction 1: every command inside the feasible set is unchanged.
    interior_err = []
    ref_err = []
    max_violation = 0.
    for _ in range(300):
        R = rng.uniform(.65, 1.8)
        z = rng.uniform(-1.5, 1.5, 2)
        r, _, _ = shield(z, R, B)
        if np.max(constraints(z, R, B)) <= 0:
            interior_err.append(np.linalg.norm(r-z))
        rr, _, ok = reference(z, R, B)
        if ok: ref_err.append(np.linalg.norm(r-rr))
        max_violation = max(max_violation, max(0., np.max(constraints(r, R, B))))
    # Prediction 2: along a ray the transition is at min(box boundary, disk boundary).
    R = 1.0
    predicted_transition = min(B[0]/direction[0], B[1]/direction[1], R)
    scales = np.linspace(.1, 1.7, 161)
    errors = []
    violations = []
    for s in scales:
        z = s*direction; r, _, _ = shield(z, R, B)
        errors.append(np.linalg.norm(r-z)); violations.append(max(0., np.max(constraints(r,R,B))))
    observed_transition = scales[np.argmax(np.asarray(errors) > 1e-8)]
    # Prediction 3: changing disk radius shifts the transition linearly while disk is limiting.
    radii = np.linspace(.45, 1.0, 12)
    measured = []
    for rad in radii:
        es = []
        for s in scales:
            rr, _, _ = shield(s*direction, rad, B); es.append(np.linalg.norm(rr-s*direction))
        measured.append(scales[np.argmax(np.asarray(es)>1e-8)])
    slope = np.polyfit(radii[:8], measured[:8], 1)[0]
    # Mini comparison on random commands: exact shield vs raw and iterative repair.
    zs = rng.uniform(-2.0,2.0,(500,2)); Rs = rng.uniform(.65,1.25,500)
    raw_v=[]; shield_v=[]; pg_v=[]; shield_dist=[]; pg_dist=[]
    t0=time.perf_counter()
    for z,rad in zip(zs,Rs):
        rr,_,_=shield(z,rad,B); shield_v.append(max(0.,np.max(constraints(rr,rad,B))))
        shield_dist.append(np.linalg.norm(rr-z)); raw_v.append(max(0.,np.max(constraints(z,rad,B))))
    shield_ms=1000*(time.perf_counter()-t0)/len(zs)
    t0=time.perf_counter()
    for z,rad in zip(zs,Rs):
        rr=pg_projection(z,rad,B); pg_v.append(max(0.,np.max(constraints(rr,rad,B)))); pg_dist.append(np.linalg.norm(rr-z))
    pg_ms=1000*(time.perf_counter()-t0)/len(zs)
    out={
      'predictions': {
        'interior_identity_max_error': float(max(interior_err)),
        'predicted_identity_error': 0.0,
        'transition_predicted': float(predicted_transition),
        'transition_observed_grid': float(observed_transition),
        'transition_abs_error': float(abs(observed_transition-predicted_transition)),
        'radius_transition_slope_predicted': 1.0,
        'radius_transition_slope_observed': float(slope),
      },
      'math_check': {'max_constraint_violation':float(max_violation), 'max_error_vs_SLSQP':float(max(ref_err)), 'n_reference':len(ref_err)},
      'comparison': {
        'raw_mean_violation':float(np.mean(raw_v)), 'shield_mean_violation':float(np.mean(shield_v)), 'pg_mean_violation':float(np.mean(pg_v)),
        'shield_mean_distance':float(np.mean(shield_dist)), 'pg_mean_distance':float(np.mean(pg_dist)),
        'shield_ms_per_command':shield_ms, 'pg_ms_per_command':pg_ms
      }, 'seed':2906}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': run()