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

SEED = 2746
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)

def get_device():
    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def boxes_from_tensors(cx, cy, yaw, length, width):
    c = torch.stack((cx, cy), -1)
    u = torch.stack((torch.cos(yaw), torch.sin(yaw)), -1)
    v = torch.stack((-torch.sin(yaw), torch.cos(yaw)), -1)
    return c, u, v, length / 2, width / 2

def sat_barrier(ca, ya, la, wa, cb, yb, lb, wb, tau=0.1):
    ca, ua, va, aa, ba = boxes_from_tensors(ca[...,0], ca[...,1], ya, la, wa)
    cb, ub, vb, ab, bb = boxes_from_tensors(cb[...,0], cb[...,1], yb, lb, wb)
    axes = torch.stack((ua, va, ub, vb), -2)
    d = cb - ca
    proj = (axes * d.unsqueeze(-2)).sum(-1).abs()
    def radius(u, v, a, b):
        pu = (u.unsqueeze(-2) * axes).sum(-1).abs()
        pv = (v.unsqueeze(-2) * axes).sum(-1).abs()
        return a.unsqueeze(-1)*pu + b.unsqueeze(-1)*pv
    gaps = proj - radius(ua, va, aa, ba) - radius(ub, vb, ab, bb)
    margin = tau * torch.logsumexp(gaps / tau, dim=-1)
    return margin, gaps

def hard_margin(ca, ya, la, wa, cb, yb, lb, wb):
    _, gaps = sat_barrier(ca, ya, la, wa, cb, yb, lb, wb, tau=1e-4)
    return gaps.max(-1).values

def soft_barrier(margin, m0=0.05, beta=0.05):
    return F.softplus((m0-margin)/beta)

def exact_signed_clearance(ca, ya, la, wa, cb, yb, lb, wb):
    # Polygon signed clearance for separated rectangles, with negative overlap
    # represented by the SAT maximum margin (the exact positive clearance).
    return hard_margin(ca, ya, la, wa, cb, yb, lb, wb)

def math_sweeps():
    out = {}
    dev = torch.device('cpu')
    z = torch.tensor([[0., 0.]], device=dev)
    one = torch.tensor([0.], device=dev)
    dims = [torch.tensor([2.]), torch.tensor([1.])]
    ds = np.linspace(0.5, 5.0, 91)
    exact=[]; smooth=[]
    for d in ds:
        cb=torch.tensor([[float(d),0.]])
        ex=hard_margin(z,one,dims[0],dims[1],cb,one,dims[0],dims[1]).item()
        sm=sat_barrier(z,one,dims[0],dims[1],cb,one,dims[0],dims[1],tau=.1)[0].item()
        exact.append(ex); smooth.append(sm)
    # Prediction 1: axis-aligned SAT crossing is d=length sum / 2 = 2m.
    i=np.argmin(np.abs(np.asarray(exact)))
    j=np.argmin(np.abs(np.asarray(smooth)))
    out['zero_crossing']={'predicted_exact_m':2.0,'observed_exact_m':float(ds[i]),'observed_smooth_zero_m':float(ds[j]),'smooth_bias_m':float(smooth[j]-exact[j])}
    # Prediction 2: logsumexp smoothing is bounded above by tau log(4), and approaches max as tau shrinks.
    gaps=torch.tensor([[1.0, .2, -.3, -.8]])
    maxg=gaps.max().item(); bounds=[]
    for tau in [.2,.1,.05,.01]:
        val=(tau*torch.logsumexp(gaps/tau,-1)).item()
        bounds.append({'tau':tau,'value':val,'error':val-maxg,'bound':tau*math.log(4)})
    out['smoothing_bound']={'predicted_error_leq_tau_log4':bounds}
    # Prediction 3: at a separated axis-aligned configuration the active-axis gradient is ~1.
    grad=[]
    for tau in [.5,.2,.1,.05,.01]:
        x=torch.tensor([[3.0,0.]],requires_grad=True)
        m,_=sat_barrier(z,one,dims[0],dims[1],x,one,dims[0],dims[1],tau=tau)
        m.backward(); grad.append({'tau':tau,'dx':float(x.grad[0,0])})
    out['gradient']={'predicted_dx_near_1':grad}
    return out

class Planner(nn.Module):
    def __init__(self, horizon=8):
        super().__init__(); self.net=nn.Sequential(nn.Linear(3,32),nn.Tanh(),nn.Linear(32, horizon*2)); self.h=horizon
    def forward(self,x): return self.net(x).view(-1,self.h,2)

def train_compare():
    torch.manual_seed(SEED)
    dev=get_device(); n=96; h=8
    x=torch.linspace(-1,1,n,device=dev).unsqueeze(1)
    inp=torch.cat((x, torch.sin(2*x), torch.cos(2*x)),1)
    target=torch.stack((torch.linspace(0,4,h,device=dev).repeat(n,1), .7*x.repeat(1,h)), -1)
    # Obstacle box centered at x=2, y=0; desired paths cross it for x near zero.
    obs_c=torch.tensor([2.,0.],device=dev); obs_y=torch.tensor(0.,device=dev)
    results={}
    for use_barrier in [False,True]:
        torch.manual_seed(SEED); model=Planner(h).to(dev); opt=torch.optim.Adam(model.parameters(),lr=.02)
        for _ in range(350):
            pred=model(inp); imitation=((pred-target)**2).mean()
            if use_barrier:
                ca=torch.zeros((n,h,2),device=dev); ca[...,0]=pred[...,0]; ca[...,1]=pred[...,1]
                cb=obs_c.view(1,1,2).expand(n,h,2)
                m,_=sat_barrier(ca, torch.zeros((n,h),device=dev), torch.full((n,h),2.,device=dev), torch.full((n,h),1.,device=dev), cb, obs_y.expand(n,h), torch.full((n,h),2.,device=dev), torch.full((n,h),1.,device=dev),tau=.1)
                loss=imitation + .25*soft_barrier(m,m0=.05,beta=.05).mean()
            else: loss=imitation
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=model(inp); ca=pred; cb=obs_c.view(1,1,2).expand(n,h,2)
            margins=hard_margin(ca,torch.zeros((n,h),device=dev),torch.full((n,h),2.,device=dev),torch.full((n,h),1.,device=dev),cb,obs_y.expand(n,h),torch.full((n,h),2.,device=dev),torch.full((n,h),1.,device=dev))
            results['barrier' if use_barrier else 'baseline']={'mse':float(((pred-target)**2).mean()),'min_margin':float(margins.min()),'collision_fraction':float((margins<=0).float().mean())}
    return results

def main():
    report={'seed':SEED,'device':str(get_device()),'math':math_sweeps(),'mini_experiment':train_compare()}
    Path('results.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()