import os, sys, json, random
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 make_report, sweep_baseline, evaluate
from bench.custom_tracks.robust_cbf_pendulum_policy import get_dataset

META = {'name': 'barrier_certified_pendulum_policy', 'domain': 'dynamics',
        'description': 'Pendulum state-to-action regression trained with differentiable zeroing-CBF residual hinge loss.'}
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
ALPHA, EPS, LAMBDA = 1.0, 0.08, 2.0

class Policy(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2,32), nn.Tanh(), nn.Linear(32,32), nn.Tanh(), nn.Linear(32,1))
    def forward(self,x): return torch.clamp(self.net(x), -1., 1.)

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

def dynamics(x,u):
    # normalized pendulum: theta_dot=omega, omega_dot=sin(theta)+u
    th, om = x[:,0:1], x[:,1:2]
    return torch.cat([om, torch.sin(th)+u], 1)

def residual(x,u):
    # Smooth ellipsoidal safe set: h=1-theta^2/1.4^2-omega^2/2^2.
    # Its omega gradient is nonzero, so the relative-degree-one residual depends on u.
    th, om=x[:,0:1], x[:,1:2]
    h=1.0-(th/1.4)**2-(om/2.0)**2
    gh=torch.cat([-2*th/(1.4**2), -2*om/(2.0**2)], 1)
    return (gh*dynamics(x,u)).sum(1,keepdim=True) + ALPHA*h

def train_metric(seed,cfg,keep=False):
    seed_all(seed)
    d=get_dataset(seed)
    xtr=torch.as_tensor(d['xtr'],device=DEVICE); ytr=torch.as_tensor(d['ytr'],device=DEVICE)
    model=Policy().to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'])
    bs=128; model.train()
    try:
        for _ in range(cfg['epochs']):
            perm=torch.randperm(len(xtr),device=DEVICE)
            for ii in perm.split(bs):
                x=xtr[ii]; target=ytr[ii]; u=model(x)
                task=(u-target).pow(2).mean()
                if cfg.get('barrier',False):
                    # one-step differentiable rollout states, including training states
                    r=residual(x,u)
                    barrier=torch.relu(EPS-r).pow(2).mean()
                    loss=task+cfg.get('weight',LAMBDA)*barrier
                else: loss=task
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10); opt.step()
        model.eval()
        with torch.no_grad(): metric=float(((model(torch.as_tensor(d['xte'],device=DEVICE))-torch.as_tensor(d['yte'],device=DEVICE))**2).mean().item())
    except Exception:
        if DEVICE=='cuda':
            torch.cuda.empty_cache(); return train_metric_cpu(seed,cfg,keep)
        raise
    return (metric,model.cpu(),d) if keep else metric

def train_metric_cpu(seed,cfg,keep=False):
    global DEVICE; old=DEVICE; DEVICE='cpu'
    try: return train_metric(seed,cfg,keep)
    finally: DEVICE=old

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

def signature(cfg,seeds=tuple(range(8))):
    observed=[]; gaps=[]; violations=[]
    for s in seeds:
        _,m,d=train_metric(s,cfg,True); x=torch.as_tensor(d['xte']); u=m(x)
        r=residual(x,u).detach().numpy().ravel()
        observed.append(float(r.min())); violations.append(float(np.mean(r<EPS)))
        # local finite-difference Lipschitz estimate on trained model residual
        xx=torch.as_tensor(np.linspace(-1.35,1.35,257)[:,None], dtype=torch.float32).repeat(1,2)
        xx[:,1]=0; rr=residual(xx,m(xx)).detach().numpy().ravel()
        gaps.append(float(np.max(np.abs(np.diff(rr))/(2.7/256))))
    L=float(np.mean(gaps)); delta=0.5*(2.7/256)
    lower=EPS-L*delta
    return {'prediction':{'lower_bound': 'epsilon - L_r delta', 'delta':delta},
            'observed_from_trained_models':{'min_residual_mean':float(np.mean(observed)),
              'violation_fraction_mean':float(np.mean(violations)), 'L_r_estimate_mean':L,
              'bound_lower_mean':lower, 'n_models':len(seeds)},
            'confirmed': bool(float(np.mean(observed)) >= lower-1e-3)}

def main():
    epochs=15; lrs=[1e-3,3e-3,6e-3]
    # union parity: all learning rates appear on both sides; baseline also sweeps its sole central knob (lr).
    base_grid=[{'lr':lr,'epochs':epochs,'barrier':False} for lr in lrs]
    idea_grid=[{'lr':lr,'epochs':epochs,'barrier':True,'weight':w} for lr,w in zip(lrs,[1.,2.,4.])]
    base=sweep_baseline(train_fn,base_grid)
    sel=[{'cfg':c,'mean':evaluate(train_fn(c),seeds=(0,1,2,3))['mean']} for c in idea_grid]
    best=min(sel,key=lambda z:z['mean'])['cfg']; idea=evaluate(train_fn(best))
    report=make_report('robust_cbf_pendulum_policy','shared_mlp_32',base,idea,signature(best))
    report['idea']['selection_sweep']=sel
    report['custom_track']={'name':META['name'],'file':'barrier_bench.py','domain':'dynamics'}
    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))

if __name__=='__main__': main()