"""gary-mesh: a MeshNCA-style graph cellular automaton in pure numpy.
One tiny SHARED local rule (~1k params) drives N cells of ANY count: parameters
decouple from neuron count. Cells learn from neighbors via direction-aware
message passing (low-order-harmonic analog of MeshNCA's spherical harmonics),
update asynchronously (Bernoulli mask), and self-organize a target field from a seed.
Trained on a small ring; runs unchanged on much larger rings -> train-small-deploy-large.
Analytic backprop-through-time; Adam. Cites: Mesh Neural Cellular Automata (Pajouheshgar
et al., TOG 2024, arXiv:2311.02820); Growing NCA (Mordvintsev et al., Distill 2020)."""
import numpy as np

def ring_graph(N):
    pos=np.stack([np.cos(2*np.pi*np.arange(N)/N), np.sin(2*np.pi*np.arange(N)/N)],1)
    nbr=[[ (i-1)%N, (i+1)%N ] for i in range(N)]   # each cell: 2 neighbors (1D periodic mesh)
    return pos, np.array(nbr)

def harmonic_weights(pos, nbr):
    """Non-parametric, direction-aware message coefficients [1, cos a, sin a] per edge
    (the l<=1 spherical-harmonic analog in 2D). Fixed, NOT trained -- like MeshNCA."""
    N=pos.shape[0]; W=np.zeros((N,nbr.shape[1],3))
    for i in range(N):
        for k,j in enumerate(nbr[i]):
            d=pos[j]-pos[i]; a=np.arctan2(d[1],d[0])
            W[i,k]=[1.0, np.cos(a), np.sin(a)]
    return W   # N x K x 3

class GaryMesh:
    def __init__(s, C=8, H=32, seed=0):
        r=np.random.default_rng(seed); s.C=C; s.H=H
        s.W1=(r.standard_normal((3*C,H))*np.sqrt(2/(3*C))).astype(np.float64); s.b1=np.zeros(H)
        s.W2=(r.standard_normal((H,C))*0.0).astype(np.float64);  s.b2=np.zeros(C)  # zero-init -> identity start (Growing-NCA trick)
    def params(s): return {"W1":s.W1,"b1":s.b1,"W2":s.W2,"b2":s.b2}
    def nparams(s): return sum(v.size for v in s.params().values())
    def perceive(s, S, nbr, Wh):
        # message_i = sum_k Wh[i,k] (x) s_{nbr[i,k]}  -> per-cell 3C perception vector
        N=S.shape[0]; Z=np.zeros((N,3*s.C))
        for k in range(nbr.shape[1]):
            sj=S[nbr[:,k]]                      # N x C
            for b in range(3):
                Z[:,b*s.C:(b+1)*s.C]+=Wh[:,k,b:b+1]*sj
        return Z
    def step(s, S, nbr, Wh, rng, train=False):
        Z=s.perceive(S,nbr,Wh)
        pre=Z@s.W1+s.b1; A=np.maximum(pre,0); dS=A@s.W2+s.b2
        M=(rng.random(S.shape[0])<0.5).astype(np.float64)[:,None]   # async Bernoulli mask
        S2=S+M*dS
        cache=(S,Z,pre,A,M) if train else None
        return S2, cache
    def rollout(s, N, nbr, Wh, T, rng, train=False):
        S=np.zeros((N,s.C)); S[:,0]=1.0   # seed: channel-0 = 1 everywhere (constant seed)
        caches=[]
        for t in range(T):
            S,c=s.step(S,nbr,Wh,rng,train)
            if train: caches.append(c)
        return S, caches
    def backward(s, S_final, target, caches, nbr, Wh):
        # loss = mean((S_final[:,0]-target)^2); BPTT through the shared rule
        N=S_final.shape[0]; g={k:np.zeros_like(v) for k,v in s.params().items()}
        dS=np.zeros_like(S_final); dS[:,0]=2*(S_final[:,0]-target)/N
        for (S,Z,pre,A,M) in reversed(caches):
            dout=dS*M                                   # through residual mask
            g["b2"]+=dout.sum(0); g["W2"]+=A.T@dout
            dA=dout@s.W2.T; dpre=dA*(pre>0)
            g["b1"]+=dpre.sum(0); g["W1"]+=Z.T@dpre
            dZ=dpre@s.W1.T
            dS_prev=dS.copy()                           # residual path: S2=S+...
            # scatter dZ back to neighbor states
            for k in range(nbr.shape[1]):
                for b in range(3):
                    contrib=Wh[:,k,b:b+1]*dZ[:,b*s.C:(b+1)*s.C]
                    np.add.at(dS_prev, nbr[:,k], contrib)
            dS=dS_prev
        loss=float(np.mean((S_final[:,0]-target)**2))
        return loss,g

class Adam:
    def __init__(s,P,lr=2e-3,b1=0.9,b2=0.999,eps=1e-8):
        s.lr=lr;s.b1=b1;s.b2=b2;s.eps=eps;s.t=0
        s.m={k:np.zeros_like(v) for k,v in P.items()}; s.v={k:np.zeros_like(v) for k,v in P.items()}
    def step(s,P,G):
        s.t+=1
        for k in P:
            s.m[k]=s.b1*s.m[k]+(1-s.b1)*G[k]; s.v[k]=s.b2*s.v[k]+(1-s.b2)*G[k]**2
            mh=s.m[k]/(1-s.b1**s.t); vh=s.v[k]/(1-s.b2**s.t)
            P[k]-=s.lr*mh/(np.sqrt(vh)+s.eps)