import sys, math, 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 make_report, permutation_pvalue

META = {'name': 'expert_balanced_regression', 'domain': 'moe-routing', 'description': 'Regression with a top-1 mixture of equally sized experts; routing balance is the intervention.'}


def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.default_rng(seed)
    n = n_train + n_test
    x = rng.uniform(-1, 1, (n, 8)).astype(np.float32)
    y = (np.sin(3*x[:, 0]) + .45*x[:, 1]**2 - .35*x[:, 2]*x[:, 3] + .2*x[:, 4] - .15*x[:, 5]**3).astype(np.float32)
    y += rng.normal(0, .025, n).astype(np.float32)
    return {'xtr': x[:n_train], 'ytr': y[:n_train,None], 'xte': x[n_train:], 'yte': y[n_train:,None], 'task':'regression', 'metric':'mse', 'out_dim':1}


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


def routes(mode, start, n, D, logits=None, rho=1.0):
    rng = np.random.default_rng(start + 991)
    out=[]
    for i in range(n):
        if mode == 'vdc' and (i == 0 or rng.random() < rho): out.append(int(D*vdc(start+i)) % D)
        elif mode == 'random': out.append(int(rng.integers(D)))
        else: out.append(int(torch.argmax(logits[i]).item()) if logits is not None else int(rng.integers(D)))
    return np.asarray(out, dtype=np.int64)


class Top1MoE(nn.Module):
    def __init__(self, mode='vdc', D=8, rho=1.0):
        super().__init__(); self.mode, self.D, self.rho = mode, D, rho
        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.assignment_index = 0
        self.last_routes = None
    def forward(self, x):
        logits = self.router(x)
        n = x.shape[0]
        if self.mode == 'vdc':
            r = routes('vdc', self.assignment_index, n, self.D, rho=self.rho)
        elif self.mode == 'random':
            r = routes('random', self.assignment_index, n, self.D)
        else: r = logits.argmax(1).detach().cpu().numpy()
        self.assignment_index += n
        self.last_routes = r
        y = torch.empty((n,1), device=x.device)
        for e in range(self.D):
            ix = np.flatnonzero(r == e)
            if len(ix): y[ix] = self.experts[e](x[ix])
        return y


def train(seed, mode, lr, epochs, rho=1.0, return_sig=False):
    torch.manual_seed(seed); np.random.seed(seed)
    d=get_dataset(seed); device='cuda' if torch.cuda.is_available() else 'cpu'
    model=Top1MoE(mode, rho=rho).to(device)
    x=torch.tensor(d['xtr'],device=device); y=torch.tensor(d['ytr'],device=device)
    opt=torch.optim.Adam(model.parameters(),lr=lr)
    for _ in range(epochs):
        model.train(); perm=torch.randperm(len(x),device=device)
        for j in range(0,len(x),128):
            ix=perm[j:j+128]; loss=((model(x[ix])-y[ix])**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
    model.eval(); xt=torch.tensor(d['xte'],device=device); yt=torch.tensor(d['yte'],device=device)
    with torch.no_grad(): metric=float(((model(xt)-yt)**2).mean())
    # Re-test trained model behavior: assignment prefix discrepancy on a fixed evaluation stream.
    model.assignment_index=0; counts=np.zeros(model.D,int); max_disc=0.0; total=0
    with torch.no_grad():
        for j in range(0,len(xt),128):
            model(xt[j:j+128]); rr=model.last_routes
            for e in rr:
                counts[e]+=1; total+=1
                max_disc=max(max_disc,float(np.max(np.abs(counts-total/model.D))))
    return metric, {'max_leaf_prefix_discrepancy':max_disc,'final_load_variance':float(np.var(counts))} if return_sig else None


def eval_cfg(mode, cfg, seeds=range(8), sig=False):
    vals=[]; signatures=[]
    for s in seeds:
        m,q=train(s,mode,cfg['lr'],cfg['epochs'],cfg.get('rho',1.0),sig); vals.append(m)
        if q: signatures.append(q)
    return {'per_seed':[float(v) for v in vals], 'mean':float(np.mean(vals)), 'std':float(np.std(vals,ddof=1)), 'config':cfg, 'signatures':signatures}


def main():
    # Union grid parity: all idea learning rates are included in baseline sweep.
    grid=[{'lr':lr,'epochs':12,'rho':rho} for lr in (1e-3,3e-3,1e-2) for rho in (0.5,1.0)]
    sweep=[]
    for cfg in grid:
        r=eval_cfg('random',cfg,seeds=range(4)); sweep.append({'config':cfg,'mean':r['mean'],'per_seed':r['per_seed']})
    best=min(sweep,key=lambda z:z['mean'])['config']
    base=eval_cfg('random',best,seeds=range(8))
    idea_runs=[eval_cfg('vdc',cfg,seeds=range(8),sig=True) for cfg in grid]
    idea=min(idea_runs,key=lambda z:z['mean'])
    diffs=[a-b for a,b in zip(idea['per_seed'],base['per_seed'])]
    cmp={'delta_mean':float(np.mean(diffs)),'idea_wins':sum(x<0 for x in diffs),'n_pairs':8,'per_seed_diffs':diffs,'p_value':permutation_pvalue(diffs)}
    if cmp['delta_mean']<0 and cmp['p_value']<.05: cmp['verdict']='idea better (significant)'
    elif cmp['delta_mean']>0 and cmp['p_value']<.05: cmp['verdict']='idea worse (significant)'
    else: cmp['verdict']='no significant win'
    cmp['system_worked']=cmp['verdict']=='idea better (significant)'
    sig=idea['signatures']; obs=float(np.mean([q['max_leaf_prefix_discrepancy'] for q in sig])); pred=1.0
    report={'bench_version':1,'track':'expert_balanced_regression','model':'top1_moe_shared','metric_direction':'lower is better','n_seeds':8,'baseline':{'best_cfg':best,'sweep':sweep,'full':base},'idea':idea,'comparison':cmp,'mechanism_signature':{'prediction':'VDC trained-model assignment prefixes have leaf discrepancy <=1, unlike random routing','predicted_max_discrepancy':pred,'observed_mean_max_discrepancy':obs,'confirmed':bool(obs<=1.05)},'custom_track':{'name':'expert_balanced_regression','file':'vdc_bench.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()