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

SEEDS = tuple(range(8))
EPOCHS = 12
NTR, NTE = 800, 300
BATCH = 128
LRS = [1e-3, 3e-3, 6e-3]
WDS = [0.0, 1e-4]
# Same lr/step-size union is used by both systems; q is the intervention knob.
QS = [2.0, 4.0, 8.0]


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.CrossEntropyLoss() if ds['task'] == 'classification' else nn.MSELoss()


def flat_params(net):
    return torch.cat([p.detach().flatten() for p in net.parameters()])


def set_flat(net, vec):
    off = 0
    with torch.no_grad():
        for p in net.parameters():
            n = p.numel(); p.copy_(vec[off:off+n].view_as(p)); off += n


def probe_barrier(net, best_vec, ds, lf, device):
    """Empirical loss barrier along current -> best parameter path."""
    cur = flat_params(net)
    if best_vec is None: return 0.0
    vals = []
    was = net.training; net.eval()
    with torch.no_grad():
        for a in (0.0, .25, .5, .75, 1.0):
            set_flat(net, cur*(1-a) + best_vec*a)
            vals.append(float(lf(net(ds['xtr'][:256].to(device)), ds['ytr'][:256].to(device))))
        set_flat(net, cur)
    if was: net.train()
    endpoint = max(vals[0], vals[-1])
    return max(0.0, max(vals) - endpoint)


def train_baseline(ds, lr, wd, seed, return_trace=False):
    seed_all(seed); device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net = make_model('rnn_small', ds['input_shape'], 1).to(device)
        opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd)
        lf = loss_fn(ds); x, y = ds['xtr'].to(device), ds['ytr'].to(device)
        for _ in range(EPOCHS):
            net.train(); perm = torch.randperm(len(x), device=device)
            for i in range(0, len(x), BATCH):
                z = net(x[perm[i:i+BATCH]]); loss = lf(z, y[perm[i:i+BATCH]])
                opt.zero_grad(); loss.backward(); opt.step()
        net.eval()
        with torch.no_grad(): metric = float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
        return metric
    except RuntimeError:
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        return train_baseline_cpu(ds, lr, wd, seed)


def train_baseline_cpu(ds, lr, wd, seed):
    seed_all(seed); net=make_model('rnn_small', ds['input_shape'],1); opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd); lf=loss_fn(ds)
    for _ in range(EPOCHS):
        p=torch.randperm(len(ds['xtr']))
        for i in range(0,len(p),BATCH):
            ix=p[i:i+BATCH]; loss=lf(net(ds['xtr'][ix]),ds['ytr'][ix]); opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad(): return float(((net(ds['xte'])-ds['yte'])**2).mean())


def train_idea(ds, lr, wd, q, seed, trace=False):
    seed_all(seed); device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net=make_model('rnn_small',ds['input_shape'],1).to(device); opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd); lf=loss_fn(ds)
        x,y=ds['xtr'].to(device),ds['ytr'].to(device); best_vec=None; best_loss=float('inf'); records=[]
        for ep in range(EPOCHS):
            net.train(); perm=torch.randperm(len(x),device=device)
            for i in range(0,len(x),BATCH):
                ix=perm[i:i+BATCH]; loss=lf(net(x[ix]),y[ix]); opt.zero_grad(); loss.backward(); opt.step()
                # The barrier estimate is made at epoch boundaries below.
            with torch.no_grad(): cur_loss=float(lf(net(x[:256]),y[:256]))
            if cur_loss < best_loss: best_loss=cur_loss; best_vec=flat_params(net).clone()
            barrier=probe_barrier(net,best_vec,ds,lf,device)
            # High-loss flat points are treated as metastable; otherwise safe mode.
            net.train(); net.zero_grad(set_to_none=True); probe=lf(net(x[:128]),y[:128]); probe.backward()
            gnorm=float(torch.sqrt(sum((p.grad.detach()**2).sum() for p in net.parameters() if p.grad is not None)))
            trapped=(cur_loss-best_loss > 0.002 and gnorm < 2.0)
            qeff=2.0 if trapped else q
            eps=float(np.clip(max(barrier,1e-5)/qeff,1e-6,0.01))
            with torch.no_grad():
                for p in net.parameters(): p.add_(torch.randn_like(p)*math.sqrt(eps))
            records.append((eps, float(barrier), int(trapped), cur_loss))
        net.eval()
        with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
        return (metric,records) if trace else metric
    except RuntimeError:
        if torch.cuda.is_available(): torch.cuda.empty_cache()
        # CPU retry is deliberately identical intervention, with a temporary CUDA disable.
        old=torch.cuda.is_available
        return train_idea_cpu(ds,lr,wd,q,seed,trace)


