from pathlib import Path
import json, math, random, sys
import numpy as np
import torch
from torch import nn
sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import train_model, evaluate, sweep_baseline, make_report

META = {"name": "sos_backup_shield_policy", "domain": "dynamics", "description": "Torque-limited pendulum action regression with a learned policy wrapped by a quadratic finite-horizon backup shield."}

# Polynomial pendulum-like local model: theta_dot=omega, omega_dot=u.
# This is the explicit quadratic certificate used as the finite-horizon SOS-style MVP.
P = np.array([[1.25, .25], [.25, .375]], dtype=np.float64)
U_MAX = 1.0
TH_MAX, OM_MAX = 1.4, 2.0
CERT_RADIUS = 0.12


def math_check():
    A = np.array([[0., 1.], [-2., -2.]])
    Q = -(A.T @ P + P @ A)
    identity_error = float(np.max(np.abs(Q - np.eye(2))))
    lam = float(np.max(np.linalg.eigvalsh(P)))
    # Boundary samples of V=1, where -Vdot/V should be >= 1/lambda_max(P).
    vals = []
    max_u = 0.
    for a in np.linspace(0, 2*np.pi, 4096, endpoint=False):
        z = np.array([np.cos(a), np.sin(a)])
        x = CERT_RADIUS * z / np.sqrt(z @ P @ z)
        vals.append(float(x @ Q @ x))
        max_u = max(max_u, abs(-2*x[0] - 2*x[1]))
    observed = min(vals) / (CERT_RADIUS**2)
    predicted = 1.0 / lam
    return {"identity_error": identity_error, "observed_min_neg_vdot_over_V": observed,
            "predicted_bound": predicted, "max_backup_u_on_V1": max_u,
            "certificate_actuator_feasible": bool(max_u <= U_MAX)}


def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(seed)
    def make(n):
        th = rng.uniform(-1.35, 1.35, n)
        om = rng.uniform(-2.0, 2.0, n)
        # Stabilizing target action, clipped to the actuator polytope.
        u = np.clip(-1.8 * th - 0.65 * om, -1., 1.)
        return np.stack([th, om], 1).astype("float32"), u[:, None].astype("float32")
    xtr, ytr = make(n_train); xte, yte = make(n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1,
            "input_shape": (2,)}


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)


def tensors(d):
    return {**d, **{k: torch.as_tensor(d[k]) for k in ("xtr", "ytr", "xte", "yte")}}


class Policy(nn.Module):
    """Same base MLP for both systems; shield is the sole intervention."""
    def __init__(self, shield=False, margin=0.0):
        super().__init__()
        self.shield, self.margin = shield, float(margin)
        self.net = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))

    def forward(self, x):
        nominal = self.net(x)
        if not self.shield:
            return torch.clamp(nominal, -U_MAX, U_MAX)
        # h=TH_MAX^2-theta^2; hdot+alpha h >= 0 gives admissible interval.
        th, om = x[:, 0:1], x[:, 1:2]
        h = TH_MAX**2 - th.square()
        alpha = 0.8
        # hdot=-2 theta omega; hdot + alpha*h + 2 theta*u >= margin.
        # For theta near zero, use the Lyapunov backup endpoint directly.
        lo = (self.margin + 2*th*om - alpha*h) / (2*th + 1e-6)
        hi = lo
        # Solve interval robustly by evaluating the affine barrier inequality.
        rhs = self.margin + 2*th*om - alpha*h
        pos_lo = rhs / (2*th + 1e-6)
        # inequality 2 theta u >= rhs; interval endpoint is lower for theta>0,
        # upper for theta<0. Pair with backup action and certificate gate.
        backup = torch.clamp(-2.0*th - 2.0*om, -U_MAX, U_MAX)
        feasible = (th.abs() > 0.04) & (rhs.abs() <= 2*th.abs()*U_MAX)
        safe_nom = torch.clamp(nominal, -U_MAX, U_MAX)
        # project toward the safe side, then choose backup outside the certificate.
        projected = torch.where(th > 0, torch.maximum(safe_nom, pos_lo), torch.minimum(safe_nom, pos_lo))
        V = 1.25*th.square() + .5*th*om + .375*om.square()
        inside = V <= CERT_RADIUS**2
        return torch.clamp(torch.where(feasible & inside, projected, backup), -U_MAX, U_MAX)


