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

SEEDS=tuple(range(8)); EPOCHS=15; BATCH=128
LRS=[1e-3,3e-3,1e-2]; WDS=[0.0,1e-4]

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def ds(seed):
    # Load the accepted track through bench; restore its documented 4-edge target
    # shape because the read-only adapter flattens all custom regression targets.
    d=get_dataset('budgeted_route_costs', seed, n_train=400, n_test=160)
    d['ytr']=d['ytr'].reshape(400,4); d['yte']=d['yte'].reshape(160,4)
    d['out_dim']=4
    return d

class Net(nn.Module):
    def __init__(self,shape):
        super().__init__(); self.m=make_model('mlp_tiny',shape,4)
    def forward(self,x): return F.softplus(self.m(x))

def base(seed,cfg,ret=False):
    seed_all(1000+seed); d=ds(seed); m=Net(d['input_shape'])
    m,metric,_=train_model(m,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['wd'],log=lambda x:None)
    return (float(metric),m,d) if ret else float(metric)

def subsets(B=1): return [()] + [(i,) for i in range(4)]

def adv_loss(pred,c,tau=.10):
    # Parallel routes: path A=(0,1), B=(2,3); enumerate all budget-1 delays.
    delays=torch.tensor([[0.,0.,0.,0.],[0.,0.,1.2,0.],[0.,1.2,0.,0.],[1.2,0.,0.,0.],[0.,0.,0.,1.2]],device=pred.device)
    vals=[]
    for d in delays:
        ct=c+d; cp=pred+d
        true=torch.stack((ct[:,0]+ct[:,1],ct[:,2]+ct[:,3]),1)
        pp=torch.softmax(-torch.stack((cp[:,0]+cp[:,1],cp[:,2]+cp[:,3]),1)/tau,1)
        vals.append((pp*true).sum(1)-true.min(1).values)
    return torch.stack(vals,1).max(1).values.mean()

def idea(seed,cfg,ret=False):
    seed_all(1000+seed); d=ds(seed); m=Net(d['input_shape']); opt=torch.optim.Adam(m.parameters(),lr=cfg['lr'],weight_decay=cfg['wd'])
    x,y=d['xtr'],d['ytr']
    for ep in range(EPOCHS):
        p=torch.randperm(len(x))
        for j in range(0,len(x),BATCH):
            q=m(x[p[j:j+BATCH]]); c=y[p[j:j+BATCH]]
            loss=adv_loss(q,c)+.01*((q-c)**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
    m.eval()
    with torch.no_grad(): metric=float(((m(d['xte'])-d['yte'])**2).mean())
    return (metric,m,d) if ret else metric

def graph_metrics(m,d):
    m = m.cpu()
    with torch.no_grad(): p=m(d['xte']).numpy(); c=d['yte'].numpy()
    worst=[]; flips=[]; nominal=[]
    for a,b in zip(p,c):
        vals=[]; fs=[]
        for sub in subsets():
            dd=np.zeros(4); dd[list(sub)]=1.2
            ta=float(b[0]+b[1]+dd[0]+dd[1]); tb=float(b[2]+b[3]+dd[2]+dd[3])
            pa=float(a[0]+a[1]+dd[0]+dd[1]); pb=float(a[2]+a[3]+dd[2]+dd[3])
            true_i=0 if ta <= tb else 1; pred_i=0 if pa <= pb else 1
            true_cost=(ta,tb); vals.append(max(0., true_cost[pred_i]-true_cost[true_i])); fs.append(int(true_i!=pred_i))
        worst.append(max(vals)); flips.append(max(fs)); nominal.append(vals[0])
    return {'worst_regret':float(np.mean(worst)),'flip_rate':float(np.mean(flips)),'nominal_regret':float(np.mean(nominal))}

def main():
    grid=[{'lr':lr,'wd':wd} for lr in LRS for wd in WDS]
    baseblock=sweep_baseline(lambda cfg: lambda s:base(s,cfg),grid)
    best=baseblock['best_cfg']
    idea_grid=[best,{'lr':1e-3,'wd':best['wd']},{'lr':1e-2,'wd':best['wd']}]
    # de-duplicate while preserving the required shared union.
    ir=[]; chosen=None
    for cfg in idea_grid:
        r=evaluate(lambda s,cfg=cfg:idea(s,cfg),SEEDS); ir.append({'cfg':cfg,'result':r})
        if chosen is None or r['mean']<chosen['result']['mean']: chosen=ir[-1]
    ideares=chosen['result']
    bm,im=[],[]
    for s in SEEDS:
        _,mb,db=base(s,best,True); _,mi,di=idea(s,chosen['cfg'],True); bm.append(graph_metrics(mb,db)); im.append(graph_metrics(mi,di))
    sig={'baseline_worst_regret_mean':float(np.mean([x['worst_regret'] for x in bm])),'idea_worst_regret_mean':float(np.mean([x['worst_regret'] for x in im])),'predicted_effect':'adversarial training reduces budgeted path flips/regret','confirmed':float(np.mean([x['worst_regret'] for x in im]))<float(np.mean([x['worst_regret'] for x in bm]))}
    rep=make_report('budgeted_route_costs','mlp_tiny',baseblock,ideares,{'custom_track':{'name':'budgeted_route_costs','file':'route_track.py','domain':'graph-decision'},'idea_sweep':ir,'mechanism_signature':sig})
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()