import json, random
import numpy as np

SEED=7
np.random.seed(SEED); random.seed(SEED)
N=16; q0=0; q1=5; r0=1; r1=2; attractor=15

def make_maps():
    A=np.full(N,attractor,dtype=int); B=np.full(N,attractor,dtype=int)
    A[q0]=r0; A[q1]=r1; A[r0]=attractor; A[r1]=attractor
    B[r0]=q1; B[r1]=q0; B[q0]=attractor; B[q1]=attractor
    return A,B

def cycles_of(f):
    seen=set(); cyc=[]
    for start in range(N):
        if start in seen: continue
        path=[]; pos={}; x=start
        while x not in pos and x not in seen:
            pos[x]=len(path); path.append(x); x=int(f[x])
        if x in pos: cyc.append(path[pos[x]:])
        seen.update(path)
    return cyc

def closure(gens):
    ident=tuple(range(N)); seen={ident}; todo=[ident]
    while todo:
        x=todo.pop()
        for g in gens:
            y=tuple(g[list(x)])
            if y not in seen: seen.add(y); todo.append(y)
    return seen

def audit():
    A,B=make_maps(); G=B[A]
    return {'q0':(0,0),'q1':(1,1),'A_q0_q1':bool(A[q0]==q1),
      'composite_q0':int(G[q0]),'composite_q1':int(G[q1]),'swap':bool(G[q0]==q1 and G[q1]==q0),
      'A_cycle_lengths':sorted(map(len,cycles_of(A))),'B_cycle_lengths':sorted(map(len,cycles_of(B))),
      'composite_cycle_lengths':sorted(map(len,cycles_of(G))),'monoid_size':len(closure([A,B]))}

def stochastic_transition(f, eps):
    P=np.full((N,N),eps/(N-1),float)
    for i,j in enumerate(f): P[i,j]=1-eps
    return P

def sweeps():
    A,B=make_maps(); G=B[A]; rows=[]
    # Prediction 1: each frozen generator is aperiodic.
    # Prediction 2: preserving both directions of the two-leg word costs four
    # independent transitions, hence reliability (1-eps)^4.
    for eps in [0,.01,.03,.05,.1,.2,.3,.4]:
        PA=stochastic_transition(A,eps); PB=stochastic_transition(B,eps)
        p=(PA[q0,r0]*PB[r0,q1])*(PA[q1,r1]*PB[r1,q0])
        rows.append({'eps':eps,'observed':float(p),'predicted_(1-eps)^4':(1-eps)**4})
    # Prediction 3: softening continuously reduces exchange fidelity, while
    # the exact deterministic limit is a perfect Z2 orbit.
    temp=[]
    for tau in [0,.01,.03,.05,.1,.2,.3,.4]:
        PA=stochastic_transition(A,tau); PB=stochastic_transition(B,tau)
        p0=(PA[q0]@PB)[q1]; p1=(PA[q1]@PB)[q0]
        temp.append({'tau':tau,'observed_fidelity':float((p0+p1)/2)})
    x=q0; orbit=[]
    for _ in range(8): x=int(G[x]); orbit.append(x)
    return rows,temp,orbit

def finite_delayed_copy(noise, n=20000, delay=6):
    # The target word is G repeated delay times. An even delay returns the bit.
    A,B=make_maps(); G=B[A]; correct=0
    for _ in range(n):
        bit=np.random.randint(2); state=q0 if bit==0 else q1
        for _ in range(delay):
            state=(np.random.randint(N) if np.random.random()<noise else int(G[state]))
        pred=0 if state==q0 else (1 if state==q1 else -1)
        correct += (pred==bit)
    return correct/n

def gru_delayed_copy(force_cpu=False):
    try:
        import torch
        torch.manual_seed(SEED); torch.set_num_threads(2)
        dev='cpu' if force_cpu else ('cuda' if torch.cuda.is_available() else 'cpu')
        T=12; n=600; bs=64
        X=torch.zeros(n,T,3,device=dev); y=torch.zeros(n,dtype=torch.long,device=dev)
        for i in range(n):
            bit=np.random.randint(2); X[i,0,bit]=1; X[i,1:,2]=1; y[i]=bit
        model=torch.nn.GRU(3,8,batch_first=True).to(dev); head=torch.nn.Linear(8,2).to(dev)
        opt=torch.optim.Adam(list(model.parameters())+list(head.parameters()),lr=.02)
        for _ in range(80):
            ix=torch.randperm(n,device=dev)
            for j in range(0,n,bs):
                h,_=model(X[ix[j:j+bs]]); loss=torch.nn.functional.cross_entropy(head(h[:,-1]),y[ix[j:j+bs]])
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): acc=(head(model(X)[0][:,-1]).argmax(1)==y).float().mean().item()
        return {'accuracy':acc,'device':dev,'parameters':sum(p.numel() for p in model.parameters())+sum(p.numel() for p in head.parameters())}
    except Exception:
        if not force_cpu: return gru_delayed_copy(True)
        return {'accuracy':None,'device':'cpu'}

def main():
    noise,temp,orbit=sweeps()
    idea={str(e):finite_delayed_copy(e) for e in [0,.05,.1,.2,.3]}
    out={'seed':SEED,'audit':audit(),'noise_sweep':noise,'temperature_sweep':temp,
      'composite_orbit_8_steps':orbit,'delayed_copy_delay':6,
      'idea_accuracy_under_transition_noise':idea,'gru_baseline':gru_delayed_copy()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()