"""Residual-tightened safety shield: math check and reproducible toy experiment."""
import json
from pathlib import Path
import numpy as np

SEED = 544
SIGMA_R = 0.20
KAPPA = 1.0
BOUND = 1.0
ACTION_LO, ACTION_HI = -0.5, 0.4
EMA_DECAY = 0.90


def normalized_residual(d, ema, sigma=SIGMA_R):
    return float(d / (sigma + ema))


def update_ema(ema, d, decay=EMA_DECAY):
    return float(decay * ema + (1.0 - decay) * d)


def nominal_model(x, a):
    return 0.85 * x + a


def shield_projection(x, a_policy, residual_r, kappa=KAPPA):
    """Exact QP projection for h(f(x,a))+kappa*r = x+a-BOUND+kappa*r <= 0."""
    # The feasible action interval is a <= BOUND-x-kappa*r.
    upper = min(ACTION_HI, BOUND - x - kappa * residual_r)
    return float(np.clip(a_policy, ACTION_LO, upper))


def policy(x):
    # Nominally feasible controller; the shifted dynamics can push it unsafe.
    return float(np.clip(0.75 - x, ACTION_LO, ACTION_HI))


def math_check():
    # Constant-error regimes have EMA[d] -> d, hence r -> d/(sigma+d).
    ds = np.array([0.0, 0.02, 0.05, 0.20, 0.80])
    rs = np.array([normalized_residual(d, d) for d in ds])
    margins = KAPPA * rs
    # Also check the exact projection satisfies the tightened constraint.
    xs = np.linspace(-0.2, 1.0, 31)
    max_violation = 0.0
    for x in xs:
        for r in rs:
            a = shield_projection(x, policy(x), r)
            max_violation = max(max_violation, x + a - BOUND + KAPPA*r)
    return {
        "d": ds.tolist(), "r_at_ema_equal_d": rs.tolist(),
        "tightening_margin": margins.tolist(),
        "r_monotone": bool(np.all(np.diff(rs) >= -1e-12)),
        "max_projection_constraint_violation": float(max_violation),
        "constraint_satisfied": bool(max_violation <= 1e-10),
    }


def run_condition(regime, method, rng, episodes=160, horizon=50):
    """Simulate true x_next=x+a+bias+noise while the model assumes x+a."""
    bias = 0.03 if regime == "nominal" else 0.50
    # Fixed margin control is deliberately the margin needed by the shifted regime.
    fixed_r = normalized_residual(0.50, 0.50)
    violations = 0
    total_dev = 0.0
    total_reward = 0.0
    min_margin = []
    for _ in range(episodes):
        x = 0.20
        ema = 0.05  # modest prior scale; prevents an artificial first-step infinite residual
        # The shield starts with a nominal-model residual prior; shifted error is learned online.
        r_prev = normalized_residual(0.03, 0.03)
        for _ in range(horizon):
            ap = policy(x)
            if method == "unshielded":
                a = float(np.clip(ap, ACTION_LO, ACTION_HI))
            elif method == "fixed":
                a = shield_projection(x, ap, fixed_r)
            elif method == "adaptive":
                r = normalized_residual(ema, ema)  # unused only for initialization below
                # r is the online residual estimate from prior measurements; EMA is a scale.
                # We retain the last normalized residual in the loop state via r_prev.
                a = shield_projection(x, ap, r_prev)
            else:
                raise ValueError(method)
            # Shifted true dynamics; small sensor/process noise makes violations nontrivial.
            noise = float(rng.normal(0.0, 0.004))
            xn = nominal_model(x, a) + bias + noise
            d = abs(xn - nominal_model(x, a))
            if xn > BOUND:
                violations += 1
            total_dev += abs(a - ap)
            # Reward preserves the requested behavior: approach the goal, penalize intervention.
            total_reward += -(0.75 - xn) ** 2 - 0.05 * (a - ap) ** 2
            ema = update_ema(ema, d)
            r_new = min(0.8, normalized_residual(d, ema))
            if method == "adaptive":
                r_prev = r_new
            x = xn
        # for unshielded/fixed no r_prev is needed; initialize adaptive before first action
    n = episodes * horizon
    return {"violation_rate": violations / n, "mean_action_deviation": total_dev / n,
            "mean_reward": total_reward / n}


def run_all():
    results = {}
    for regime in ("nominal", "shifted"):
        for method in ("unshielded", "fixed", "adaptive"):
            rng = np.random.default_rng(SEED + (0 if regime == "nominal" else 100) +
                                        {"unshielded": 0, "fixed": 10, "adaptive": 20}[method])
            results[f"{regime}_{method}"] = run_condition(regime, method, rng)
    return results


if __name__ == "__main__":
    out = {"seed": SEED, "sigma_r": SIGMA_R, "kappa": KAPPA,
           "math_check": math_check(), "experiment": run_all()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))