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

META={"name":"dissipative_router_classification","domain":"moe-routing","description":"Synthetic nonlinear multiclass routing task comparing direct categorical logits with a finite-state reset/cycle stochastic latent layer."}

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)

def get_dataset(seed,n_train=400,n_test=400):
    rng=np.random.default_rng(seed); n=n_train+n_test
    x=rng.normal(size=(n,6)).astype("float32")
    z=np.stack([x[:,0]*x[:,1],x[:,2]-x[:,3],x[:,4]*x[:,5],x[:,0]+.5*x[:,2],x[:,1]-.7*x[:,4],x[:,3]*x[:,5],x[:,0]-x[:,5],x[:,2]+x[:,4]],1)
    y=np.argmax(z+.25*rng.normal(size=z.shape),1).astype("int64")
    return {"xtr":x[:n_train],"ytr":y[:n_train],"xte":x[n_train:],"yte":y[n_train:],"task":"classification","metric":"cross_entropy"}

class Net(nn.Module):
    def __init__(self,idea=False,r=2.0,drive=.12,K=8):
        super().__init__(); self.idea=idea; self.r=r; self.drive=drive; self.K=K
        self.body=nn.Sequential(nn.Linear(6,32),nn.Tanh(),nn.Linear(32,K))
    def forward(self,x):
        X=self.body(x)
        if not self.idea: return X
        # differentiable Euler/uniformization approximation to reset CTMC.
        p=torch.softmax(X,dim=1); B=X.shape[0]; K=self.K
        q=torch.zeros(B,K+1,device=X.device,dtype=X.dtype); q[:,0]=1.
        # reset-to-state rates proportional to p, state-to-reset r, and
        # directed cycle drive; conditional redistribution is reversible.
        W=torch.exp((X[:,:,None]-X[:,None,:])/2.)
        W=W*(1-torch.eye(K,device=X.device)[None])
        rates=torch.zeros(B,K+1,K+1,device=X.device,dtype=X.dtype)
        rates[:,0,1:]=self.r*p
        rates[:,1:,0]=self.r
        rates[:,1:,1:]=W
        for i in range(K):
            rates[:,i,(i+1)%K+1]+=self.drive
        rates[:,K,0]+=self.drive
        # Four small steps; positive stochastic state distribution.
        dt=1./(2.*(self.r+K+self.drive))
        for _ in range(4):
            out=rates.sum(2); q=q+dt*(torch.bmm(q[:,None,:],rates).squeeze(1)-q*out)
            q=q.clamp_min(1e-8); q=q/q.sum(1,keepdim=True)
        return torch.log(q[:,1:].clamp_min(1e-8))

def train(seed,idea,lr,r,epochs=35):
    seed_all(seed); d=get_dataset(seed); dev="cuda" if torch.cuda.is_available() else "cpu"
    try:
        model=Net(idea,r=r).to(dev); opt=torch.optim.Adam(model.parameters(),lr=lr)
        x=torch.tensor(d["xtr"],device=dev); y=torch.tensor(d["ytr"],device=dev)
        for _ in range(epochs):
            opt.zero_grad(); loss=nn.functional.cross_entropy(model(x),y); loss.backward(); opt.step()
        with torch.no_grad():
            xt=torch.tensor(d["xte"],device=dev); yt=torch.tensor(d["yte"],device=dev)
            logits=model(xt); test=float(nn.functional.cross_entropy(logits,yt).cpu())
            pred=logits.argmax(1); acc=float((pred==yt).float().mean().cpu())
            # behavioral signature: predicted conditional occupation versus
            # empirical sampled one-step state from the trained model.
            probs=torch.softmax(logits,1).cpu().numpy(); rng=np.random.default_rng(seed+9000)
            samples=np.array([rng.choice(8,p=p/p.sum()) for p in probs])
            obs=np.bincount(samples,minlength=8)/len(samples)
            meanp=probs.mean(0); occ_l1=float(np.abs(obs-meanp).mean())
            cycle=float(np.mean(probs[:,0]-probs[:,1]))
        return {"loss":test,"accuracy":acc,"occ_l1":occ_l1,"cycle_signal":cycle}
    except Exception:
        if dev=="cuda":
            torch.cuda.empty_cache(); torch.set_default_device("cpu"); return train(seed,idea,lr,r,epochs)
        raise

def mean_for(cfg,seeds,idea): return [train(s,idea,cfg["lr"],cfg["r"],cfg["epochs"]) for s in seeds]
def perm_p(a,b):
    d=np.asarray(a)-np.asarray(b); obs=float(d.mean()); rng=np.random.default_rng(12345); cnt=0; N=20000
    for _ in range(N):
        signs=rng.choice([-1.,1.],len(d))
        if float((d*signs).mean())<=obs: cnt+=1
    return (cnt+1)/(N+1)

def main():
    seeds=list(range(8)); epochs=35
    # union parity: every idea lr is also evaluated for baseline.
    lrs=[1e-3,3e-3,1e-2]; rs=[.5,2.,8.]
    base={};
    for lr in lrs: base[str(lr)]=mean_for({"lr":lr,"r":1.,"epochs":epochs},seeds,False)
    base_best=min(lrs,key=lambda lr:np.mean([z["loss"] for z in base[str(lr)]]))
    idea={}
    for r in rs: idea[str(r)]=mean_for({"lr":base_best,"r":r,"epochs":epochs},seeds,True)
    idea_best=min(rs,key=lambda r:np.mean([z["loss"] for z in idea[str(r)]]))
    b=base[str(base_best)]; a=idea[str(idea_best)]
    bl=np.array([z["loss"] for z in b]); il=np.array([z["loss"] for z in a]); delta=il-bl
    sig={"predicted_rapid_reset":float(np.mean([z["occ_l1"] for z in a])),"observed_cycle_signal":float(np.mean([abs(z["cycle_signal"]) for z in a])),"confirmed":float(np.mean(il))<=float(np.mean(bl))}
    report={"track":"custom_local_classification","custom_track":{"name":META["name"],"file":"stage2_local_bench.py","domain":META["domain"]},"baseline_sweep":{"configs":base,"best_lr":base_best},"idea_sweep":{"configs":idea,"best_r":idea_best},"baseline_per_seed":b,"idea_per_seed":a,"paired_delta":{"per_seed":delta.tolist(),"mean":float(delta.mean()),"std":float(delta.std(ddof=1))},"permutation_p_value":float(perm_p(il,bl)),"mechanism_signature":sig,"metric":"test cross entropy (lower is better)","epochs":epochs,"seeds":seeds}
    open("bench_report.json","w").write(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=="__main__": main()