import json, random
import numpy as np
import torch
from torch import nn
import sys
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report

TRACK='token_expert_sequence'
MODEL='mlp_tiny'
SEEDS=tuple(range(8))
EPOCHS=25

class DissipativeNet(nn.Module):
    """Same mlp_tiny backbone, with a differentiable reset/cycle latent readout."""
    def __init__(self, input_shape, out_dim=2, reset_rate=2.0, drive=.12, steps=16):
        super().__init__()
        # Exact mlp_tiny: 24 -> 64 -> 64 -> 2.
        dim=int(np.prod(input_shape))
        self.body=nn.Sequential(nn.Linear(dim,64),nn.ReLU(),nn.Linear(64,64),nn.ReLU(),nn.Linear(64,out_dim))
        self.r=float(reset_rate); self.drive=float(drive); self.steps=int(steps)
    def forward(self,x):
        X=self.body(x.reshape(x.shape[0],-1)); K=X.shape[1]
        p=torch.softmax(X,dim=1)
        B=x.shape[0]; 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
        d=(X[:,None,:]-X[:,:,None])/2
        W=torch.exp(d)*(1-torch.eye(K,device=x.device,dtype=x.dtype)[None])
        rates[:,1:,1:]=W
        # Directed reset-sector cycle 0 -> 1 -> ... -> K-1 -> 0.
        for i in range(K-1): rates[:,i,i+1]+=self.drive
        rates[:,K-1,0]+=self.drive
        total=rates.sum(2)
        dt=.45/(total.max().detach()+1e-6)
        P=torch.eye(K+1,device=x.device,dtype=x.dtype)[None]-dt*torch.diag_embed(total)+dt*rates
        q=torch.zeros(B,K+1,device=x.device,dtype=x.dtype); q[:,0]=1.
        for _ in range(self.steps): q=torch.bmm(q[:,None,:],P).squeeze(1).clamp_min(1e-8); q=q/q.sum(1,keepdim=True)
        return torch.log(q[:,1:].clamp_min(1e-8))

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

def baseline_fn(cfg):
    def run(seed):
        set_seed(seed); d=get_dataset(TRACK,seed,n_train=400,n_test=400)
        net=nn.Sequential(nn.Flatten(), make_model(MODEL,d['input_shape'],d['out_dim']))
        _,metric,_=train_model(net,d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *_:None)
        return metric
    return run

def idea_fn(cfg):
    def run(seed):
        set_seed(seed); d=get_dataset(TRACK,seed,n_train=400,n_test=400)
        net=DissipativeNet(d['input_shape'],d['out_dim'],cfg['r'])
        _,metric,_=train_model(net,d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *_:None)
        return metric
    return run

def signature(cfg):
    vals=[]; currents=[]
    for seed in SEEDS:
        set_seed(seed); d=get_dataset(TRACK,seed,n_train=400,n_test=400)
        net=DissipativeNet(d['input_shape'],d['out_dim'],cfg['r'])
        net,_,_=train_model(net,d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *_:None)
        with torch.no_grad():
            dev=next(net.parameters()).device; xt=torch.as_tensor(d['xte'],dtype=torch.float32,device=dev)
            X=net.body(xt.reshape(len(d['xte']),-1)); p=torch.softmax(X,1).cpu().numpy()
            q=torch.exp(net(xt)).cpu().numpy()
        vals.append(float(np.abs(q-p).mean()))
        currents.append(float(np.abs(q[:,0]-q[:,1]).mean()))
    # A model-behavior check, not an analytical identity: rapid reset should
    # reduce occupation mismatch relative to the nearby low-reset setting.
    low=[]
    for seed in SEEDS:
        set_seed(seed); d=get_dataset(TRACK,seed,n_train=400,n_test=400)
        net=DissipativeNet(d['input_shape'],d['out_dim'],.5)
        net,_,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=128,log=lambda *_:None)
        with torch.no_grad():
            dev=next(net.parameters()).device; xt=torch.as_tensor(d['xte'],dtype=torch.float32,device=dev)
            X=net.body(xt.reshape(len(d['xte']),-1)); p=torch.softmax(X,1); q=torch.exp(net(xt))
        low.append(float(torch.abs(q-p).mean()))
    return {'predicted_effect':'occupation mismatch decreases as reset_rate increases', 'mismatch_at_r':float(np.mean(vals)), 'mismatch_at_r_0.5':float(np.mean(low)), 'observed_cycle_proxy':float(np.mean(currents)), 'confirmed':bool(np.mean(vals)<np.mean(low))}

def main():
    # Union parity: all idea learning rates are also baseline candidates.
    grid=[{'lr':lr,'epochs':EPOCHS} for lr in (1e-3,3e-3,1e-2)]
    base=sweep_baseline(baseline_fn,grid,seeds=(0,1,2,3))
    # Three idea settings at the selected baseline lr; reset is the method knob.
    lr=base['best_cfg']['lr']
    idea_cfgs=[{'lr':lr,'epochs':EPOCHS,'r':r} for r in (.5,2.,8.)]
    idea_runs=[(cfg,evaluate(idea_fn(cfg),seeds=SEEDS)) for cfg in idea_cfgs]
    best_cfg,idea_res=min(idea_runs,key=lambda z:z[1]['mean'])
    extra=signature(best_cfg)
    rep=make_report(TRACK,MODEL,base,idea_res,extra=extra)
    rep['idea_sweep']=[{'cfg':cfg,'result':res} for cfg,res in idea_runs]
    rep['custom_track']={'name':TRACK,'file':'/home/maxwelhelp/all/math2nn/bench/custom_tracks/token_expert_sequence.py','domain':'moe-routing'}
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()