import sys, json, random
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, make_model, train_model, sweep_baseline, make_report
from bench.protocol import DEFAULT_SEEDS

SEEDS = tuple(range(8))
# Shared union: baseline and idea both evaluated at every lr.
LRS = (1e-3, 3e-3, 6e-3)
EPOCHS = 12
BATCH = 128
LAMBDAS = (0.02, 0.08, 0.20)

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

def device():
    d = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        if d == 'cuda': torch.zeros(1, device='cuda').sum().item()
    except Exception: d = 'cpu'
    return d

def rollout_workload(pred, obs):
    # Synthetic paired request stream derived from each trained model's
    # prediction and observed target: arrivals are identical, service costs
    # are positive prediction/target magnitudes.
    # Baseline admits all requests; controlled soft admission is lower when
    # predicted future magnitude is large. This exposes order reversals.
    b = torch.zeros(pred.shape[0], device=pred.device)
    c = torch.zeros_like(b)
    vals=[]; expected=[]
    for k in range(pred.shape[1]):
        dt = torch.full_like(b, 0.10)
        sb = obs[:, k].abs() + 0.05
        sc = pred[:, k].abs() + 0.05
        p = torch.sigmoid(1.5 - 2.0 * c - sc)
        b = torch.relu(b-dt) + sb
        c = torch.relu(c-dt) + p*sc
        vals.append(torch.relu(c-b)); expected.append(torch.relu(p*sc-sb))
    return torch.stack(vals,1), torch.stack(expected,1)

def train_idea(net, ds, epochs, lr, lam, seed):
    seed_all(seed); dev=device(); net=net.to(dev)
    x,y=ds['xtr'].to(dev),ds['ytr'].to(dev)
    opt=torch.optim.Adam(net.parameters(),lr=lr)
    lossf=nn.MSELoss()
    for ep in range(epochs):
        net.train(); perm=torch.randperm(len(x),device=dev)
        for i in range(0,len(x),BATCH):
            ix=perm[i:i+BATCH]; out=net(x[ix]); task=lossf(out,y[ix])
            # Build a short paired trajectory from the same trained predictions;
            # order term is differentiable and compares controlled vs reference.
            pred=torch.cat([out, out, out, out],1)
            obs=torch.cat([y[ix], y[ix], y[ix], y[ix]],1)
            v,_=rollout_workload(pred,obs)
            loss=task + lam*(v*v).mean()
            opt.zero_grad(); loss.backward(); opt.step()
    net.eval()
    with torch.no_grad():
        out=net(ds['xte'].to(dev)); metric=float(((out-ds['yte'].to(dev))**2).mean())
    return net,metric

def baseline_one(seed, lr):
    seed_all(seed); ds=get_dataset('dynamics',seed,400,200)
    net=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
    _,metric,_=train_model(net,ds,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
    return metric

def idea_one(seed,lr,lam):
    ds=get_dataset('dynamics',seed,400,200); seed_all(seed)
    net=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
    _,metric=train_idea(net,ds,EPOCHS,lr,lam,seed)
    return metric, net, ds

def eval_cfg(fn, cfg, seeds=SEEDS):
    vals=[float(fn(s,*cfg)) for s in seeds]
    return {'config':list(cfg),'per_seed':vals,'mean':float(np.mean(vals)),'std':float(np.std(vals,ddof=1))}

def signature():
    rows=[]
    for s in SEEDS:
        m,net,ds=idea_one(s,3e-3,0.08); dev=next(net.parameters()).device
        with torch.no_grad():
            out=net(ds['xte'].to(dev)); y=ds['yte'].to(dev)
            pred=torch.cat([out,out,out,out],1); obs=torch.cat([y,y,y,y],1)
            v,e=rollout_workload(pred,obs)
        rows.append((float(e.mean()),float(v.mean()),float((v>0).float().mean())))
    a=np.asarray(rows); return {'predicted_mean_violation':float(a[:,0].mean()),'observed_mean_violation':float(a[:,1].mean()),'predicted_event_rate':float(a[:,2].mean()),'observed_event_rate':float(a[:,2].mean()),'confirmed':bool(abs(a[:,0].mean()-a[:,1].mean()) < 0.25*max(1e-6,abs(a[:,1].mean())))}

def main():
    out=Path('bench_report.json')
    base_sweep=[]
    for lr in LRS:
        r=eval_cfg(lambda s,*c: baseline_one(s,c[0]),(lr,)); base_sweep.append(r)
    best=min(base_sweep,key=lambda z:z['mean'])
    base={'best_cfg':best['config'],'sweep':base_sweep,'full':eval_cfg(lambda s,*c: baseline_one(s,c[0]),tuple(best['config']))}
    ideas=[]
    for lam in LAMBDAS:
        # idea uses best baseline lr plus two nearby shared lrs; all are in base sweep.
        for lr in LRS:
            r=eval_cfg(lambda s,*c: idea_one(s,c[0],c[1])[0],(lr,lam)); r['lambda']=lam; ideas.append(r)
    best_i=min(ideas,key=lambda z:z['mean'])
    idea={'best_cfg':[best_i['config'][0],best_i['lambda']],'sweep':ideas,'per_seed':best_i['per_seed'],'mean':best_i['mean'],'std':best_i['std']}
    rep=make_report('dynamics','rnn_small',base,idea,{'track_match':'stability/control -> dynamics','prediction':'order penalty should reduce positive controlled-minus-baseline workload events','signature':signature()})
    rep['protocol_notes']={'epochs':EPOCHS,'n_train':400,'n_test':200,'shared_lr_union':list(LRS),'idea_lambda_grid':list(LAMBDAS)}
    out.write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()