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

SEED = 234
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def math_check():
    # Two invertible, noncommuting edge operators.
    A = np.array([[1., 1.], [0., 1.]])
    B = np.array([[1., 0.], [1., 1.]])
    ab, ba = A @ B, B @ A
    noncomm = float(np.linalg.norm(ab - ba))
    # Composition law for identity sigma/tau is ordinary ordered multiplication.
    C = np.array([[2., 0.], [0., .5]])
    direct = A @ B @ C
    composed = (A @ B) @ C
    composition_error = float(np.linalg.norm(direct - composed))
    # Same additive aggregate, different ordered holonomy.
    additive_collision = float(np.linalg.norm((A+B) - (B+A)))
    return {"AB_minus_BA_fro": noncomm, "composition_error": composition_error,
            "additive_collision": additive_collision, "AB": ab.tolist(), "BA": ba.tolist()}

class Additive(nn.Module):
    def __init__(self, d=3):
        super().__init__(); self.edge = nn.Parameter(torch.randn(2,d,d)*.35)
        self.head = nn.Sequential(nn.Flatten(), nn.Linear(d*d, 16), nn.Tanh(), nn.Linear(16,1))
    def forward(self, x):
        return self.head(self.edge[x].sum(1)).squeeze(-1)

class Holonomy(nn.Module):
    def __init__(self, d=3):
        super().__init__(); self.edge = nn.Parameter(torch.randn(2,d,d)*.35)
        # Learned reversal and color-switch maps, implemented as conjugations.
        self.R = nn.Parameter(torch.eye(d) + .03*torch.randn(d,d))
        self.S = nn.Parameter(torch.eye(d) + .03*torch.randn(d,d))
        self.head = nn.Sequential(nn.Flatten(), nn.Linear(d*d, 16), nn.Tanh(), nn.Linear(16,1))
    def transform(self, X, step):
        # T is reversal; Sigma is a learned color switch. Alternation follows T^r.
        Ri = torch.linalg.pinv(self.R)
        X = self.R @ X @ Ri if step % 2 else X
        if step % 2: # deterministic two-color channel for this toy path
            Si = torch.linalg.pinv(self.S); X = Si @ X @ self.S
        return X
    def forward(self, x):
        B,L = x.shape; d=self.edge.shape[-1]
        H = torch.eye(d, device=x.device).expand(B,d,d).clone()
        for r in range(L): H = H @ self.transform(self.edge[x[:,r]], r)
        return self.head(H).squeeze(-1)
    def regularizer(self):
        d=self.S.shape[0]; I=torch.eye(d,device=self.S.device)
        return ((self.S@self.S-I)**2).mean()

def make_data(n, seed):
    g=np.random.default_rng(seed); x=g.integers(0,2,size=(n,3),dtype=np.int64)
    # Ordered relation: first two edges must be 0 then 1. Same-count permutations conflict.
    y=((x[:,0]==0)&(x[:,1]==1)).astype(np.float32)
    return torch.tensor(x), torch.tensor(y)

def train(model, xt, yt, xv, yv, epochs=450):
    opt=torch.optim.Adam(model.parameters(),lr=.025,weight_decay=1e-4)
    lossfn=nn.BCEWithLogitsLoss()
    for _ in range(epochs):
        opt.zero_grad(); z=model(xt); loss=lossfn(z,yt)
        if isinstance(model,Holonomy): loss=loss + .01*model.regularizer()
        loss.backward(); opt.step()
    with torch.no_grad():
        pred=(torch.sigmoid(model(xv))>.5).float(); acc=float((pred==yv).float().mean())
        train_acc=float(((torch.sigmoid(model(xt))>.5).float()==yt).float().mean())
    return train_acc,acc

def main():
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        xt,yt=make_data(2048,10); xv,yv=make_data(2048,11)
        xt,yt,xv,yv=[z.to(device) for z in (xt,yt,xv,yv)]
        # identical initialization scale and exact same data for fair comparison
        torch.manual_seed(SEED); base=Additive().to(device)
        torch.manual_seed(SEED); idea=Holonomy().to(device)
        ba, bv=train(base,xt,yt,xv,yv); ia,iv=train(idea,xt,yt,xv,yv)
        result={'device':device,'math':math_check(),
                'baseline':{'train_accuracy':ba,'test_accuracy':bv},
                'idea':{'train_accuracy':ia,'test_accuracy':iv},
                'n_train':len(xt),'n_test':len(xv),'epochs':450}
    except Exception as e:
        if device=='cuda':
            torch.cuda.empty_cache(); device='cpu'
            xt,yt=make_data(2048,10); xv,yv=make_data(2048,11)
            torch.manual_seed(SEED); ba,bv=train(Additive(),xt,yt,xv,yv)
            torch.manual_seed(SEED); ia,iv=train(Holonomy(),xt,yt,xv,yv)
            result={'device':device,'math':math_check(),'baseline':{'train_accuracy':ba,'test_accuracy':bv},'idea':{'train_accuracy':ia,'test_accuracy':iv},'fallback_reason':str(e)}
        else: raise
    print(json.dumps(result,sort_keys=True))

if __name__=='__main__': main()