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

TRACK = 'expert_balanced_regression'
D = 8

def vdc(m):
    z, f = 0.0, 0.5
    while m:
        z += (m & 1) * f
        m >>= 1; f *= 0.5
    return z

def route(mode, start, n, D, x=None, rho=1.0):
    rng = np.random.default_rng(start + 17015)
    if mode == 'vdc':
        return np.asarray([int(D*vdc(start+i)) % D if (i == 0 or rng.random() < rho) else int(rng.integers(D)) for i in range(n)])
    if mode == 'random':
        return rng.integers(0, D, size=n)
    return x.argmax(1).detach().cpu().numpy()

class MoE(nn.Module):
    def __init__(self, mode, rho):
        super().__init__(); self.mode=mode; self.rho=rho; self.pos=0
        self.router=nn.Linear(8,D)
        self.experts=nn.ModuleList([nn.Sequential(nn.Linear(8,32),nn.ReLU(),nn.Linear(32,1)) for _ in range(D)])
        self.last=None
    def forward(self,x):
        logits=self.router(x); n=len(x)
        if self.mode == 'greedy': r=route('greedy',self.pos,n,D,logits)
        else: r=route(self.mode,self.pos,n,D,rho=self.rho)
        self.pos += n; self.last=r
        y=torch.empty((n,1),device=x.device)
        for e in range(D):
            ii=np.flatnonzero(r==e)
            if len(ii): y[ii]=self.experts[e](x[ii])
        return y

def run(seed, mode, cfg, signature=False):
    torch.manual_seed(seed); np.random.seed(seed)
    d=get_dataset(TRACK, seed=seed, n_train=400, n_test=400)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net=MoE(mode,cfg.get('rho',1.0)).to(dev)
        x,y=d['xtr'].to(dev),d['ytr'].to(dev)
        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
        for _ in range(cfg['epochs']):
            p=torch.randperm(len(x),device=dev)
            for j in range(0,len(x),128):
                ii=p[j:j+128]; loss=((net(x[ii])-y[ii])**2).mean()
                opt.zero_grad(); loss.backward(); opt.step()
        net.eval(); xt,yt=d['xte'].to(dev),d['yte'].to(dev)
        with torch.no_grad(): metric=float(((net(xt)-yt)**2).mean())
        if not signature: return metric
        net.pos=0; counts=np.zeros(D,int); mx=0.; total=0
        with torch.no_grad():
            for j in range(0,len(xt),128):
                net(xt[j:j+128])
                for e in net.last:
                    counts[e]+=1; total+=1
                    mx=max(mx,float(np.max(np.abs(counts-total/D))))
        return metric, {'max_leaf_prefix_discrepancy':mx,'final_load_variance':float(np.var(counts))}
    except RuntimeError:
        dev='cpu'; net=MoE(mode,cfg.get('rho',1.0)).to(dev)
        x,y=d['xtr'],d['ytr']; opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
        for _ in range(cfg['epochs']):
            p=torch.randperm(len(x))
            for j in range(0,len(x),128):
                ii=p[j:j+128]; loss=((net(x[ii])-y[ii])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        net.eval()
        with torch.no_grad(): return float(((net(d['xte'])-d['yte'])**2).mean())

def fn(mode,cfg,signature=False):
    return lambda seed: run(seed,mode,cfg,signature)

def main():
    # Every idea learning rate is present in the baseline grid (union parity).
    grid=[{'lr':lr,'epochs':12,'rho':rho} for lr in (0.001,0.003,0.01) for rho in (0.5,1.0)]
    base=sweep_baseline(lambda c: fn('random',c),grid)
    # Three idea settings: baseline best and two nearby method settings.
    idea_cfgs=[base['best_cfg'],{'lr':0.003,'epochs':12,'rho':1.0},{'lr':0.01,'epochs':12,'rho':1.0}]
    ideas=[evaluate(fn('vdc',c,False)) for c in idea_cfgs]
    best=min(zip(idea_cfgs,ideas),key=lambda z:z[1]['mean'])
    cfg,idea=best; idea['config']=cfg
    sigs=[run(s,'vdc',cfg,True)[1] for s in range(8)]
    extra={'prediction':'trained VDC assignments maintain max leaf prefix discrepancy <=1','predicted_max_discrepancy':1.0,'observed_mean_max_discrepancy':float(np.mean([s['max_leaf_prefix_discrepancy'] for s in sigs])),'observed_final_load_variance':float(np.mean([s['final_load_variance'] for s in sigs])),'confirmed':bool(max(s['max_leaf_prefix_discrepancy'] for s in sigs)<=1.05)}
    report=make_report(TRACK,'top1_moe_shared',base,idea,extra)
    report['custom_track']={'name':TRACK,'file':'expert_balanced_regression.py','domain':'moe-routing'}
    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__': main()