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

# Fixed a priori settings. The benchmark's input is 8 (theta, omega, u) tuples;
# the model predicts a constant control for the next target rollout.
DT = 0.05 / 4.0
SUBSTEPS = 32                 # eight benchmark intervals, four integration substeps each
THETA_SAFE = 1.40
K1, K2 = 2.0, 2.0
U_MAX = 1.5

class RolloutSystem(nn.Module):
    def __init__(self, input_shape, shield=False, k1=K1, k2=K2):
        super().__init__()
        self.base = make_model('rnn_small', input_shape, 1)
        self.shield = shield
        self.k1, self.k2 = float(k1), float(k2)

    def project(self, th, om, u_nom):
        # h = theta_safe^2 - theta^2; hdot=-2 theta omega.
        # psi2 = Lf2 h + LgLf h*u + k1*Lf h + k2*psi1 >= 0.
        h = THETA_SAFE**2 - th * th
        hd = -2.0 * th * om
        # hessian contribution: -2*om^2; drift contribution from pendulum.
        drift = -2.0 * om * om + 2.0 * th * (9.81 / 10.0) * torch.sin(th)
        psi1 = hd + self.k1 * h
        b = drift + self.k1 * hd + self.k2 * psi1
        a = -4.0 * th                  # Lg Lf h
        # Projection onto [-U_MAX,U_MAX] intersecting a*u+b >= 0.
        unclipped = torch.clamp(u_nom, -U_MAX, U_MAX)
        req = -b
        # scalar exact projection; where a is nearly zero, retain bounded action.
        candidate = torch.where(a.abs() > 1e-7, req / a, unclipped)
        candidate = torch.clamp(candidate, -U_MAX, U_MAX)
        feasible = (a * candidate + b >= 0)
        # If actuator bounds make the constraint infeasible, safest available bound.
        safe = torch.where(a > 0, torch.full_like(candidate, U_MAX),
                           torch.full_like(candidate, -U_MAX))
        safe = torch.where(a.abs() > 1e-7, safe, unclipped)
        candidate = torch.where(feasible, candidate, safe)
        # The min/max composition gives an exact scalar QP solution while remaining
        # differentiable almost everywhere for train_model's autograd.
        return torch.where((a * unclipped + b >= 0), unclipped, candidate)

    def forward(self, x):
        # rnn_small accepts flattened windows in the bench implementation.
        u_nom = torch.tanh(self.base(x)).reshape(-1) * U_MAX
        th = x[:, -3]
        om = x[:, -2]
        for _ in range(SUBSTEPS):
            u = self.project(th, om, u_nom) if self.shield else u_nom
            om = om + (-9.81 / 10.0 * torch.sin(th) - 0.25 * om + 2.0 * u) * DT
            th = th + om * DT
        return th.reshape(-1, 1)

def train_one(seed, lr, shield, k1=K1, k2=K2, epochs=12):
    torch.manual_seed(seed); np.random.seed(seed)
    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
    model = RolloutSystem(ds['input_shape'], shield=shield, k1=k1, k2=k2)
    _, metric, _ = train_model(model, ds, epochs=epochs, lr=lr, batch=128, log=lambda *_: None)
    return float(metric) if metric is not None else float('inf')

def run_cfg(cfg, shield):
    return lambda seed: train_one(seed, cfg['lr'], shield, cfg.get('k1', K1), cfg.get('k2', K2))

def signature(seed, lr, k1, k2):
    # Measure trained-model behavior, not an analytic toy identity.
    torch.manual_seed(seed); np.random.seed(seed)
    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
    m = RolloutSystem(ds['input_shape'], shield=True, k1=k1, k2=k2)
    m, _, _ = train_model(m, ds, epochs=12, lr=lr, batch=128, log=lambda *_: None)
    m.eval(); x = ds['xte'].to(next(m.parameters()).device)
    with torch.no_grad():
        u_nom = torch.tanh(m.base(x)).reshape(-1)*U_MAX
        th=x[:,-3]; om=x[:,-2]; raw=[]; sh=[]
        for _ in range(SUBSTEPS):
            h=THETA_SAFE**2-th*th; hd=-2*th*om
            drift=-2*om*om+2*th*(9.81/10)*torch.sin(th)
            psi1=hd+k1*h; b=drift+k1*hd+k2*psi1; a=-4*th
            raw.append((a*u_nom+b).cpu().numpy())
            u=m.project(th,om,u_nom); sh.append((a*u+b).cpu().numpy())
            om=om+(-9.81/10*torch.sin(th)-.25*om+2*u)*DT; th=th+om*DT
    raw=np.concatenate(raw); sh=np.concatenate(sh)
    predicted = float(np.mean(raw < 0))
    observed = float(np.mean(sh < -1e-5))
    return {'trained_model_raw_constraint_violation_rate':predicted,
            'trained_model_shielded_constraint_violation_rate':observed,
            'projection_activation_rate':float(np.mean(np.abs(raw-sh)>1e-5)),
            'predicted_vs_observed': {'predicted': predicted, 'observed': observed},
            'confirmed': bool(observed <= 0.01 and predicted > observed + 0.02)}

def main():
    # Union of all idea learning rates is included in baseline sweep: parity.
    lrs=[1e-3,3e-3,1e-2]
    baseline_grid=[{'lr':lr} for lr in lrs]
    base=sweep_baseline(lambda c: run_cfg(c, False), baseline_grid, seeds=tuple(range(4)))
    # Re-evaluate the selected baseline on all eight paired seeds.
    base['full']=evaluate(run_cfg(base['best_cfg'], False), seeds=tuple(range(8)))
    # Same three-config budget for idea; includes baseline-best and nearby settings.
    idea_cfg=[{'lr':lr,'k1':k1,'k2':k2} for lr,k1,k2 in [(base['best_cfg']['lr'],2.,2.),(1e-3,2.,2.),(1e-2,2.,2.)]]
    idea_trials=[]
    for c in idea_cfg:
        r=evaluate(run_cfg(c, True), seeds=tuple(range(4)))
        idea_trials.append({'cfg':c,'mean':r['mean']})
    best_idea_cfg=min(idea_cfg, key=lambda c: next(t['mean'] for t in idea_trials if t['cfg']==c))
    idea=evaluate(run_cfg(best_idea_cfg, True), seeds=tuple(range(8)))
    rep=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,signature(0,best_idea_cfg['lr'],2.,2.))
    rep['idea']['sweep']=idea_trials; rep['custom_track']=None
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()