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

SEED=362
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)
device='cuda' if torch.cuda.is_available() else 'cpu'
try:
    if device=='cuda': torch.cuda.set_device(0)
except Exception:
    device='cpu'

class MLP(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(1,24),nn.Tanh(),nn.Linear(24,24),nn.Tanh(),nn.Linear(24,1))
    def forward(self,x): return self.net(x)

def deriv(y,x,create=True):
    return torch.autograd.grad(y,x,torch.ones_like(y),create_graph=create,retain_graph=True)[0]

def make_model(): return nn.ModuleList([MLP().to(device),MLP().to(device)])

def eval_u(net,x,need2=False):
    x=x.clone().detach().requires_grad_(True); u=net(x); ux=deriv(u,x,True)
    if need2: return x,u,ux,deriv(ux,x,True)
    return x,u,ux

def setup():
    # Gauss-like uniform quadrature; exact solution x(1-x), -u''=2.
    xL=torch.linspace(0.005,0.495,48,device=device).reshape(-1,1)
    xR=torch.linspace(0.505,0.995,48,device=device).reshape(-1,1)
    xf=torch.tensor([[0.5]],device=device)
    return xL,xR,xf

def nitsche_loss(m, xL,xR,xf, gamma=20., ghost=1.):
    _,uL,dL=eval_u(m[0],xL); _,uR,dR=eval_u(m[1],xR)
    bulk=(0.5*(dL*dL)-2*uL).mean()*0.5 + (0.5*(dR*dR)-2*uR).mean()*0.5
    # exterior homogeneous Dirichlet: - n P u + gamma/(2h)u^2, h=.5
    xb=torch.tensor([[0.0],[1.0]],device=device)
    _,ul,dl=eval_u(m[0],xb[:1]); _,ur,dr=eval_u(m[1],xb[1:])
    # n=-1 left, +1 right
    ext=(dl*ul + gamma/(1.0)*ul.square()/2 - dr*ur + gamma/(1.0)*ur.square()/2).mean()
    # symmetric interface Nitsche for jump uL-uR, average flux; h=.5
    _,uiL,diL=eval_u(m[0],xf); _,uiR,diR=eval_u(m[1],xf)
    jump=uiL-uiR; avg=.5*(diL+diR)
    iface=(-avg*jump + gamma/(1.0)*jump.square()/2).mean()
    # p=1 ghost: gamma_A h^(1) [u']^2
    g=(ghost*0.5*(diL-diR).square()).mean()
    return bulk+ext+iface+g, {'bulk':bulk.item(),'boundary':(ul.square().mean()+ur.square().mean()).item(),'jump':jump.abs().item(),'djump':(diL-diR).abs().item()}

def baseline_loss(m,xL,xR,xf):
    _,uL,dL,rL=eval_u(m[0],xL,True); _,uR,dR,rR=eval_u(m[1],xR,True)
    # standard PINN strong residual and pointwise penalties, including patch interface
    loss=(rL.add(2).square().mean()+rR.add(2).square().mean())
    xb=torch.tensor([[0.0],[1.0]],device=device)
    _,ul,_=eval_u(m[0],xb[:1]); _,ur,_=eval_u(m[1],xb[1:])
    _,uiL,diL=eval_u(m[0],xf); _,uiR,diR=eval_u(m[1],xf)
    loss=loss+100*(ul.square().mean()+ur.square().mean())+100*(uiL-uiR).square().mean()+1*(diL-diR).square().mean()
    return loss, {'boundary':(ul.square().mean()+ur.square().mean()).item(),'jump':(uiL-uiR).abs().item(),'djump':(diL-diR).abs().item()}

def metrics(m):
    xx=torch.linspace(0,1,401,device=device).reshape(-1,1); vals=[]
    for a,b in [(0,.5),(.5,1)]:
        q=xx[(xx[:,0]>=a)&(xx[:,0]<=b)].reshape(-1,1); vals.append(m[0 if a==0 else 1](q))
    pred=torch.cat(vals).flatten(); q=torch.cat([xx[(xx[:,0]>=0)&(xx[:,0]<=.5)],xx[(xx[:,0]>=.5)&(xx[:,0]<=1)]])
    truth=q*(1-q); err=(pred-truth).abs().max().item()
    xL,xR,xf=setup(); _,uL,dL=eval_u(m[0],torch.tensor([[0.0]],device=device)); _,uR,dR=eval_u(m[1],torch.tensor([[1.0]],device=device)); _,a,da=eval_u(m[0],xf); _,b,db=eval_u(m[1],xf)
    return {'max_error':err,'boundary_rms':math.sqrt((uL.square()+uR.square()).mean().item()),'interface_value_jump':abs((a-b).item()),'interface_derivative_jump':abs((da-db).item())}

def run(kind,steps=900):
    torch.manual_seed(SEED+ (0 if kind=='baseline' else 1)); m=make_model(); opt=torch.optim.Adam(m.parameters(),lr=2e-3)
    xL,xR,xf=setup(); curve=[]; t=time.time()
    for it in range(steps):
        opt.zero_grad(); loss,parts=(baseline_loss(m,xL,xR,xf) if kind=='baseline' else nitsche_loss(m,xL,xR,xf)); loss.backward(); opt.step()
        if it in [0,99,299,599,899]: curve.append(float(loss.detach().cpu()))
    out=metrics(m); out.update({'loss_curve':curve,'final_objective':float(loss.detach().cpu()),'seconds':time.time()-t})
    return out

def math_check():
    # Constitutive AD identity plus a coercivity scan of the 1D symmetric
    # Nitsche quadratic form on u(x)=a*x+b.
    x=torch.tensor([[.2]],requires_grad=True)
    u=x*x/3; ux=deriv(u,x,True); F=1+ux
    W=.5*(F-1).square()
    P=torch.autograd.grad(W,F,torch.ones_like(W),retain_graph=True)[0]
    fp_err=abs((P-(F-1)).item())
    vals={}
    for gamma in [0.1,1.,2.,4.,10.,20.]:
        def q(v):
            aa,bb=v
            return .5*aa*aa-(aa*(aa+bb)+aa*bb)+gamma/2*(bb*bb+(aa+bb)**2)
        H=np.zeros((2,2)); H[0,0]=q((1,0)); H[1,1]=q((0,1))
        H[0,1]=H[1,0]=(q((1,1))-H[0,0]-H[1,1])/2
        vals[str(gamma)]=float(np.linalg.eigvalsh(H).min())
    jump=1.7-(-.4); ghost_energy=.8*.5*jump**2
    return {'F_identity_and_P_error':fp_err,
            'nitsche_min_eigenvalue_by_gamma':vals,
            'ghost_energy_for_derivative_jump':ghost_energy,
            'ghost_is_nonnegative':bool(ghost_energy>=0)}

if __name__=='__main__':
    try:
        chk=math_check(); base=run('baseline'); idea=run('idea')
    except Exception as e:
        if device=='cuda':
            device='cpu'; torch.cuda.empty_cache(); chk=math_check(); base=run('baseline'); idea=run('idea')
        else: raise
    print(json.dumps({'device':device,'math_check':chk,'baseline':base,'idea':idea},indent=2))