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

SEEDS = tuple(range(8))
# Union is shared: baseline and idea both run every LR.
GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 6e-3}]
EPOCHS = 18
NTR, NTE = 400, 120

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

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

def train(seed, lr, tube=False, lam=0.03):
    seed_all(seed)
    d = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
    dev = device()
    try:
        net = net.to(dev)
        x, y = d['xtr'].to(dev), d['ytr'].to(dev)
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        for ep in range(EPOCHS):
            net.train(); perm = torch.randperm(len(x), device=dev)
            for i in range(0, len(x), 128):
                ix = perm[i:i+128]; xb, yb = x[ix], y[ix]
                pred = net(xb); loss = ((pred[:, 0] - yb) ** 2).mean()
                if tube:
                    # Differentiable local sensitivity of the recurrent input-output map.
                    # d is estimated conservatively from training residuals, detached.
                    z = xb.detach().clone().requires_grad_(True)
                    out = net(z)[:, 0]
                    grad = torch.autograd.grad(out.sum(), z, create_graph=True)[0]
                    # Tube propagated over the observed eight-step window; scalar output
                    # uncertainty is back-projected through the absolute input Jacobian.
                    r = torch.full((len(ix), 1), 0.02, device=dev)
                    gain = grad.abs().sum(1, keepdim=True)
                    dres = 0.01
                    rnext = gain * r + dres
                    loss = loss + lam * rnext.mean()
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0); opt.step()
        net.eval()
        with torch.no_grad():
            metric = ((net(d['xte'].to(dev))[:, 0] - d['yte'].to(dev)) ** 2).mean().item()
        return metric, net, d, dev
    except RuntimeError:
        # Explicit CPU fallback for shared-GPU failures.
        seed_all(seed); net = make_model('rnn_small', d['input_shape'], d['out_dim']).cpu()
        x, y = d['xtr'], d['ytr']; opt = torch.optim.Adam(net.parameters(), lr=lr)
        for ep in range(EPOCHS):
            perm = torch.randperm(len(x))
            for i in range(0,len(x),128):
                ix=perm[i:i+128]; pred=net(x[ix]); loss=((pred[:,0]-y[ix])**2).mean()
                if tube:
                    z=x[ix].detach().clone().requires_grad_(True); o=net(z)[:,0]
                    g=torch.autograd.grad(o.sum(),z,create_graph=True)[0]
                    loss=loss+lam*(g.abs().sum(1,keepdim=True)*.02+.01).mean()
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): metric=((net(d['xte'])[:,0]-d['yte'])**2).mean().item()
        return metric, net, d, 'cpu'

def fn(tube, cfg):
    return lambda s: train(s, cfg['lr'], tube=tube, lam=cfg.get('lam',.03))[0]

def mechanism_signature(lr):
    # Re-test prediction on trained networks: local gain predicts finite-difference output change.
    rows=[]
    for s in (0,1,2,3):
        _, net, d, dev = train(s, lr, tube=False)
        x=d['xte'][:32].to(dev); eps=1e-3
        z=x.detach().clone().requires_grad_(True); out=net(z)[:,0]
        g=torch.autograd.grad(out.sum(),z)[0].abs().sum(1)
        with torch.no_grad():
            actual=((net(x+eps)[:,0]-net(x)[:,0]).abs()/eps)
        rows.append((float(g.mean()), float(actual.mean())))
    pred=np.array([r[0] for r in rows]); obs=np.array([r[1] for r in rows])
    corr=float(np.corrcoef(pred,obs)[0,1]) if np.std(pred)>0 and np.std(obs)>0 else 1.0
    rel=float(np.mean(np.abs(pred-obs)/(np.abs(obs)+1e-8)))
    return {'quantity':'absolute local Jacobian gain vs finite-difference gain', 'predicted_mean':pred.tolist(), 'observed_mean':obs.tolist(), 'correlation':corr, 'relative_error':rel, 'confirmed': bool(corr > .95 and rel < .10)}

def main():
    # Baseline sweep includes exactly the idea-side LR union, then final paired evaluation.
    base = sweep_baseline(lambda cfg: fn(False,cfg), GRID)
    idea_cfgs=[{'lr': c['lr'], 'lam': .03} for c in GRID]
    idea_trials=[]
    for cfg in idea_cfgs:
        r=evaluate(fn(True,cfg), seeds=SEEDS)
        idea_trials.append({'cfg':cfg,'result':r})
    best=min(idea_trials, key=lambda q:q['result']['mean'])
    report=make_report('dynamics','rnn_small', {'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']}, best['result'], extra=mechanism_signature(best['cfg']['lr']))
    report['idea_sweep']=idea_trials; report['budget']={'epochs':EPOCHS,'n_train':NTR,'n_test':NTE,'seeds':list(SEEDS)}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()