def train_idea_cpu(ds,lr,wd,q,seed,trace=False):
    seed_all(seed); net=make_model('rnn_small',ds['input_shape'],1); opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd); lf=loss_fn(ds); x,y=ds['xtr'],ds['ytr']; best=None; bl=1e9; rec=[]
    for ep in range(EPOCHS):
        p=torch.randperm(len(x))
        for i in range(0,len(x),BATCH):
            ix=p[i:i+BATCH]; z=lf(net(x[ix]),y[ix]); opt.zero_grad();z.backward();opt.step()
        with torch.no_grad(): cl=float(lf(net(x[:256]),y[:256]))
        if cl<bl: bl=cl;best=flat_params(net).clone()
        bar=probe_barrier(net,best,ds,lf,'cpu'); net.zero_grad();lf(net(x[:128]),y[:128]).backward(); gn=float(torch.sqrt(sum((p.grad**2).sum() for p in net.parameters() if p.grad is not None))); tr=cl-bl>.002 and gn<2; eps=float(np.clip(max(bar,1e-5)/(2 if tr else q),1e-6,.01))
        with torch.no_grad():
            for p in net.parameters():p.add_(torch.randn_like(p)*math.sqrt(eps))
        rec.append((eps,bar,int(tr),cl))
    with torch.no_grad(): m=float(((net(ds['xte'])-ds['yte'])**2).mean())
    return (m,rec) if trace else m


def run_cfg(cfg, seeds=SEEDS, idea=False):
    vals=[]
    for s in seeds:
        ds=get_dataset('dynamics',s,n_train=NTR,n_test=NTE)
        vals.append(train_idea(ds,cfg['lr'],cfg['weight_decay'],cfg.get('q_safe',8),s) if idea else train_baseline(ds,cfg['lr'],cfg['weight_decay'],s))
    return {'per_seed':[float(v) for v in vals], 'mean':float(np.mean(vals))}


def main():
    # Baseline central knobs and all idea step sizes share the same union.
    grid=[{'lr':lr,'weight_decay':wd} for lr in LRS for wd in WDS]
    base=sweep_baseline(lambda c: (lambda seed: train_baseline(get_dataset('dynamics', seed, n_train=NTR, n_test=NTE), c['lr'], c['weight_decay'], seed)), grid)
    best=base['best_cfg']
    nearby=list(LRS)
    idea_cfgs=[{'lr':lr,'weight_decay':best['weight_decay'],'q_safe':q} for lr in nearby for q in QS]
    idea_runs=[(c,run_cfg(c,SEEDS,True)) for c in idea_cfgs]
    ic,best_idea=min(idea_runs,key=lambda z:z[1]['mean'])
    # NN-scale signature: measured barrier/noise settings and observed trapped events.
    ds=get_dataset('dynamics',0,n_train=NTR,n_test=NTE); _,trace=train_idea(ds,ic['lr'],ic['weight_decay'],ic['q_safe'],0,True)
    usable=[r for r in trace if r[0]>0 and r[1]>0]; slope=float(np.polyfit([1/r[0] for r in usable],[math.log(max(r[2],1e-3)) for r in usable],1)[0]) if len(usable)>=3 else float('nan'); pred=-float(np.median([r[1] for r in usable])) if usable else float('nan')
    sig={'predicted_slope':pred,'observed_slope':slope,'n_checkpoints':len(usable),'confirmed':bool(np.isfinite(slope) and np.isfinite(pred) and abs(slope-pred)<=.2*max(abs(pred),1e-6))}
    rep=make_report('dynamics','rnn_small',base,best_idea,{'barrier_noise_calibration':sig,'q_safe':ic['q_safe'],'custom_track':None})
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))

if __name__=='__main__': main()