import json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

SEED = 2823
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Nonlinear controllable discrete plant. The task asks for a nonzero action,
# whereas the certificate asks for a stable closed loop around x*=0.
def plant(x, u):
    x1, x2 = x[:, :1], x[:, 1:2]
    return torch.cat((x1 + .05*(-x1**3 + u + .5*x2),
                      .95*x2 + .05*x1), dim=1)

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):
        # Enforce the equilibrium condition pi(0)=0 exactly.
        z=torch.zeros(1,2,device=x.device)
        return 2.0*(torch.tanh(self.net(x)/2.0)-torch.tanh(self.net(z)/2.0))

class Lyapunov(nn.Module):
    # V=||g(x)-g(0)||^2 + eps ||x||^2, hence V(0)=0 and V>=eps||x||^2.
    def __init__(self, eps=.08):
        super().__init__(); self.eps=eps
        self.g=nn.Sequential(nn.Linear(2,24), nn.Tanh(), nn.Linear(24,2))
    def forward(self,x):
        z0=self.g(torch.zeros(1,2,device=x.device)).expand(x.shape[0],-1)
        return ((self.g(x)-z0)**2).sum(1,keepdim=True)+self.eps*(x*x).sum(1,keepdim=True)

def batch(n=512, radius=1.0):
    # fixed distribution makes lambda comparisons less noisy
    return (torch.rand(n,2,device=device)*2-1)*radius

def train(lam_delta, steps=1800):
    torch.manual_seed(SEED+int(lam_delta*1000))
    p, v = Policy().to(device), Lyapunov().to(device)
    opt=torch.optim.Adam(list(p.parameters())+list(v.parameters()),lr=3e-3)
    c2=.025
    for k in range(steps):
        x=batch(512,1.0)
        u=p(x); xn=plant(x,u); vv=v(x); vn=v(xn)
        # Task: imitate a stabilizing reference while retaining a policy-learning objective.
        target=-0.9*x[:, :1]-0.25*x[:, 1:2]
        task=((u-target)**2).mean()
        pos=torch.nn.functional.softplus(.002-vv).mean()
        dec=torch.nn.functional.softplus(vn-vv+c2*(x*x).sum(1,keepdim=True)).mean()
        loss=task + 0.2*pos + lam_delta*dec
        opt.zero_grad(); loss.backward(); opt.step()
    return p,v

def evaluate(p,v,radius=1.0,n=101):
    a=torch.linspace(-radius,radius,n,device=device); X,Y=torch.meshgrid(a,a,indexing='ij')
    x=torch.stack((X.flatten(),Y.flatten()),1)
    with torch.no_grad():
        u=p(x); xn=plant(x,u); V=v(x); vn=v(xn); d=vn-V; norm=(x*x).sum(1,keepdim=True)
        margin=(d+.025*norm).flatten(); positivity=(V-.08*norm).flatten()
        cert=(margin<=0)&(positivity>=-.000001)
        # Largest centered grid radius where every point inside passes.
        rr=x.norm(dim=1)
        radii=torch.linspace(.05,radius,20,device=device)
        good=[]
        for r in radii:
            inside=rr<=r
            good.append(bool((~inside | cert).all()))
        cr=float(radii[max(i for i,z in enumerate(good) if z)]) if any(good) else 0.
        ratios=(vn/(V+1e-8)).flatten()
        certified_ratios=ratios[cert]
        return dict(task=float(((u-(-.9*x[:, :1]-.25*x[:, 1:2]))**2).mean()), violation=float(torch.relu(margin).mean()),
                    min_margin=float(margin.min()), pos_min=float(positivity.min()),
                    certified_fraction=float(cert.float().mean()), certified_radius=cr,
                    certified_ratio_max=float(certified_ratios.max()) if certified_ratios.numel() else float('nan'),
                    certified_ratio_median=float(certified_ratios.median()) if certified_ratios.numel() else float('nan'),
                    x=x.detach().cpu().numpy(), V=V.flatten().cpu().numpy(), margin=margin.detach().cpu().numpy())

def rollout(p,v,x0,steps=30):
    x=torch.tensor(x0,dtype=torch.float32,device=device).view(1,2); vals=[]
    with torch.no_grad():
        for _ in range(steps): vals.append(float(v(x).item())); x=plant(x,p(x))
    return vals

def math_checks():
    # Directly verify the advertised implication for synthetic values, including
    # the predicted boundary q=1-c2/c1 and its failure when c2>=c1.
    c1,c2=0.08,.025; q=1-c2/c1
    rng=np.random.default_rng(SEED); V0=rng.uniform(.1,10,5000)
    V=V0.copy(); ratios=[]
    for _ in range(8):
        V=(1-c2/c1)*V; ratios.append(V.copy())
    bound_errors=[float(np.max(ratios[t]-q**(t+1)*V0)) for t in range(8)]
    # A scalar parameter sweep demonstrates the exact transition at c2/c1=1.
    sweep=[]
    for c in [.01,.025,.079,.081,.10]:
        qq=1-c/c1; sweep.append({'c2':c,'predicted_q':qq,'stable_bound':bool(0<qq<1),
                                  'observed_q':qq})
    return {'c1':c1,'c2':c2,'predicted_contraction_q':q,
            'max_bound_errors':bound_errors,'contraction_bound_verified':max(map(abs,bound_errors))<1e-10,
            'boundary_sweep':sweep}

def main():
    out={'device':device,'math':math_checks(),'runs':[]}
    # lambda=0 is the standard task-only baseline; increasing lambda tests
    # predicted monotone reduction in certificate violations.
    for lam in [0., .05, .2, 1.0, 4.0]:
        p,v=train(lam); ev=evaluate(p,v)
        ev['lambda_delta']=lam
        ev['rollout_V_ratio']=rollout(p,v,[.8,.7])[-1]/max(rollout(p,v,[.8,.7])[0],1e-12)
        out['runs'].append({k:z for k,z in ev.items() if k not in ('x','V','margin')})
    # Independent radius sweep on the strongest certificate model.
    p,v=train(4.0)
    out['radius_sweep']=[{'radius':r,**{k:z for k,z in evaluate(p,v,radius=r,n=81).items() if k in ('violation','min_margin','certified_fraction','certified_radius')}} for r in [.2,.4,.6,.8,1.0]]
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()