import json, math, random
import numpy as np
import torch
from torch import nn

SEED=1234
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # Probe CUDA so an unavailable/broken shared device falls back cleanly.
    if device.type=='cuda':
        torch.zeros(1,device=device).sum().item()
except Exception:
    device=torch.device('cpu')

lam=0.5

def field_np(x):
    r=np.sqrt(x[:,0]**2+x[:,1]**2)
    th=np.arctan2(x[:,1],x[:,0])
    sing=r**lam*np.sin(th/2)
    smooth=0.2*np.sin(np.pi*x[:,0])*np.sin(np.pi*x[:,1])+0.1*x[:,0]*x[:,1]
    return sing+smooth

def sample_uniform(n): return np.random.rand(n,2)
def sample_graded(n):
    # Concentrate radial mass near the junction while retaining the full square.
    # Independent Beta(.55,1) coordinates gives a near-corner radial grading.
    return np.random.beta(.55,1.0,size=(n,2))

class MLP(nn.Module):
    def __init__(self, enriched=False):
        super().__init__(); self.enriched=enriched
        self.net=nn.Sequential(nn.Linear(2,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh(),nn.Linear(64,1))
        if enriched: self.a=nn.Parameter(torch.tensor([0.8]))
    def forward(self,x):
        z=self.net(x)
        if self.enriched:
            r=torch.sqrt((x*x).sum(1)+1e-12); th=torch.atan2(x[:,1],x[:,0])
            phi=torch.sin(th/2).unsqueeze(1)
            r=r.unsqueeze(1)
            # cutoff is one in the relevant near-corner region and smoothly vanishes
            r0=.7; q=torch.clamp(r/r0,0,1)
            chi=1-3*q*q+2*q*q*q
            sing=chi*r.pow(lam)*phi*self.a
            return sing+(1-chi)*z+chi*z*r.pow(lam)
        return z

def train(enriched, sampler, n=1600, steps=1400):
    x=torch.tensor(sampler(n),dtype=torch.float32,device=device)
    y=torch.tensor(field_np(x.detach().cpu().numpy()),dtype=torch.float32,device=device).view(-1,1)
    model=MLP(enriched).to(device)
    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    for i in range(steps):
        # fixed points makes the equal-evaluation comparison deterministic
        pred=model(x); loss=((pred-y)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return model, float(loss.detach().cpu())

def evaluate(model):
    # fixed dense test grid, plus a corner-only region
    g=np.linspace(.001,1,121); xx,yy=np.meshgrid(g,g); X=np.c_[xx.ravel(),yy.ravel()]
    with torch.no_grad(): p=model(torch.tensor(X,dtype=torch.float32,device=device)).cpu().numpy().ravel()
    y=field_np(X)
    rel=np.linalg.norm(p-y)/np.linalg.norm(y)
    mask=np.sqrt((X*X).sum(1))<.12
    corner=np.linalg.norm(p[mask]-y[mask])/np.linalg.norm(y[mask])
    return float(rel),float(corner),int(mask.sum())

def math_check():
    # Along a fixed ray the singular part must have log-log slope lambda.
    r=np.logspace(-6,-1,80); X=np.c_[r,r]
    s=np.sqrt(2*r*r)**lam*np.sin(np.pi/8)
    slope=np.polyfit(np.log(r),np.log(np.abs(s)),1)[0]
    # Smooth remainder has bounded first radial derivative numerically at the corner.
    # Compare its radial increments at two scales (increments should be O(r)).
    def rem(t): return .2*np.sin(np.pi*t/np.sqrt(2))**2+.1*(t/np.sqrt(2))**2
    inc1=abs(rem(2e-3)-rem(1e-3)); inc2=abs(rem(4e-3)-rem(2e-3))
    return {'loglog_slope':float(slope),'expected_lambda':lam,'slope_abs_error':float(abs(slope-lam)), 'smooth_increment_ratio':float(inc2/inc1)}

if __name__=='__main__':
    check=math_check()
    results={}
    configs=[('baseline_uniform',False,sample_uniform),('baseline_graded',False,sample_graded),('enriched_uniform',True,sample_uniform),('enriched_graded',True,sample_graded)]
    for name,en,samp in configs:
        m,trainloss=train(en,samp)
        rel,corner,count=evaluate(m)
        results[name]={'train_mse':trainloss,'relative_l2':rel,'corner_relative_l2':corner,'corner_test_points':count}
    # Verify grading actually changes allocation without changing point count.
    u=sample_uniform(100000); g=sample_graded(100000)
    grading={'n_uniform':len(u),'n_graded':len(g),'fraction_r_lt_.12_uniform':float((np.linalg.norm(u,axis=1)<.12).mean()),'fraction_r_lt_.12_graded':float((np.linalg.norm(g,axis=1)<.12).mean())}
    out={'device':str(device),'math_check':check,'grading_check':grading,'results':results,'seed':SEED,'train_points':1600,'steps':1400}
    print(json.dumps(out,indent=2))