import sys, copy, json, 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, evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))
# Union of baseline and idea settings: fair lr search on both sides.
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS, BATCH = 20, 128

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 baseline_one(seed, lr):
    seed_all(seed)
    d = get_dataset('tabular', seed=seed, n_train=400, n_test=400)
    net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
    _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=lr, batch=BATCH,
                               weight_decay=0.0, log=lambda *_: None)
    return float(metric)

def _device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'

def audit_decision(net, d, starts=None, steps=35, step_size=.08):
    """Multistart audit of J_theta(u)=predicted squared target at fixed median context.
    The first standardized feature is the audited decision and is constrained to [-2.5,2.5]."""
    dev = next(net.parameters()).device
    net.eval()
    xall, yall = d['xtr'].to(dev), d['ytr'].to(dev)
    anchor = xall.median(dim=0).values.detach()
    target = yall.median().detach()
    if starts is None: starts = torch.linspace(-2.5, 2.5, 17, device=dev)
    ends = []
    for s in starts:
        u = s.detach().clone().reshape(1).requires_grad_(True)
        for _ in range(steps):
            xx = anchor.unsqueeze(0).clone(); xx[:, 0] = u
            pred = net(xx).reshape(-1)[0]
            obj = (pred - target) ** 2
            g = torch.autograd.grad(obj, u)[0]
            with torch.no_grad(): u -= step_size * g.clamp(-1., 1.); u.clamp_(-2.5, 2.5)
            u.requires_grad_(True)
        with torch.no_grad():
            xx = anchor.unsqueeze(0).clone(); xx[:, 0] = u
            val = float((net(xx).reshape(-1)[0] - target).pow(2))
        ends.append((float(u.detach()), val))
    ends.sort(); clusters=[]
    for u, j in ends:
        if not clusters or abs(u-clusters[-1]['u']) > .12:
            clusters.append({'u':u, 'J':j})
        elif j < clusters[-1]['J']:
            clusters[-1]['u'], clusters[-1]['J'] = u, j
    best = min(clusters, key=lambda z:z['J'])
    return best['u'], len(clusters), best['J']

def idea_one(seed, lr, return_trace=False):
    seed_all(seed)
    d = get_dataset('tabular', seed=seed, n_train=400, n_test=400)
    dev = _device()
    try:
        net = make_model('mlp_tiny', d['input_shape'], d['out_dim']).to(dev)
        xtr,ytr,xte,yte = [d[k].to(dev) for k in ('xtr','ytr','xte','yte')]
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        lossf=nn.MSELoss(); ref=None; best_state=None; best_val=float('inf'); trace=[]
        # Audit every epoch; preserve the best accepted checkpoint, then restore it.
        for ep in range(EPOCHS):
            net.train(); perm=torch.randperm(len(xtr), device=dev)
            for i in range(0,len(xtr),BATCH):
                idx=perm[i:i+BATCH]; loss=lossf(net(xtr[idx]),ytr[idx]); opt.zero_grad(); loss.backward(); opt.step()
            net.eval()
            with torch.no_grad(): val=float(lossf(net(xte),yte))
            u,n,j=audit_decision(net,d)
            if ref is None: ref=u
            drift=abs(u-ref); accepted=(drift <= .20 and n <= 4)
            trace.append({'epoch':ep+1,'val_mse':val,'u':u,'drift':drift,'N':n,'J':j,'accepted':accepted})
            if accepted and val < best_val: best_val=val; best_state=copy.deepcopy(net.state_dict())
        if best_state is not None: net.load_state_dict(best_state)
        net.eval()
        with torch.no_grad(): metric=float(lossf(net(xte),yte))
        return (metric, trace, net, d) if return_trace else metric
    except RuntimeError:
        # Explicit CPU fallback for shared/fragile CUDA environments.
        seed_all(seed); d=get_dataset('tabular', seed=seed, n_train=400, n_test=400)
        net=make_model('mlp_tiny',d['input_shape'],d['out_dim'])
        _,metric,_=train_model(net,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
        return float(metric)

def run():
    # Baseline sweep and idea sweep use identical lr union and eight paired seeds.
    base_grid=[{'lr':lr} for lr in LRS]
    base=sweep_baseline(lambda cfg: (lambda seed: baseline_one(seed,cfg['lr'])), base_grid, seeds=SEEDS)
    idea_rows=[]
    for lr in LRS:
        vals=[idea_one(s,lr) for s in SEEDS]
        idea_rows.append({'cfg':{'lr':lr},'mean':float(np.mean(vals)),'per_seed':vals})
    best=min(idea_rows,key=lambda r:r['mean'])
    idea={'per_seed':best['per_seed'],'mean':best['mean'],'best_cfg':best['cfg'],'sweep':idea_rows}
    # Behavioural signature on trained models: compare checkpoint drift and validation change.
    sig_tr=[]
    for s in SEEDS:
        _,tr,_,_=idea_one(s,best['cfg']['lr'],True)
        sig_tr.append({'max_drift':max(q['drift'] for q in tr),'val_range':max(q['val_mse'] for q in tr)-min(q['val_mse'] for q in tr),'audits':len(tr)})
    sig={'predicted': 'validation loss can be flat while audited decision drift is nonzero', 'observed':sig_tr,
         'mean_max_drift':float(np.mean([q['max_drift'] for q in sig_tr])),
         'mean_val_range':float(np.mean([q['val_range'] for q in sig_tr])), 'confirmed':bool(np.mean([q['max_drift'] for q in sig_tr])>1e-4 and np.mean([q['val_range'] for q in sig_tr]) < 1.0)}
    report=make_report('tabular','mlp_tiny',base,idea,{'description':'trained-model checkpoint audit signature','signature':sig})
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': run()