import json, math
from pathlib import Path
import numpy as np

EPS = 0.03
R_OBS = 0.60
D_SAFE = 0.25
K1 = 2.0
K2 = 2.0
AMAX = 3.0


def barrier(p, v, c=np.zeros(2), k1=K1, k2=K2):
    q = np.asarray(p, float) - c
    rho = math.sqrt(float(q @ q) + EPS**2)
    h = rho - R_OBS - D_SAFE
    g = q / rho
    H = np.eye(2) / rho - np.outer(q, q) / rho**3
    hd = float(g @ v)
    psi1 = hd + k1 * h
    # psi2 = g @ a + b, because p_dot=v and v_dot=a
    b = float(v @ H @ v + k1 * hd + k2 * psi1)
    return h, psi1, g, H, b


def project_box_halfspace(a_nom, g, b, amax=AMAX):
    """Exact small 2-D Euclidean projection onto box and g.a+b >= 0."""
    a_nom = np.asarray(a_nom, float)
    box = np.clip(a_nom, -amax, amax)
    if float(g @ box + b) >= -1e-12:
        return box, True, False
    req = -b
    cand = []
    gg = float(g @ g)
    z = a_nom + (req - float(g @ a_nom)) / gg * g
    if np.all(z <= amax + 1e-10) and np.all(z >= -amax - 1e-10):
        cand.append(z)
    for i in range(2):
        j = 1 - i
        for xi in (-amax, amax):
            if abs(g[j]) > 1e-12:
                zj = (req - g[i] * xi) / g[j]
                if -amax - 1e-10 <= zj <= amax + 1e-10:
                    z = np.zeros(2); z[i] = xi; z[j] = zj; cand.append(z)
    for x in (-amax, amax):
        for y in (-amax, amax):
            z = np.array([x, y], float)
            if float(g @ z) >= req - 1e-10: cand.append(z)
    if not cand:
        return box, False, True
    best = min(cand, key=lambda z: float((z-a_nom) @ (z-a_nom)))
    return best, True, True


def shield(p, v, a_nom, amax=AMAX, k1=K1, k2=K2):
    h, psi1, g, H, b = barrier(p, v, k1=k1, k2=k2)
    a, feasible, active = project_box_halfspace(a_nom, g, b, amax)
    return a, feasible, active, h, psi1, float(g @ a + b)


def derivative_check(seed=4):
    rng = np.random.default_rng(seed)
    errs_h, errs_psi = [], []
    dt = 1e-6
    for _ in range(100):
        p = rng.normal(size=2); p *= 1.3 / np.linalg.norm(p)
        v = rng.normal(size=2)
        a = rng.normal(size=2)
        h, psi, g, H, b = barrier(p, v)
        hp = barrier(p + dt*v, v + dt*a)[0]
        psip = barrier(p + dt*v, v + dt*a)[1]
        errs_h.append(abs((hp-h)/dt - float(g@v)))
        errs_psi.append(abs((psip-psi)/dt - float(g@a + v@H@v + K1*float(g@v))))
    return {"max_abs_hdot_error": float(max(errs_h)), "max_abs_psi1dot_error": float(max(errs_psi))}


def rollout(dt, shielded, scale, T=4.0):
    # Fixed nominal policy drives through the obstacle; shield should brake/deflect.
    p = np.array([-2.0, 0.0]); v = np.zeros(2); c = np.zeros(2)
    min_h = 1e9; activ = 0; infeas = 0; proj = 0.0; n = int(T/dt)
    for _ in range(n):
        a_nom = np.array([scale*2.2, 0.0]) - 0.8*v
        h, psi, g, H, b = barrier(p, v, c)
        min_h = min(min_h, h)
        if shielded:
            a, feasible, active, _, _, residual = shield(p,v,a_nom)
            activ += int(active); infeas += int(not feasible)
            proj += float(np.linalg.norm(a-a_nom))
        else:
            a = np.clip(a_nom, -AMAX, AMAX)
        # Semi-implicit Euler, matching the continuous double-integrator model.
        v = v + dt*a; p = p + dt*v
    min_h = min(min_h, barrier(p,v,c)[0])
    return {"min_h": float(min_h), "violation": float(max(0,-min_h)),
            "activation_rate": activ/n, "infeasible_rate": infeas/n,
            "mean_projection": proj/n}


def activation_sweep():
    out=[]
    p=np.array([-0.9,0.0]); v=np.array([0.0,0.0])
    for scale in np.linspace(0, 8, 17):
        a_nom=np.array([scale,0.0])
        a, feasible, active, h, psi, res=shield(p,v,a_nom)
        out.append({"scale":float(scale), "active":bool(active), "feasible":bool(feasible),
                    "projection":float(np.linalg.norm(a-a_nom)), "residual":float(res)})
    return out


def gain_sweep():
    # At a fixed near-boundary state, required normal acceleration grows with k2*psi1.
    p=np.array([-0.86,0.0]); v=np.array([-0.15,0.0])
    rows=[]
    for k2 in [0.0, 0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 24.0]:
        h,psi,g,H,b=barrier(p,v,k1=K1,k2=k2)
        req=-b
        rows.append({"k2":k2,"h":h,"psi1":psi,"normal_required":float(req),
                     "predicted_linear_term":float(-k2*psi), "feasible_under_amax": bool(req <= AMAX + 1e-12)})
    return rows


def main():
    results={"derivative_check":derivative_check(), "activation_sweep":activation_sweep(),
             "gain_sweep":gain_sweep(), "activation_threshold_exact": float(barrier(np.array([-0.9,0.0]), np.zeros(2))[4]), "rollouts":{}}
    for scale in [0.8,1.2,1.6,2.0]:
        results["rollouts"][str(scale)]={}
        for dt in [0.02,0.01,0.005]:
            results["rollouts"][str(scale)][str(dt)]={
                "unshielded":rollout(dt,False,scale), "shielded":rollout(dt,True,scale)}
    Path("results.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    main()