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

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

def kernel_np(x, y, d, s, sigma, c=None, eps=1e-8):
    if c is None: c = np.zeros(x.shape[-1])
    r = np.linalg.norm(x-y, axis=-1) + eps
    ux = np.linalg.norm(x-c, axis=-1) + eps
    uy = np.linalg.norm(y-c, axis=-1) + eps
    m = np.minimum(np.minimum(ux/r, uy/r), 1.0)
    return r**(s-d) * m**(-sigma)

def math_checks():
    d, s, sigma = 3, 0.8, 0.35
    x = np.array([[0.31, 0.22, 0.17], [0.47, -0.19, 0.28], [0.18, 0.41, -0.23]])
    y = np.array([[-0.24, 0.16, 0.29], [-0.13, 0.39, -0.31], [-0.36, -0.11, 0.21]])
    k1 = kernel_np(x, y, d, s, sigma)
    scale = 3.0
    k2 = kernel_np(scale*x, scale*y, d, s, sigma)
    hom_err = float(np.max(np.abs(k2/(k1*scale**(s-d))-1)))
    xa, ya = np.array([[1e-5, 0.0]]), np.array([[0.8, 0.0]])
    plain = (np.linalg.norm(xa-ya, axis=1)+1e-8)**(s-d)
    full = kernel_np(xa, ya, d, s, sigma)
    amplification = float(full[0]/plain[0])
    predicted = float(((np.linalg.norm(xa[0])+1e-8)/(np.linalg.norm(xa[0]-ya[0])+1e-8))**(-sigma))
    radii = np.logspace(-7, -2, 3000)
    integ = np.trapz(radii**(d-1-sigma), radii)
    expected = (1e-2**(d-sigma)-1e-7**(d-sigma))/(d-sigma)
    return {'homogeneity_max_relative_error': hom_err,
            'near_origin_amplification': amplification,
            'predicted_amplification': predicted,
            'origin_integral_relative_error': float(abs(integ/expected-1)),
            'admissible_d3_demo': bool(0 < s < d-2*sigma)}

class AttnRegressor(nn.Module):
    def __init__(self, biased, n=64, dim=32, s=.8, sigma=.35):
        super().__init__(); self.biased=biased; self.n=n; self.dim=dim
        self.inp=nn.Linear(2, dim); self.q=nn.Linear(dim, dim); self.k=nn.Linear(dim, dim)
        self.v=nn.Linear(dim, dim); self.out=nn.Sequential(nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, 1))
        self.s=s; self.sigma=sigma
    def forward(self, xy):
        z=self.inp(xy); q=self.q(z); k=self.k(z); v=self.v(z)
        logits=q@k.transpose(-1,-2)/math.sqrt(self.dim)
        if self.biased:
            diff=xy[:,:,None,:]-xy[:,None,:,:]
            r=torch.sqrt((diff*diff).sum(-1)+1e-8)+1e-5
            u=torch.sqrt((xy*xy).sum(-1)+1e-8)+1e-5
            m=torch.minimum(torch.minimum(u[:,:,None]/r,u[:,None,:]/r),torch.ones_like(r))
            K=r**(self.s-2)*m**(-self.sigma)
            logits=logits+0.7*torch.log(K+1e-8)
        a=torch.softmax(logits, -1)
        return self.out(a@v)

def make_data(num, n=64):
    xy=np.random.uniform(-1,1,(num,n,2)).astype('float32')
    r=np.sqrt((xy*xy).sum(-1)+1e-5)
    # clipped cusp with a smooth multiscale component
    target=np.minimum(1.0, r**(-0.65))/2 + .15*np.sin(8*xy[:,:,0])*np.cos(6*xy[:,:,1])
    return torch.tensor(xy), torch.tensor(target[:,:,None], dtype=torch.float32)

def run():
    checks=math_checks(); device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        train_x, train_y=make_data(256); test_x, test_y=make_data(64)
        train_x,train_y,test_x,test_y=[z.to(device) for z in (train_x,train_y,test_x,test_y)]
        results={}
        for biased in (False, True):
            torch.manual_seed(SEED); model=AttnRegressor(biased).to(device)
            opt=torch.optim.Adam(model.parameters(),lr=3e-3); loss_fn=nn.MSELoss()
            curve=[]
            for step in range(180):
                idx=torch.randint(0,train_x.shape[0],(16,),device=device)
                loss=loss_fn(model(train_x[idx]),train_y[idx])
                opt.zero_grad(); loss.backward(); opt.step()
                if step in (0,29,89,179):
                    with torch.no_grad(): curve.append(float(loss_fn(model(test_x),test_y).cpu()))
            with torch.no_grad(): final=float(loss_fn(model(test_x),test_y).cpu())
            results['idea' if biased else 'baseline']={'test_mse_curve':curve,'final_test_mse':final}
    except Exception as e:
        device='cpu'; results={'error':repr(e)}
    out={'seed':SEED,'device':device,'checks':checks,'results':results}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': run()