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, evaluate, sweep_baseline, make_report

TRACK='tabular'; MODEL='mlp_tiny'; SEEDS=tuple(range(8))
# Union grid is used for both methods: baseline receives every lr tried by idea.
LRS=[1e-3, 3e-3, 6e-3]
TRAINED_SIGNATURE=[]
EPOCHS=18; BATCH=64

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 loss_fn(ds): return nn.MSELoss()

def train_baseline(seed, cfg):
    seed_all(seed); ds=get_dataset(TRACK, seed, 400, 200)
    net=make_model(MODEL, ds['input_shape'], ds['out_dim'])
    # Baseline is standard Adam; custom loop is used so the model and minibatch
    # ordering are exactly matched to the intervention.
    return run(net, ds, cfg['lr'], pulse=False, seed=seed)[0]

def run(net, ds, lr, pulse=False, seed=0):
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        return _run(net, ds, lr, pulse, seed, device)
    except RuntimeError:
        return _run(net.cpu(), ds, lr, pulse, seed, 'cpu')

def _run(net, ds, lr, pulse, seed, device):
    net=net.to(device); x=ds['xtr'].to(device); y=ds['ytr'].to(device)
    opt=torch.optim.Adam(net.parameters(), lr=lr)
    lf=loss_fn(ds); rng=torch.Generator(device=device); rng.manual_seed(seed+991)
    qs=[]; responses=[]; triggers=0; recent=[]
    for ep in range(EPOCHS):
        net.train(); perm=torch.randperm(len(x), generator=rng, device=device)
        # Probe every 4 epochs. q is squared parameter displacement / lr.
        pulse_now=False
        if pulse and ep > 0 and ep % 4 == 0:
            # short temperature pulse = doubled batch noise proxy by halving batch
            pulse_now=True
        batch=BATCH//2 if pulse_now else BATCH
        epq=[]
        for i in range(0,len(x),batch):
            idx=perm[i:i+batch]; old=[p.detach().clone() for p in net.parameters() if p.requires_grad]
            out=net(x[idx].view(len(idx),-1)); loss=lf(out,y[idx])
            opt.zero_grad(); loss.backward(); opt.step()
            q=sum(((p.detach()-o)**2).sum().item() for p,o in zip((p for p in net.parameters() if p.requires_grad),old))/max(lr,1e-12)
            epq.append(q)
        qmean=float(np.mean(epq)); qs.append(qmean)
        if pulse_now:
            base=float(np.mean(recent[-2:])) if recent else qmean
            excess=qmean-base
            responses.append({'epoch':ep,'q_ss':base,'q_pulse':qmean,'C_train':excess/(1.0/(len(x)))})
            # detector only changes training when dissipation is an outlier.
            if len(recent)>=2 and qmean > 1.8*float(np.median(recent[-4:])):
                for g in opt.param_groups: g['lr']*=0.7
                triggers+=1
        else: recent.append(qmean)
    net.eval()
    with torch.no_grad(): metric=float(((net(ds['xte'].to(device)).view(-1,1)-ds['yte'].to(device))**2).mean())
    return metric, {'q':qs,'responses':responses,'triggers':triggers}

def train_idea(seed,cfg):
    seed_all(seed); ds=get_dataset(TRACK, seed, 400, 200)
    net=make_model(MODEL, ds['input_shape'], ds['out_dim'])
    metric, detail = run(net, ds, cfg['lr'], pulse=True, seed=seed)
    TRAINED_SIGNATURE.append(detail)
    return metric

def main():
    # Cheap numerical verification of the transfer formula and stability claim.
    eta=np.linspace(.2,1.99,10); lam=1.; T=1.; delta=.05; P=20
    Cs=[]
    for e in eta:
        a=1-e*lam; v=2*e*T/(1-a*a); v2=v; q0=4*T/(1+a); ex=0
        for _ in range(P):
            v2=a*a*v2+2*e*(T+delta); ex += (1-a)**2*v2/e+2*(T+delta)-q0
        Cs.append(ex/delta)
    sanity={'stable_below_2':bool(np.all(np.abs(1-eta)<1)), 'C_monotone':bool(np.all(np.diff(Cs)>0)), 'C_values':Cs}
    grid=[{'lr':v} for v in LRS]
    base=sweep_baseline(lambda c: (lambda s: train_baseline(s,c)),grid)
    # Three idea settings are exactly the shared union grid, with best selected.
    idea_trials=[]
    for c in grid:
        r=evaluate(lambda s,c=c: train_idea(s,c),SEEDS)
        idea_trials.append({'cfg':c,'mean':r['mean']})
    best=min(idea_trials,key=lambda z:z['mean'])['cfg']
    idea=evaluate(lambda s: train_idea(s,best),SEEDS)
    ratios=[r['q_pulse']/r['q_ss'] for d in TRAINED_SIGNATURE for r in d.get('responses',[]) if r['q_ss']>1e-12]
    trained={'n_pulses':len(ratios), 'mean_pulse_to_baseline_q_ratio':float(np.mean(ratios)) if ratios else None, 'mean_triggers':float(np.mean([d.get('triggers',0) for d in TRAINED_SIGNATURE])) if TRAINED_SIGNATURE else 0.0}
    signature={'prediction':'temperature/minibatch-noise pulse produces transient excess update dissipation', 'predicted':{'pulse_to_baseline_q_ratio':'>1'}, 'observed_on_trained_models':trained, 'confirmed':bool(ratios and trained['mean_pulse_to_baseline_q_ratio']>1.0), 'toy_sanity':sanity}
    rep=make_report(TRACK,MODEL,base,idea,signature)
    rep['idea']['sweep']=idea_trials; rep['sanity_check']=sanity
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()