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

SEED = 17
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')

# Core recurrence from the proposal, with a compact GRU-like parameterization.
class EdgeGRU(nn.Module):
    def __init__(self, inp=2, mem=8):
        super().__init__(); self.mem = mem
        self.z = nn.Linear(inp + mem, mem); self.q = nn.Linear(inp + mem, mem)
        self.h = nn.Linear(inp + mem, mem)
    def forward(self, m, u):
        v = torch.cat([u, m], -1)
        z = torch.sigmoid(self.z(v)); q = torch.sigmoid(self.q(v))
        cand = torch.tanh(self.h(torch.cat([u, q*m], -1)))
        return (1-z)*m + z*cand

class PersistentEdge(nn.Module):
    def __init__(self, mem=8):
        super().__init__(); self.rnn=EdgeGRU(2,mem); self.out=nn.Linear(mem,1)
    def forward(self, vals, active):
        # vals/active: [batch,T], one stable pair per sequence; dictionary is per batch key.
        m=torch.zeros(vals.size(0), self.rnn.mem, device=vals.device); outputs=[]
        for t in range(vals.size(1)):
            u=torch.stack([active[:,t], vals[:,t]],-1)
            # Inactive pairs are not updated: this is dictionary retention.
            new=self.rnn(m,u); m=torch.where(active[:,t,None].bool(),new,m)
            outputs.append(self.out(m).squeeze(-1))
        return torch.stack(outputs,1)

class ActiveMLP(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,1))
    def forward(self, vals, active):
        return self.net(torch.stack([active,vals],-1)).squeeze(-1)

def math_check():
    torch.manual_seed(3)
    g=EdgeGRU(2,16); m=torch.randn(1000,16)*20; u=torch.randn(1000,2)*50
    with torch.no_grad():
        mp=g(m,u)
        # Since candidate is in [-1,1], each coordinate is a convex combination.
        bound=(mp.abs() <= torch.maximum(m.abs(), torch.ones_like(m))).float().mean().item()
        inactive=torch.randn(4,16); update=g(inactive,torch.randn(4,2))
        retained=torch.where(torch.zeros(4,1).bool(),update,inactive)
        exact=(retained==inactive).all().item()
    return {'convex_coordinate_bound_fraction':bound,'inactive_retention_exact':bool(exact),
            'max_abs_before':float(m.abs().max()),'max_abs_after':float(mp.abs().max())}

def data(n):
    # Pair-specific latent contact history. At t=0..2 the value is observable;
    # after a gap it reappears at t=9..11 with no informative current value.
    latent=torch.randint(0,2,(n,)).float()*2-1
    T=12; active=torch.zeros(n,T); vals=torch.zeros(n,T)
    active[:,0:3]=1; active[:,9:12]=1; vals[:,0:3]=latent[:,None]
    target=latent[:,None].expand(n,T)
    return vals.to(device),active.to(device),target.to(device)

def train(model, vals, active, target, steps=700):
    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    model.train()
    for _ in range(steps):
        # fresh batch each step, avoiding memorization of sequence ordering
        v,a,y=data(128); pred=model(v,a)
        loss=((pred[:,9:12]-y[:,9:12])**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        pred=model(vals,active); react=((pred[:,9:12]-target[:,9:12])**2).mean().item()
        one=((pred[:,0:3]-target[:,0:3])**2).mean().item()
        long=((pred-target)**2).mean().item()
    return react,one,long

def timed(model, vals, active, reps=100):
    with torch.no_grad():
        for _ in range(10): model(vals,active)
        if device.type=='cuda': torch.cuda.synchronize()
        t=time.perf_counter()
        for _ in range(reps): model(vals,active)
        if device.type=='cuda': torch.cuda.synchronize()
    return (time.perf_counter()-t)/reps*1000

def main():
    check=math_check(); vals,active,target=data(2048)
    torch.manual_seed(SEED); base=ActiveMLP().to(device)
    b=train(base,vals,active,target)
    torch.manual_seed(SEED); idea=PersistentEdge(8).to(device)
    e=train(idea,vals,active,target)
    base_ms=timed(base,vals[:128],active[:128]); edge_ms=timed(idea,vals[:128],active[:128])
    result={'device':str(device),'math_check':check,
      'baseline':{'reactivation_mse':b[0],'early_mse':b[1],'all_timestep_mse':b[2], 'params':sum(p.numel() for p in base.parameters()),'ms_batch128':base_ms},
      'persistent_edge_gru':{'reactivation_mse':e[0],'early_mse':e[1],'all_timestep_mse':e[2], 'params':sum(p.numel() for p in idea.parameters()),'ms_batch128':edge_ms},
      'reactivation_mse_ratio_idea_over_baseline':e[0]/max(b[0],1e-12),
      'latency_overhead_ratio':edge_ms/max(base_ms,1e-12)}
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()