import json, time
import numpy as np
import torch

SEED = 2025
torch.manual_seed(SEED); np.random.seed(SEED)
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if device == "cuda": torch.cuda.set_device(0); torch.cuda.empty_cache()
except Exception:
    device = "cpu"


def _work_dtype(x):
    return torch.float64 if x.dtype == torch.float64 else torch.float32

def truncated(x, c, xmin=-2.0, h=0.25):
    """Vectorized fixed truncated-power cubic cardinal spline contraction."""
    n = c.shape[-1]; dt = _work_dtype(x)
    xx, cc = x.to(dt), c.to(dt)
    u = ((xx - xmin) / h).clamp(0., float(n))
    i = torch.arange(n, device=x.device, dtype=dt)
    r = torch.arange(5, device=x.device, dtype=dt)
    a = torch.tensor([1., -4., 6., -4., 1.], device=x.device, dtype=dt)
    z = u[..., None, None] - i[None, :, None] - r[None, None, :]
    basis = (torch.relu(z).pow(3) * a).sum(-1) / 6.
    return (basis * cc).sum(-1).to(x.dtype)

def cox_de_boor(x, c, xmin=-2.0, h=0.25):
    """Reference Cox-de Boor evaluation of B_3(u-i), i=0,...,N-1."""
    n = c.shape[-1]; dt = _work_dtype(x)
    xx, cc = x.to(dt), c.to(dt)
    u = ((xx - xmin) / h).clamp(0., float(n))
    m = n + 3
    knots = torch.arange(m, device=x.device, dtype=dt)
    b = ((u[..., None] >= knots) & (u[..., None] < knots + 1.)).to(dt)
    for d in range(1, 4):
        ii = torch.arange(m-d, device=x.device, dtype=dt)
        b = ((u[..., None]-ii)/d)*b[..., :m-d] + ((ii+d+1.-u[..., None])/d)*b[..., 1:m-d+1]
    return (b[..., :n] * cc).sum(-1).to(x.dtype)

def sync():
    if device == "cuda": torch.cuda.synchronize()

def timed(fn, x, c, warm=30, reps=150):
    for _ in range(warm): fn(x,c)
    sync(); t0=time.perf_counter()
    for _ in range(reps): fn(x,c)
    sync(); return (time.perf_counter()-t0)*1000/reps

def main():
    n=16
    x=torch.randn(4096,device=device)*1.8-.3; c=torch.randn(4096,n,device=device)
    x1=x.clone().requires_grad_(True); x2=x.clone().requires_grad_(True)
    y1,y2=truncated(x1,c),cox_de_boor(x2,c)
    g1=torch.autograd.grad(y1.sum(),x1)[0]; g2=torch.autograd.grad(y2.sum(),x2)[0]
    xo=torch.tensor([-100.,-10.,0.,10.,100.],device=device); yo=truncated(xo,torch.ones(5,n,device=device))
    xb=torch.randn(128,32,32,device=device)*1.5-.2; cb=torch.randn(128,32,32,n,device=device)
    ft=timed(truncated,xb,cb); fc=timed(cox_de_boor,xb,cb)
    def step(fn):
        xx=xb.detach().requires_grad_(True); cc=cb.detach().requires_grad_(True)
        (fn(xx,cc).square().mean()).backward()
    st=timed(lambda a,b:step(truncated),xb,cb,warm=8,reps=30)
    sc=timed(lambda a,b:step(cox_de_boor),xb,cb,warm=8,reps=30)
    # Valid precision comparison; unlike the original version, float64 is preserved.
    xd=(torch.rand(20000,device=device,dtype=torch.float64)*16-0)*.25-2
    cd=torch.randn(20000,n,device=device,dtype=torch.float64)
    ed=(truncated(xd,cd)-cox_de_boor(xd,cd)).abs()
    result={"device":device,"n":n,"forward_max_abs_error":(y1-y2).abs().max().item(),
      "forward_relative_mean_error":(y1-y2).abs().mean().item()/(y2.abs().mean().item()+1e-12),
      "gradient_max_abs_error":(g1-g2).abs().max().item(),"out_of_range_finite":bool(torch.isfinite(yo).all()),
      "out_of_range_outputs":yo.detach().cpu().tolist(),"forward_ms_truncated":ft,"forward_ms_cox_de_boor":fc,
      "forward_speedup":fc/ft,"step_ms_truncated":st,"step_ms_cox_de_boor":sc,"step_speedup":sc/st,
      "float64_max_abs_error":ed.max().item(),"float64_mean_abs_error":ed.mean().item()}
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()