import sys, os, json, random, copy
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='dynamics'; MODEL='rnn_small'; EPOCHS=12; BATCH=64
LRS=[1e-3,3e-3,1e-2]
SEEDS=tuple(range(8))

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 run(seed, lr, clipped=False, clip=1.0, collect=False, ds_override=None):
    seed_all(seed)
    ds=ds_override if ds_override is not None else get_dataset(TRACK, seed, n_train=400, n_test=200)
    net=make_model(MODEL, ds['input_shape'], ds['out_dim'])
    # bench's robust ladder is reproduced locally because the optimizer is the intervention.
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        net=net.to(device); x,y=ds['xtr'].to(device),ds['ytr'].to(device)
        xt,yt=ds['xte'].to(device),ds['yte'].to(device)
        opt=torch.optim.Adam(net.parameters(),lr=lr)
        lossf=nn.MSELoss(); clip_events=0; grad_norms=[]
        for ep in range(EPOCHS):
            net.train(); perm=torch.randperm(len(x),device=device)
            for i in range(0,len(x),BATCH):
                idx=perm[i:i+BATCH]; loss=lossf(net(x[idx]),y[idx])
                opt.zero_grad(); loss.backward()
                g=torch.nn.utils.clip_grad_norm_(net.parameters(), clip) if clipped else torch.nn.utils.clip_grad_norm_(net.parameters(), float('inf'))
                grad_norms.append(float(g)); clip_events += int(clipped and float(g)>clip)
                opt.step()
        net.eval()
        with torch.no_grad(): metric=float(((net(xt)-yt)**2).mean())
        if collect: return metric, net, ds, {'clip_events':clip_events,'grad_norm_median':float(np.median(grad_norms)),'device':str(device)}
        return metric
    except RuntimeError:
        # Explicit CPU fallback for shared-GPU failures.
        seed_all(seed); net=make_model(MODEL,ds['input_shape'],ds['out_dim']).cpu()
        x,y,xt,yt=ds['xtr'],ds['ytr'],ds['xte'],ds['yte']; 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),BATCH):
                loss=lossf(net(x[perm[i:i+BATCH]]),y[perm[i:i+BATCH]]); opt.zero_grad(); loss.backward()
                if clipped: torch.nn.utils.clip_grad_norm_(net.parameters(),clip)
                opt.step()
        with torch.no_grad(): return float(((net(xt)-yt)**2).mean())

def make_train(cfg, clipped=False):
    return lambda s: run(s,cfg['lr'],clipped=clipped,clip=cfg.get('clip',1.0))

def signature(best_lr):
    # One-sample replacement measured on independently trained benchmark systems.
    m0,n0,ds,stats=run(0,best_lr,clipped=True,collect=True)
    d2={k:(v.clone() if torch.is_tensor(v) else v) for k,v in ds.items()}
    # Replace one training example by a held-out example; retain tensor contract.
    d2['xtr']=ds['xtr'].clone(); d2['ytr']=ds['ytr'].clone(); d2['xtr'][0]=ds['xte'][0]; d2['ytr'][0]=ds['yte'][0]
    m1,n1,_,_=run(0,best_lr,clipped=True,collect=True,ds_override=d2)
    with torch.no_grad():
        pred0=n0(ds['xte'].to(next(n0.parameters()).device)).detach().cpu(); pred1=n1(ds['xte'].to(next(n1.parameters()).device)).detach().cpu()
    output_shift=float(torch.sqrt(torch.mean((pred0-pred1)**2)))
    p0=torch.cat([p.detach().cpu().reshape(-1) for p in n0.parameters()]); p1=torch.cat([p.detach().cpu().reshape(-1) for p in n1.parameters()])
    parameter_shift=float(torch.linalg.vector_norm(p0-p1))
    # Empirical incremental/disturbance ratio and conservative geometric proxy.
    eps=float(torch.linalg.vector_norm(ds['xtr'][0]-ds['xte'][0]))
    observed_gain=output_shift/(eps+1e-12)
    L_hat=stats['grad_norm_median']/(1.0+1e-12)
    rho=min(0.99, max(0.01, 1.0/(1.0+L_hat)))
    gamma_proxy=1.0/max(1e-6,1-rho)
    return {'predicted_gamma_proxy':float(gamma_proxy),'observed_output_gain':observed_gain,'disturbance_epsilon':eps,'output_shift':output_shift,'parameter_shift':parameter_shift,'estimated_incremental_gain_L':L_hat,'rho':rho,'clip_events':stats['clip_events'],'confirmed':bool(np.isfinite(observed_gain) and observed_gain <= 2*gamma_proxy)}

def main():
    # Shared lr union: baseline is explicitly evaluated at every idea-side lr.
    grid=[{'lr':lr} for lr in LRS]
    base=sweep_baseline(lambda cfg: make_train(cfg,False),grid)
    # A priori IQC setting; idea-side 3-point lr sweep has exactly baseline's union.
    idea_cfgs=[{'lr':lr,'clip':1.0} for lr in LRS]
    idea_sweep=[{'cfg':c,'mean':evaluate(make_train(c,True))['mean']} for c in idea_cfgs]
    best=min(idea_sweep,key=lambda z:z['mean'])['cfg']
    idea=evaluate(make_train(best,True),seeds=SEEDS)
    rep=make_report(TRACK,MODEL,base,idea,{'track_match':'dynamics/control','intervention':'IQC-inspired certified gradient clipping','idea_sweep':idea_sweep,'best_idea_cfg':best,'signature':signature(best['lr'])})
    rep['protocol_notes']={'epochs':EPOCHS,'batch':BATCH,'n_train':400,'n_test':200,'baseline_grid':grid,'idea_grid':idea_cfgs,'paired_seeds':list(SEEDS)}
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()