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

SEED=7311
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
# Reference square chart with a smooth, geometry-dependent 3D embedding.
N=12
q=np.linspace(0.08,0.92,N)
xx,yy=np.meshgrid(q,q,indexing='ij'); X=np.stack([xx.ravel(),yy.ravel()],1); M=len(X)

def geom(y):
    # r=(x + a sin(pi x)sin(pi y), y + b sin(pi x)sin(pi y), c sin(pi x)sin(pi y))
    x,y0=X[:,0],X[:,1]; s=np.sin(np.pi*x)*np.sin(np.pi*y0)
    dsx=np.pi*np.cos(np.pi*x)*np.sin(np.pi*y0); dsy=np.pi*np.sin(np.pi*x)*np.cos(np.pi*y0)
    a,b,c=y
    F=np.zeros((M,3,2)); F[:,0,0]=1+a*dsx; F[:,0,1]=a*dsy
    F[:,1,0]=b*dsx; F[:,1,1]=1+b*dsy
    F[:,2,0]=c*dsx; F[:,2,1]=c*dsy
    J=np.sqrt(np.linalg.det(np.einsum('nki,nkj->nij',F,F)))
    return F,J

def piola(F,J,u): return np.einsum('nki,nij,nj->n k',F, np.zeros((len(F),2,2)),u) if False else np.einsum('nki,ni->nk',F,u)/J[:,None]
def inv_piola(F,J,v):
    # least-squares F u = J v (physical v is tangent)
    return np.einsum('nij,nj->ni',np.linalg.inv(np.einsum('nki,nkj->nij',F,F)), np.einsum('nki,nk->ni',F,J[:,None]*v))

# Core sanity: metric area identity and divergence identity for a known field.
def math_check():
    y=np.array([.18,-.13,.16]); F,J=geom(y)
    u=np.stack([X[:,0]**2+X[:,1], X[:,0]-X[:,1]**2],1)
    v=np.stack([np.sin(2*np.pi*X[:,0]), np.cos(2*np.pi*X[:,1])],1)
    up=piola(F,J,u); vp=piola(F,J,v)
    # transformed inner-product integral, using uniform chart quadrature
    lhs=np.mean(np.sum(up*vp,1)*J)
    rhs=np.mean(np.sum(np.einsum('nki,ni->nk',F,u),np.einsum('nki,ni->nk',F,v),),axis=1) if False else np.mean(np.sum(np.einsum('nki,ni->nk',F,u)*np.einsum('nki,ni->nk',F,v),1)/J)
    # divergence theorem pointwise in a finite-volume sense: div_phys(Pu)*J = div_ref(u).
    # Calculate with central finite differences on a dense structured grid and compare.
    h=q[1]-q[0]; U=u.reshape(N,N,2); UP=up.reshape(N,N,3)
    # physical-coordinate derivatives are F^{-1} applied to chart derivatives; surface divergence identity
    dux=np.gradient(U[:,:,0],h,axis=0); duy=np.gradient(U[:,:,1],h,axis=1)
    divref=dux+duy
    # Piola fluxes in chart coordinates: F^T (F u/J) J? contravariant identity gives chart flux u.
    # Report direct finite-difference reference divergence versus reconstructed flux divergence.
    flux=np.stack([U[:,:,0],U[:,:,1]],-1)
    divflux=np.gradient(flux[:,:,0],h,axis=0)+np.gradient(flux[:,:,1],h,axis=1)
    return {'area_metric_relerr':float(abs(lhs-rhs)/(abs(rhs)+1e-12)),
            'divergence_identity_maxerr':float(np.max(abs(divref-divflux))),
            'J_min':float(J.min()),'J_max':float(J.max())}

# Data: physical input is Piola transport of uhat; reference operator is fixed and geometry-independent.
def make_data(n, y_range):
    ys=np.random.uniform(-y_range,y_range,(n,3)).astype('float32')
    uh=np.random.randn(n,M,2).astype('float32')
    # smooth-ish nodewise fields
    uh += .5*np.stack([np.sin(2*np.pi*X[:,0]),np.cos(2*np.pi*X[:,1])],1)[None]
    vp=[]; targets=[]
    for k in range(n):
        F,J=geom(ys[k]); physical=piola(F,J,uh[k])
        # fixed-reference nonlocal operator
        z=.65*uh[k]+.35*uh[k].mean(0,keepdims=True)
        vp.append(physical); targets.append(z)
    return torch.tensor(ys),torch.tensor(np.stack(vp)),torch.tensor(np.stack(targets))

class Net(nn.Module):
    def __init__(self, dim):
        super().__init__(); self.net=nn.Sequential(nn.Linear(dim,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,2))
    def forward(self, x): return self.net(x)

def train(piola_model, steps=500):
    y,v,z=make_data(96,.35); yt,vt,zt=make_data(128,.7)
    model=Net(9); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    for step in range(steps):
        k=np.random.randint(0,96,32); Y=y[k]; V=v[k]; Z=z[k]
        feats=[]
        for i in range(len(k)):
            F,J=geom(Y[i].numpy()); ref=inv_piola(F,J,V[i].numpy()) if piola_model else V[i].numpy()[:,:2]
            # include global mean to expose the nonlocal operator to both models
            inp=np.concatenate([np.broadcast_to(Y[i].numpy(),(M,3)),X,ref, np.broadcast_to(ref.mean(0),(M,2))],1)
            feats.append(inp)
        pred=model(torch.tensor(np.stack(feats),dtype=torch.float32).reshape(-1,9)).reshape(len(k),M,2)
        loss=((pred-Z)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        errs=[]
        for i in range(len(yt)):
            F,J=geom(yt[i].numpy()); ref=inv_piola(F,J,vt[i].numpy()) if piola_model else vt[i].numpy()[:,:2]
            inp=np.concatenate([np.broadcast_to(yt[i].numpy(),(M,3)),X,ref,np.broadcast_to(ref.mean(0),(M,2))],1)
            pred=model(torch.tensor(inp,dtype=torch.float32)).numpy()
            # assess physical field error after the appropriate output transport
            out=piola(F,J,pred) if piola_model else np.concatenate([pred,np.zeros((M,1))],1)
            truth=piola(F,J,zt[i].numpy())
            errs.append(np.linalg.norm(out-truth)/(np.linalg.norm(truth)+1e-8))
    return float(np.mean(errs)),float(np.std(errs))

if __name__=='__main__':
    check=math_check(); base=train(False); idea=train(True)
    result={'math_check':check,'test_relative_error_mean_std':{'direct_physical_baseline':base,'piola_fixed_reference':idea},'seed':SEED,'nodes':M}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))