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

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

torch.set_num_threads(8)

# Smooth random Fourier fields and a known operator: output is a shifted,
# mildly nonlinear transform of the input field.
def fields(n, m, phase=0):
    x=np.linspace(0,1,m,endpoint=False)[None,:]
    rng=np.random.RandomState(SEED+phase)
    a=rng.randn(n,5)*np.array([1,.6,.4,.3,.2])[None,:]
    b=rng.randn(n,5)*np.array([1,.6,.4,.3,.2])[None,:]
    u=np.zeros((n,m));
    for k in range(1,6): u += a[:,k-1,None]*np.sin(2*np.pi*k*x)+b[:,k-1,None]*np.cos(2*np.pi*k*x)
    u=u/(np.std(u,axis=1,keepdims=True)+1e-6)*.7
    # operator is pointwise nonlinear plus shifted low-frequency component
    y=np.tanh(u)+.25*np.roll(u,2,axis=1)
    return x.astype(np.float32),u.astype(np.float32),y.astype(np.float32)

def interp_periodic(x0, v, q):
    xx=np.mod(x0,1); order=np.argsort(xx); xx=xx[order]; vv=v[...,order]
    xx2=np.r_[xx,xx[0]+1]; vv2=np.concatenate([vv,vv[...,0:1]],axis=-1)
    qq=np.mod(q,1); j=np.searchsorted(xx2,qq,side='right')-1; j=np.clip(j,0,len(xx2)-2)
    t=(qq-xx2[j])/(xx2[j+1]-xx2[j]+1e-8)
    return vv2[...,j]*(1-t)+vv2[...,j+1]*t

class MeasurementNet(nn.Module):
    def __init__(self,z=32):
        super().__init__(); self.enc=nn.Sequential(nn.Linear(2,48),nn.Tanh(),nn.Linear(48,z),nn.Tanh())
        self.g=nn.Sequential(nn.Linear(z,64),nn.Tanh(),nn.Linear(64,z),nn.Tanh())
        self.dec=nn.Sequential(nn.Linear(z+1,64),nn.Tanh(),nn.Linear(64,1))
    def encode(self, coords, vals):
        h=self.enc(torch.stack((coords,vals),-1)); return h.mean(1)
    def forward(self, coords, vals, q):
        z=self.g(self.encode(coords,vals)); zz=z[:,None,:].expand(-1,q.shape[1],-1)
        return self.dec(torch.cat((zz,q[:,:,None]),-1)).squeeze(-1)

class FixedNet(nn.Module):
    def __init__(self,m=32):
        super().__init__(); self.net=nn.Sequential(nn.Linear(m,96),nn.Tanh(),nn.Linear(96,128),nn.Tanh(),nn.Linear(128,64))
    def forward(self,v): return self.net(v)

