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

SEED=19
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.zeros(1,device=device)
except Exception:
    device=torch.device('cpu')

N=2048
rng=np.random.default_rng(SEED)
X=torch.tensor(rng.uniform(-1,1,(N,2)),dtype=torch.float32)
Y=(X[:,0:1]*X[:,1:2]+.25*torch.sin(2*X[:,0:1]-X[:,1:2]))
train_slice=slice(0,1536); test_slice=slice(1536,None)

class Field(nn.Module):
    def __init__(self,d,h=32):
        super().__init__(); self.net=nn.Sequential(nn.LayerNorm(d),nn.Linear(d,h),nn.GELU(),nn.Linear(h,d))
    def forward(self,x): return self.net(x)

class Baseline(nn.Module):
    def __init__(self,d=16):
        super().__init__(); self.inp=nn.Linear(2,d); self.f=Field(d); self.out=nn.Linear(d,1); self.gate=nn.Parameter(torch.tensor(.1))
    def forward(self,x):
        z=torch.tanh(self.inp(x)); return self.out(z+self.gate*self.f(z))

class Cyclic(nn.Module):
    def __init__(self,d=16, random_order=True):
        super().__init__(); self.inp=nn.Linear(2,d); self.f1=Field(d); self.f2=Field(d); self.out=nn.Linear(d,1)
        self.gate=nn.Parameter(torch.tensor(.1)); self.random_order=random_order
    def forward(self,x):
        z=torch.tanh(self.inp(x)); base=z
        # Center fields pointwise, preserving noncommutativity while canceling
        # their first-order average. Random centered run times have variance 1.
        u=self.f1(z); v=self.f2(z); mean=(u+v)/2
        fields=(u-mean,v-mean)
        a=torch.randn((),device=z.device); b=torch.randn((),device=z.device)
        if self.training and self.random_order and torch.rand((),device=z.device)<.5:
            z=z+self.gate*.35*a*fields[1]; z=z+self.gate*.35*b*(self.f1(z)-self.f2(z))/2
        else:
            z=z+self.gate*.35*a*fields[0]; z=z+self.gate*.35*b*(self.f2(z)-self.f1(z))/2
        return self.out(z)

def train(kind):
    torch.manual_seed(SEED)
    model=Baseline() if kind=='baseline' else Cyclic()
    model.to(device); xx=X.to(device); yy=Y.to(device)
    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    losses=[]; t=time.time()
    for step in range(350):
        model.train(); opt.zero_grad(); pred=model(xx[train_slice]); loss=((pred-yy[train_slice])**2).mean(); loss.backward(); opt.step()
        losses.append(float(loss.detach().cpu()))
    model.eval()
    with torch.no_grad():
        tr=float(((model(xx[train_slice])-yy[train_slice])**2).mean().cpu())
        te=float(((model(xx[test_slice])-yy[test_slice])**2).mean().cpu())
    return {'train_mse':tr,'test_mse':te,'loss_step_50':losses[49],'loss_step_350':losses[-1], 'seconds':time.time()-t, 'parameters':sum(p.numel() for p in model.parameters())}

# Explicit order signal using the learned fields and shared input, with fixed
# unit run times; this is separate from stochastic training forward passes.
def order_signal(model):
    if not isinstance(model,Cyclic): return None
    model.eval(); z=torch.tanh(model.inp(X[test_slice].to(device)))
    with torch.no_grad():
        u=model.f1(z); v=model.f2(z); m=(u+v)/2; u=u-m; v=v-m; e=.1
        ab=z+e*u; ab=ab+e*(model.f2(ab)-model.f1(ab))/2
        ba=z+e*v; ba=ba+e*(model.f1(ba)-model.f2(ba))/2
    return float((ab-ba).norm(dim=1).mean().cpu())

if __name__=='__main__':
    b=train('baseline'); c=train('cyclic')
    # Retrain one cyclic model only to measure order gap consistently.
    torch.manual_seed(SEED); cm=Cyclic().to(device); opt=torch.optim.Adam(cm.parameters(),lr=2e-3)
    for _ in range(350):
        opt.zero_grad(); l=((cm(X.to(device))[train_slice]-Y.to(device)[train_slice])**2).mean(); l.backward(); opt.step()
    c['order_signal']=order_signal(cm)
    out={'device':str(device),'baseline':b,'cyclic':c}
    print(json.dumps(out,indent=2)); open('toy_results.json','w').write(json.dumps(out,indent=2))