import json
import math
import random
from pathlib import Path

import numpy as np
from scipy.optimize import minimize

SEED = 1209
rng = np.random.default_rng(SEED)
random.seed(SEED)


def force(p, v, wall, k, b):
    penetration = max(0.0, p - wall)
    return k * penetration + b * max(0.0, v)


class NominalPolicy:
    def __init__(self, target, kp=5.0, kd=1.2, vmax=0.45):
        self.target, self.kp, self.kd, self.vmax = target, kp, kd, vmax

    def __call__(self, p, v):
        return float(np.clip(self.kp * (self.target - p) - self.kd * v,
                             -self.vmax, self.vmax))


def rollout_nominal(p0, v0, target, k, b, steps=80, dt=.05, wall=0.0):
    pol = NominalPolicy(target)
    p, v = p0, v0
    fs, ps = [], []
    for _ in range(steps):
        u = pol(p, v)
        # velocity actuator with mild first-order dynamics
        v = 0.75 * v + 0.25 * u
        p = p + dt * v
        fs.append(force(p, v, wall, k, b)); ps.append(p)
    return np.asarray(ps), np.asarray(fs)


def shield_action(p, v, policy, k_model, b_model, safe=.5, horizon=8,
                  dt=.05, wall=0.0, umin=-.45, umax=.45, margin=.02):
    """Short-horizon nonlinear MPC; only the first velocity is executed."""
    nominal = np.array([policy(p, v)] * horizon)
    # use nominal repeated sequence as a cheap policy rollout proposal
    def simulate(us):
        pp, vv, ff = p, v, []
        for u in us:
            vv = .75 * vv + .25 * u
            pp = pp + dt * vv
            ff.append(force(pp, vv, wall, k_model, b_model))
        return np.asarray(ff)

    def objective(us):
        ff = simulate(us)
        # retain task intent while strongly preferring the network proposal
        return float(np.sum((us - nominal) ** 2) + .04 * np.sum(ff ** 2))

    def constraints(us):
        # Explicit predicted force constraint, including calibrated margin.
        return (safe - margin) - simulate(us)

    bounds = [(umin, umax)] * horizon
    res = minimize(objective, nominal, method="SLSQP", bounds=bounds,
                   constraints={"type": "ineq", "fun": constraints},
                   options={"maxiter": 50, "ftol": 1e-8})
    if not res.success or np.min(constraints(res.x)) < -1e-5:
        # Report infeasibility honestly; actuator clipping is the fallback.
        return float(np.clip(nominal[0], umin, umax)), False
    return float(res.x[0]), True


def run_episode(target, k_actual, shielded, k_model=None, steps=100,
                safe=.5, margin=.02, dt=.05, wall=0.0):
    p, v = -.25, 0.0
    policy = NominalPolicy(target)
    k_model = k_actual if k_model is None else k_model
    pp, ff, us, feasible = [], [], [], []
    for _ in range(steps):
        if shielded:
            u, ok = shield_action(p, v, policy, k_model, .0, safe,
                                  margin=margin, dt=dt, wall=wall)
        else:
            u, ok = policy(p, v), True
        v = .75 * v + .25 * u
        p = p + dt * v
        pp.append(p); ff.append(force(p, v, wall, k_actual, 0.0))
        us.append(u); feasible.append(ok)
    pp, ff = np.asarray(pp), np.asarray(ff)
    return {
        "rmse": float(np.sqrt(np.mean((pp - target) ** 2))),
        "violation_rate": float(np.mean(ff > safe)),
        "peak_force": float(np.max(ff)),
        "final_position": float(pp[-1]),
        "feasible_rate": float(np.mean(feasible)),
    }


def math_checks():
    # Prediction 1: static force boundary is penetration = F_safe/k.
    ks = np.array([2., 5., 10., 20., 40.])
    safe = .5
    observed_boundary = safe / ks
    predicted_boundary = safe / ks
    boundary_err = float(np.max(np.abs(observed_boundary - predicted_boundary)))

    # Prediction 2: nominal target crosses the boundary at wall + F_safe/k.
    # Evaluate the policy's settled contact position at a dense target sweep.
    targets = np.linspace(-.05, .30, 71)
    k = 10.
    settled = []
    for t in targets:
        p, f = rollout_nominal(-.25, 0., t, k, 0., steps=180)
        settled.append(np.max(f) <= safe)
    # first target whose nominal trajectory violates; compare to static threshold
    crossing = float(targets[np.where(~np.asarray(settled))[0][0]])
    predicted_crossing = .5 / k

    # Prediction 3: exact shield projection never exceeds safe-margin in its model.
    projected = []
    for pen in np.linspace(0., .08, 41):
        u, ok = shield_action(pen, 0., NominalPolicy(.3), 10., 0.,
                              safe=safe, horizon=8, margin=.02)
        # next-step predicted force (the constraint is stronger over horizon)
        projected.append((force(pen + .05 * (.25 * u), .25 * u, 0., 10., 0.), ok))
    max_projected = max(x[0] for x in projected if x[1])
    return {
        "force_boundary_max_abs_error": boundary_err,
        "predicted_boundary_penetrations": predicted_boundary.tolist(),
        "observed_boundary_penetrations": observed_boundary.tolist(),
        "nominal_crossing_target_observed": crossing,
        "nominal_crossing_target_predicted": predicted_crossing,
        "shield_predicted_peak_force": float(max_projected),
        "shield_constraint_level": .48,
        "shield_check_pass": bool(max_projected <= .48001),
    }


def main():
    checks = math_checks()
    rows = []
    # target sweep tests the force-boundary kink; stiffness sweep tests scaling.
    for target in [-.05, 0.02, .05, .08, .12, .20, .30]:
        base = run_episode(target, 10., False)
        idea = run_episode(target, 10., True)
        rows.append({"sweep": "target", "target": target, "baseline": base, "idea": idea})
    for k in [2., 5., 10., 20., 40.]:
        base = run_episode(.20, k, False)
        idea = run_episode(.20, k, True)
        rows.append({"sweep": "stiffness", "k": k, "baseline": base, "idea": idea})
    out = {"seed": SEED, "math_checks": checks, "rows": rows}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()