def train_models():
    x,u,y=fields(900,64,1); xt,ut,yt=fields(180,64,2)
    # canonical measurements are first 32 regular points, output 64 grid
    ci=np.linspace(0,63,32).round().astype(int); qi=np.arange(64)
    mc=MeasurementNet().to(device); base=FixedNet().to(device)
    opt=torch.optim.Adam(mc.parameters(),lr=3e-3); opb=torch.optim.Adam(base.parameters(),lr=3e-3)
    X=torch.tensor(x[:,ci],device=device); U=torch.tensor(u[:,ci],device=device); Q=torch.tensor(x[:,qi],device=device)
    Y=torch.tensor(y,device=device); UB=torch.tensor(u[:,ci],device=device)
    for step in range(300):
        ix=torch.randint(0,len(x),(96,),device=device)
        pred=mc(X[ix],U[ix],Q[ix]); loss=((pred-Y[ix])**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        pb=base(UB[ix]); lb=((pb-Y[ix])**2).mean(); opb.zero_grad(); lb.backward(); opb.step()
    return mc.cpu().eval(),base.cpu().eval(),(x,u,y,xt,ut,yt,ci)

def predict_mesh(model, coords, vals, q):
    with torch.no_grad():
        return model(torch.tensor(coords[None],dtype=torch.float32),torch.tensor(vals[None],dtype=torch.float32),torch.tensor(q[None],dtype=torch.float32)).numpy()[0]

def main():
    t=time.time(); mc,base,data=train_models(); x,u,y,xt,ut,yt,ci=data
    # Mesh transfer: input sensors are alternative layouts; output is arbitrary query mesh.
    rng=np.random.RandomState(77); rows=[]
    for n in [8,16,32,64]:
        errs=[]
        ss=np.sort(rng.rand(len(xt),n),axis=1); vals=np.stack([interp_periodic(x[0],ut[i],ss[i]) for i in range(len(xt))])
        q=np.sort(rng.rand(len(xt),48),axis=1); truth=np.stack([interp_periodic(x[0],yt[i],q[i]) for i in range(len(xt))])
        with torch.no_grad():
            pred=mc(torch.tensor(ss,dtype=torch.float32),torch.tensor(vals,dtype=torch.float32),torch.tensor(q,dtype=torch.float32)).numpy()
        rows.append((n,float(np.sqrt(np.mean((pred-truth)**2)))))
    # baseline canonical error on its supported fixed grid, and baseline cannot accept new mesh.
    with torch.no_grad():
        bp=base(torch.tensor(ut[:,ci],dtype=torch.float32)).numpy()
    base_rmse=float(np.sqrt(np.mean((bp-yt)**2)))
    # Core prediction 1: permutation invariance of E_X.
    i=0; s=np.sort(rng.rand(20)); vv=interp_periodic(x[0],u[i],s)
    p1=predict_mesh(mc,s,vv,np.linspace(0,1,20,endpoint=False)); perm=rng.permutation(20)
    p1p=predict_mesh(mc,s[perm],vv[perm],np.linspace(0,1,20,endpoint=False))
    perm_diff=float(np.max(np.abs(p1-p1p)))
    # Core prediction 2: Lipschitz composition discrepancy scales linearly with reconstruction perturbation.
    # Hold a measurement set fixed and perturb reconstructed function by amplitude delta.
    s=np.linspace(0,1,32,endpoint=False); vv=interp_periodic(x[0],u[0],s); q=np.linspace(0,1,40,endpoint=False)
    direction=np.sin(2*np.pi*3*s); deltas=np.array([0,.01,.02,.04,.08,.16]); diffs=[]
    for d in deltas:
        a=predict_mesh(mc,s,vv,q); b=predict_mesh(mc,s,vv+d*direction,q); diffs.append(float(np.sqrt(np.mean((a-b)**2))))
    slope=float(np.polyfit(deltas[1:],diffs[1:],1)[0]); r2=float(np.corrcoef(deltas,diffs)[0,1])
    # Core prediction 3: more measurements improve reconstruction/transfer, with diminishing returns.
    # Compare fixed Fourier interpolation reconstruction error against true input field.
    recon=[]
    for n in [4,8,16,32,64]:
        ee=[]
        for i in range(50):
            s=np.sort(rng.rand(n)); vv=interp_periodic(x[0],u[i],s); q=x[0]
            ee.append(np.mean((interp_periodic(s,vv,q)-u[i])**2))
        recon.append((n,float(np.sqrt(np.mean(ee)))))
    out={'device':device,'seconds':round(time.time()-t,2),'baseline_fixed32_rmse':base_rmse,
         'idea_mesh_transfer_rmse_by_input_sensors':rows,'prediction_permutation_max_difference':perm_diff,
         'prediction_lipschitz_sweep':{'deltas':deltas.tolist(),'output_discrepancies':diffs,'linear_slope':slope,'correlation':r2},
         'prediction_sensor_density_reconstruction_rmse':recon,
         'note':'Measurement model uses coordinate/value mean pooling and coordinate-conditioned output decoder; baseline is fixed-grid MLP.'}
    print(json.dumps(out,indent=2))
if __name__=='__main__':
    try: main()
    except Exception as e:
        if device=='cuda':
            print('CUDA failed, rerun CPU:',repr(e)); device='cpu'; torch.cuda.empty_cache(); main()
        else: raise