#!/usr/bin/env python3
"""CAffNet (arXiv:2605.24437) claim [4]: safety-critical CONTROL. CAffNet's exact affine projection acts
as a control-barrier-function (CBF) safety filter -- the robot provably avoids the obstacle (0 violations)
while reaching the goal, whereas an unconstrained / soft-penalty baseline collides.

2D single-integrator robot dx = u dt. Nominal controller u_nom = k(goal - x). Obstacle: disk center c,
radius r. Safety margin h(x)=||x-c||^2 - r^2 (>=0 safe). CBF condition ḣ + alpha h >= 0 gives the AFFINE
control constraint  a(x)^T u >= b(x),  a(x)=2(x-c),  b(x) = -alpha h(x). CAffNet projects u_nom onto this
half-space exactly (closed form); the soft baseline adds a repulsion penalty but is not guaranteed safe.
Deterministic seeds.
"""
import numpy as np, json, hashlib

def simulate(controller, x0, goal, c, r, alpha=1.0, dt=0.02, T=1500, umax=2.0):
    x = x0.copy(); traj = [x.copy()]; min_h = np.inf; reached = False
    for _ in range(T):
        u = controller(x, goal, c, r, alpha, umax)
        x = x + dt * u; traj.append(x.copy())
        h = float(np.sum((x - c) ** 2) - r ** 2); min_h = min(min_h, h)
        if np.linalg.norm(x - goal) < 0.1: reached = True; break
    return np.array(traj), min_h, reached

def u_nominal(x, goal, c, r, alpha, umax):
    u = 2.0 * (goal - x); n = np.linalg.norm(u); return u if n <= umax else u * umax / n

def u_caffnet(x, goal, c, r, alpha, umax):
    """Project the nominal control onto the CBF safe half-space a^T u >= b (exact affine projection)."""
    u = u_nominal(x, goal, c, r, alpha, umax)
    a = 2.0 * (x - c); h = float(np.sum((x - c) ** 2) - r ** 2); b = -alpha * h
    slack = a @ u - b
    if slack < 0:                                          # violates CBF -> project onto {u: a^T u = b}
        u = u + (b - a @ u) / (a @ a + 1e-12) * a
    n = np.linalg.norm(u); return u if n <= umax else u * umax / n

def u_soft(x, goal, c, r, alpha, umax):
    """Soft baseline: nominal control + a repulsion penalty (gradient of a barrier); NOT guaranteed safe."""
    u = 2.0 * (goal - x)
    d2 = float(np.sum((x - c) ** 2)); rep = (x - c) / (d2 ** 2 + 1e-6)   # repulsion, decays with distance
    u = u + 0.05 * rep
    n = np.linalg.norm(u); return u if n <= umax else u * umax / n

def main():
    R = {"claim": "CAffNet_safety_critical_control", "paper": "arXiv:2605.24437"}
    rng = np.random.default_rng(0)
    c = np.array([1.0, 0.0]); r = 0.4                       # obstacle disk between start and goal

    # sweep start/goal pairs that force the path THROUGH the obstacle region
    caff_safe = 0; caff_reached = 0; soft_safe = 0; soft_reached = 0; nom_safe = 0
    caff_min_h = []; soft_min_h = []; N = 40
    for s in range(N):
        ang = rng.uniform(0, 2 * np.pi)
        x0 = c + np.array([np.cos(ang), np.sin(ang)]) * 1.6
        goal = c - np.array([np.cos(ang), np.sin(ang)]) * 1.6      # goal opposite the obstacle -> path crosses it
        _, mh_c, rc = simulate(u_caffnet, x0, goal, c, r); caff_min_h.append(mh_c)
        _, mh_s, rs = simulate(u_soft, x0, goal, c, r); soft_min_h.append(mh_s)
        _, mh_n, rn = simulate(u_nominal, x0, goal, c, r)
        caff_safe += (mh_c >= -1e-6); caff_reached += rc
        soft_safe += (mh_s >= -1e-6); soft_reached += rs
        nom_safe += (mh_n >= -1e-6)
    R["caffnet_safe_frac"] = round(caff_safe / N, 3); R["caffnet_reached_frac"] = round(caff_reached / N, 3)
    R["soft_safe_frac"] = round(soft_safe / N, 3); R["soft_reached_frac"] = round(soft_reached / N, 3)
    R["nominal_safe_frac"] = round(nom_safe / N, 3)
    R["caffnet_min_safety_margin"] = round(float(np.min(caff_min_h)), 5)      # >= 0 => never entered obstacle
    R["soft_min_safety_margin"] = round(float(np.min(soft_min_h)), 5)

    R["caffnet_always_safe"] = caff_safe == N                                 # 0 violations by construction
    R["caffnet_reaches_goal"] = caff_reached >= 0.9 * N
    R["baseline_collides"] = soft_safe < N                                    # soft baseline violates sometimes
    R["caffnet_strictly_safer_than_baseline"] = (caff_safe / N) > (soft_safe / N)

    R["verdict"] = "supports" if (R["caffnet_always_safe"] and R["caffnet_reaches_goal"]
                                  and R["baseline_collides"]) else "inconclusive"

    print("claim: " + R["claim"])
    print(f"2D navigation through an obstacle disk (center {c.tolist()}, r={r}); {N} start/goal pairs crossing it.")
    print(f"  CAffNet (CBF affine projection): safe={R['caffnet_safe_frac']}, reached goal={R['caffnet_reached_frac']}, "
          f"min safety margin={R['caffnet_min_safety_margin']} (>=0 => never enters obstacle)")
    print(f"  Soft-penalty baseline:            safe={R['soft_safe_frac']}, reached={R['soft_reached_frac']}, min margin={R['soft_min_safety_margin']}")
    print(f"  Unconstrained nominal:            safe={R['nominal_safe_frac']}")
    print(f"  -> CAffNet always safe: {R['caffnet_always_safe']}; reaches goal: {R['caffnet_reaches_goal']}; baseline collides: {R['baseline_collides']}")
    print(f"verdict: {R['verdict']}")

    def _np(o):
        if isinstance(o, np.bool_): return bool(o)
        if isinstance(o, np.integer): return int(o)
        if isinstance(o, np.floating): return float(o)
        raise TypeError
    import os; os.makedirs("outputs", exist_ok=True)
    open("outputs/control_results.json", "w").write(json.dumps(R, indent=2, default=_np))
    print("RESULTS_SHA256=" + hashlib.sha256(json.dumps(R, sort_keys=True, default=_np).encode()).hexdigest())
    return 0 if R["verdict"] == "supports" else 1

if __name__ == "__main__":
    raise SystemExit(main())