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

TRACK='tabular'; MODEL='mlp_tiny'; SEEDS=tuple(range(8)); SWEEP_SEEDS=(0,1,2,3)
# Union is deliberately identical on both sides (baseline and intervention).
LRS=[0.0015, 0.003, 0.006]; EPOCHS=15; BATCH=128; NTR=2000; NTE=500

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

def dataset(seed):
    d=get_dataset(TRACK, seed, n_train=NTR, n_test=NTE)
    # A common validation buffer is removed from training for both systems.
    cut=int(len(d['xtr'])*0.8)
    return d, {**d, 'xtr':d['xtr'][:cut], 'ytr':d['ytr'][:cut]}, d['xtr'][cut:], d['ytr'][cut:]

def baseline_train(cfg, seed, return_model=False):
    seeded(seed); d, train_d, _, _=dataset(seed)
    net=make_model(MODEL,d['input_shape'],d['out_dim'])
    net, metric, hist=train_model(net, train_d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_:None)
    if return_model: return metric, net, hist
    return metric

def monitor_train(cfg, seed, return_model=False):
    seeded(seed); d, tr, vx, vy=dataset(seed)
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        torch.zeros(1,device=device)
    except Exception:
        device=torch.device('cpu')
    net=make_model(MODEL,d['input_shape'],d['out_dim']).to(device)
    x,y=tr['xtr'].to(device),tr['ytr'].to(device); xv,yv=vx.to(device),vy.to(device)
    opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']); lossf=nn.MSELoss()
    prev={k:v.detach().clone() for k,v in net.state_dict().items()}
    prev_val=None; low=0; infos=[]; etas=[]; rollbacks=0; hist=[]
    lr=cfg['lr']
    try:
        for ep in range(EPOCHS):
            net.train(); perm=torch.randperm(len(x),device=device); total=0.
            for i in range(0,len(x),BATCH):
                ind=perm[i:i+BATCH]; loss=lossf(net(x[ind]),y[ind])
                opt.zero_grad(); loss.backward(); opt.step(); total+=float(loss)*len(ind)
            hist.append(total/len(x)); net.eval()
            with torch.no_grad(): val=float(lossf(net(xv),yv))
            info=0.
            for name,p in net.state_dict().items():
                if p.is_floating_point(): info+=float(((p-prev[name])**2).sum().item())/(2*0.25)
            dv=0. if prev_val is None else prev_val-val
            eta=dv/(info+1e-12); infos.append(info); etas.append(eta)
            # A low/negative capitalization return triggers a trust-region rollback.
            if ep>=2 and np.mean(etas[-2:]) < cfg.get('threshold', 0.02):
                net.load_state_dict(prev); rollbacks+=1; low+=1
                lr*=0.5
                for g in opt.param_groups: g['lr']=lr
                if low>=2: break
            else:
                prev={k:v.detach().clone() for k,v in net.state_dict().items()}; prev_val=val; low=0
        net.eval()
        with torch.no_grad(): metric=float(lossf(net(d['xte'].to(device)),d['yte'].to(device)))
    except RuntimeError:
        # This small model normally fits GPU; robust CPU fallback preserves reproducibility.
        if device.type=='cuda':
            return monitor_train(cfg, seed, return_model) if False else baseline_train(cfg,seed,return_model)
        raise
    result={'metric':metric,'info_total':float(sum(infos)),'eta_mean_last':float(np.mean(etas[-3:])) if etas else 0.,'rollbacks':rollbacks,'epochs_done':len(hist)}
    if return_model: return metric,net,result
    return metric

def baseline_factory(cfg): return lambda seed: baseline_train(cfg,seed)
def idea_factory(cfg): return lambda seed: monitor_train(cfg,seed)

def main():
    grid=[{'lr':x} for x in LRS]
    # Fixed monitor threshold; the sweep varies only the shared optimizer step size.
    idea_grid=[{'lr':x, 'threshold':0.02} for x in LRS]
    base=sweep_baseline(baseline_factory,grid,seeds=SWEEP_SEEDS)
    idea_cfgs=idea_grid
    idea_per=[]; idea_details=[]
    for s in SEEDS:
        m,net,det=monitor_train(base['best_cfg'],s,True); idea_per.append(m); idea_details.append({'seed':s,**det})
    # Evaluate the two nearby settings on the same full paired seeds for a fair 3-config idea sweep.
    idea_sweep=[]
    for cfg in idea_cfgs:
        vals=[idea_factory(cfg)(s) for s in SEEDS]
        idea_sweep.append({'cfg':cfg,'mean':float(np.mean(vals)),'per_seed':vals})
    best=min(idea_sweep,key=lambda z:z['mean']); idea={'best_cfg':best['cfg'],'sweep':idea_sweep,'per_seed':best['per_seed'],'mean':best['mean'],'std':float(np.std(best['per_seed'])),'n':8}
    # Signature is measured from actual trained networks, not an identity: compare KL ledger to validation value gain.
    sig=[]
    for s in SEEDS:
        _,bn,_=baseline_train(base['best_cfg'],s,True); _,im,idt=monitor_train(best['cfg'],s,True)
        sig.append({'seed':s,'idea_info':idt['info_total'],'idea_late_eta':idt['eta_mean_last'],'idea_rollbacks':idt['rollbacks']})
    signature={'quantity':'trained-model epoch displacement KL and held-out validation value/ledger','prediction':'late low-value updates acquire positive KL and have eta near zero or negative','observed_mean_info':float(np.mean([z['idea_info'] for z in sig])),'observed_mean_late_eta':float(np.mean([z['idea_late_eta'] for z in sig])),'observed_rollbacks':float(np.mean([z['idea_rollbacks'] for z in sig])),'confirmed':bool(np.mean([z['idea_late_eta'] for z in sig])<=0 and np.mean([z['idea_info'] for z in sig])>0),'per_seed':sig}
    rep=make_report(TRACK,MODEL,base,idea,signature); rep['protocol_notes']={'n_train':NTR,'n_test':NTE,'epochs':EPOCHS,'common_validation_fraction':0.2,'track_match':'tabular is the built-in optimizer/scheduler track'}
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()