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

SEEDS=tuple(range(8)); SWEEP_SEEDS=tuple(range(4)); NTR=400; NTE=200; EPOCHS=5; BATCH=128
# CPU is selected intentionally: this workload is tiny and avoids shared-GPU fallback variability.
DEVICE='cpu'

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

def estimate_gains(net,x,device):
    z=x[:32].detach().to(device).requires_grad_(True)
    out=net(z).sum(); (gr,)=torch.autograd.grad(out,z)
    a=gr.detach().abs().reshape(-1,8,3).mean((0,2)).cpu().numpy()+1e-8
    return np.outer(a,1.0/a)

def greedy_clusters(G,max_size=3):
    cs=[{i} for i in range(8)]
    while True:
        best=None
        for a in range(len(cs)):
            for b in range(a):
                if len(cs[a]|cs[b])>max_size: continue
                score=max((math.log(float(G[i,j])+1e-8)+math.log(float(G[j,i])+1e-8)
                           for i in cs[a] for j in cs[b]),default=-1e99)
                if best is None or score>best[0]: best=(score,a,b)
        if best is None or best[0]<math.log(.01): break
        _,a,b=best; m=cs[a]|cs[b]
        cs=[c for k,c in enumerate(cs) if k not in (a,b)]+[m]
    return sorted([sorted(c) for c in cs],key=lambda z:z[0])

def baseline_run(seed,lr):
    seed_all(seed); d=get_dataset('dynamics',seed,n_train=NTR,n_test=NTE)
    net=make_model('rnn_small',d['input_shape'],d['out_dim'])
    _,m,_=train_model(net,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
    return float(m)

def idea_run(seed,lr,return_sig=False):
    seed_all(seed); d=get_dataset('dynamics',seed,n_train=NTR,n_test=NTE)
    net=make_model('rnn_small',d['input_shape'],d['out_dim']).to(DEVICE)
    x,y=d['xtr'].to(DEVICE),d['ytr'].to(DEVICE)
    q=nn.Parameter(torch.zeros(8,device=DEVICE)); opt=torch.optim.Adam(list(net.parameters())+[q],lr=lr); lossf=nn.MSELoss()
    for ep in range(EPOCHS):
        net.eval(); G=estimate_gains(net,x,DEVICE); gt=torch.as_tensor(G,dtype=torch.float32)
        dd=torch.exp(q-q.mean()); norm=dd.view(1,8,1)
        P=greedy_clusters(G); owner={i:k for k,c in enumerate(P) for i in c}
        cut=sum(G[i,j] for i in range(8) for j in range(8) if owner[i]!=owner[j])
        perm=torch.randperm(len(x)); net.train()
        for st in range(0,len(x),BATCH):
            ix=perm[st:st+BATCH]
            dd_batch=torch.exp(q-q.mean()); norm_batch=dd_batch.view(1,8,1)
            mu_batch=torch.max((gt@dd_batch)/dd_batch)
            pred=net(x[ix].view(-1,8,3)*norm_batch)
            loss=lossf(pred,y[ix])+0.01*torch.nn.functional.softplus(mu_batch-1.0+0.02)+1e-4*cut
            opt.zero_grad(); loss.backward(); opt.step()
    net.eval(); xt=d['xte'].to(DEVICE)
    with torch.no_grad(): metric=float(((net(xt)-d['yte'].to(DEVICE))**2).mean())
    G2=estimate_gains(net,xt,DEVICE); dd2=torch.exp(q-q.mean()).detach().numpy()
    sig={'mu':float(np.max((G2@dd2)/dd2)),'unscaled_mu':float(np.max(G2.sum(1))),'clusters':greedy_clusters(G2),'gain_mean':float(G2.mean())}
    return (metric,sig) if return_sig else metric

def main():
    grid=[{'lr':v} for v in (1e-3,3e-3,6e-3)]
    base=sweep_baseline(lambda c: lambda s: baseline_run(int(s),c['lr']),grid,seeds=SWEEP_SEEDS)
    irs=[(c,evaluate(lambda s: idea_run(int(s),c['lr']),SWEEP_SEEDS)['mean']) for c in grid]
    best_cfg=min(irs,key=lambda z:z[1])[0]; vals=[]; sigs=[]
    for s in SEEDS:
        v,sg=idea_run(s,best_cfg['lr'],True); vals.append(v); sigs.append(sg)
    idea={'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':8,'best_cfg':best_cfg,'sweep':[{'cfg':c,'mean':m} for c,m in irs]}
    mm=float(np.mean([z['mu'] for z in sigs])); uu=float(np.mean([z['unscaled_mu'] for z in sigs]))
    extra={'prediction':'trained-model gain weighting lowers normalized local certificate versus q=0','observed_normalized_mu':mm,'observed_unscaled_mu':uu,'relative_reduction':float((uu-mm)/max(uu,1e-12)),'clusters_from_trained_models':sigs[0]['clusters'],'confirmed':bool(mm<uu)}
    print(json.dumps(make_report('dynamics','rnn_small',base,idea,extra),indent=2))
if __name__=='__main__': main()