import json, math, random, time
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)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda':
        torch.cuda.manual_seed_all(SEED)
        # force a tiny allocation to detect unusable CUDA contexts
        torch.zeros(1, device=device)
except Exception:
    device=torch.device('cpu')

a=0.75
s=0.10
lam_g=0.01
EPS=1e-6

class MLP(nn.Module):
    def __init__(self, width=48, depth=3):
        super().__init__()
        layers=[nn.Linear(1,width), nn.Tanh()]
        for _ in range(depth-1): layers += [nn.Linear(width,width), nn.Tanh()]
        layers += [nn.Linear(width,1)]
        self.net=nn.Sequential(*layers)
        for m in self.modules():
            if isinstance(m,nn.Linear):
                nn.init.xavier_uniform_(m.weight); nn.init.zeros_(m.bias)
    def forward(self,x): return self.net(x)

def d_exact(x): return torch.minimum(x,1-x)
def target(x):
    # Smooth quotient times the fractional boundary factor.
    d=d_exact(x)
    return d.pow(a)*(1.0+0.35*torch.sin(2*math.pi*x)+0.15*x)
def weighted_grad(v,x):
    grad=torch.autograd.grad(v.sum(),x,create_graph=True)[0]
    d=d_exact(x).clamp_min(EPS)
    return d.pow(1-a+s)*grad

def train(kind, steps=1800, n=96):
    torch.manual_seed(SEED)
    model=MLP().to(device)
    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    hist=[]; t0=time.time()
    for it in range(steps):
        x=torch.rand(n,1,device=device)
        # avoid the cusp of min(x,1-x) at the midpoint for this gradient regularizer
        x.requires_grad_(kind=='idea')
        y=target(x)
        raw=model(x)
        if kind=='idea':
            u=d_exact(x).clamp_min(EPS).pow(a)*raw
            wg=weighted_grad(raw,x)
            loss=((u-y)**2).mean()+lam_g*(wg**2).mean()
        else:
            u=raw
            # Standard free-output PINN-style boundary penalty.
            xb=torch.tensor([[0.0],[1.0]],device=device)
            loss=((u-y)**2).mean()+10.0*(model(xb)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if it in (0,99,499,999,1799): hist.append(float(loss.detach().cpu()))
    with torch.no_grad():
        xx=torch.linspace(1e-5,1-1e-5,4000,device=device).view(-1,1)
        pred=(d_exact(xx).clamp_min(EPS).pow(a)*model(xx) if kind=='idea' else model(xx))
        yy=target(xx)
        rel=float(torch.sqrt(((pred-yy)**2).mean())/torch.sqrt((yy**2).mean()))
        strip=((xx<0.05)|(xx>0.95))
        striperr=float(torch.sqrt(((pred[strip]-yy[strip])**2).mean())/torch.sqrt((yy[strip]**2).mean()))
        maxb=float(torch.max(torch.abs(pred[strip]-yy[strip])))
    return {'loss_checkpoints':hist,'relative_L2':rel,'boundary_strip_relative_L2':striperr,'boundary_strip_max_abs':maxb,'seconds':time.time()-t0}

def math_check():
    # For smooth v, the claimed weighted derivative should vanish at boundary,
    # whereas u'=a*d^(a-1)*v has the fractional singular scaling.
    x=torch.logspace(-6,-2,80,dtype=torch.float64)
    v=1+0.4*x
    vp=torch.full_like(x,0.4)
    up=a*x**(a-1)*v+x**a*vp
    w=x**(1-a+s)*vp
    # Fit log-log slopes; expected u' slope a-1 and weighted smooth term slope 1-a+s.
    su=float(np.polyfit(np.log(x.numpy()),np.log(torch.abs(up).numpy()),1)[0])
    sw=float(np.polyfit(np.log(x.numpy()),np.log(torch.abs(w).numpy()),1)[0])
    ratio=float((w[-1]/w[0]).item())
    return {'unweighted_uprime_log_slope':su,'expected_unweighted_slope':a-1,
            'weighted_vprime_log_slope':sw,'expected_weighted_slope':1-a+s,
            'weighted_growth_across_range':ratio,'weighted_decreases_to_boundary':bool(w[0]<w[-1])}

def main():
    out={'device':str(device),'seed':SEED,'a':a,'s':s,'lambda_g':lam_g,
         'math_check':math_check()}
    out['baseline']=train('baseline'); out['idea']=train('idea')
    out['idea_minus_baseline_relative_L2']=out['idea']['relative_L2']-out['baseline']['relative_L2']
    out['idea_minus_baseline_boundary_strip']=out['idea']['boundary_strip_relative_L2']-out['baseline']['boundary_strip_relative_L2']
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()