def train_metric(seed, cfg, keep=False):
    seed_all(seed); d = get_dataset(seed)
    model = Policy(shield=cfg.get("margin", 0.) > 0, margin=cfg.get("margin", 0.))
    model, metric, hist = train_model(model, tensors(d), epochs=cfg["epochs"], lr=cfg["lr"], batch=128, log=lambda *_: None)
    return (metric, model, d) if keep else metric


def train_fn(cfg):
    return lambda seed: train_metric(seed, cfg)


def mechanism_signature(cfg, seeds=tuple(range(8))):
    # Measured on outputs of trained models, not an analytic-only calculation.
    shifts, dissipation, interventions, accepted = [], [], [], []
    for s in seeds:
        _, model, d = train_metric(s, cfg, keep=True)
        model = model.cpu()
        model.eval(); x = torch.as_tensor(d["xte"])
        with torch.no_grad():
            nominal = model.net(x); shielded = model(x)
        th, om = x[:,0:1], x[:,1:2]
        V = 1.25*th.square()+.5*th*om+.375*om.square()
        u = shielded
        # Euler certificate derivative for actual trained shield output.
        vdot = th*om + (0.5*th + .75*om)*u
        backup = torch.clamp(-2*th-2*om, -1., 1.)
        shifts.append(float(torch.mean(torch.abs(shielded-nominal)).item()))
        dissipation.append(float(torch.mean(vdot[V <= CERT_RADIUS**2]).item()))
        interventions.append(float(torch.mean(torch.abs(shielded-backup)).item()))
        accepted.append(float(torch.mean(((V <= CERT_RADIUS**2) & (torch.abs(shielded-nominal)<1e-5)).float()).item()))
    # Prediction: inside the certified set the shield should be non-increasing;
    # observed mean derivative should be non-positive on trained-model outputs.
    obs = float(np.mean(dissipation))
    return {"prediction": {"certificate_vdot_mean_on_inside": "<=0", "backup_intervention_nonzero": True},
            "observed_from_trained_models": {"mean_vdot_inside": obs, "mean_action_shift": float(np.mean(shifts)),
              "mean_distance_to_backup": float(np.mean(interventions)), "mean_nominal_acceptance": float(np.mean(accepted)), "n_models": len(seeds)},
            "confirmed": bool(obs <= 1e-6 and float(np.mean(interventions)) > 1e-4)}


def main():
    check = math_check()
    print(json.dumps({"math_check": check}))
    if check["identity_error"] > 1e-10 or not check["certificate_actuator_feasible"]:
        raise RuntimeError("certificate sanity check failed")
    epochs = 15
    lrs = [1e-3, 3e-3, 6e-3]
    # Union parity: baseline sees every lr used by idea; method knob margin is swept
    # on the idea side and zero-margin standard clipping is baseline.
    base_grid = [{"lr": lr, "margin": 0.0, "epochs": epochs} for lr in lrs]
    idea_grid = [{"lr": lr, "margin": margin, "epochs": epochs} for lr, margin in zip(lrs, [0.02, 0.05, 0.10])]
    base = sweep_baseline(train_fn, base_grid)
    idea_sweep = [{"cfg": c, "mean": evaluate(train_fn(c), seeds=(0,1,2,3))["mean"]} for c in idea_grid]
    best = min(idea_sweep, key=lambda z: z["mean"])["cfg"]
    idea = evaluate(train_fn(best))
    report = make_report("sos_backup_shield_policy", "local_mlp_tiny", base, idea,
                         mechanism_signature(best))
    report["idea"]["selection_sweep"] = idea_sweep
    report["custom_track"] = {"name": META["name"], "file": "sos_shield_bench.py", "domain": "dynamics"}
    report["math_check"] = check
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()