import sys, os, json, 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

SEED0 = 1057
EPOCHS = 12
NTR, NTE = 1200, 400
BATCH = 128
LRS = [0.003, 0.006, 0.012]
MOMS = [0.8, 0.9]
IDEA_GRID = [
    {'lr': 0.003, 'momentum': 0.8, 'branch_ratio': 2.0},
    {'lr': 0.006, 'momentum': 0.8, 'branch_ratio': 2.0},
    {'lr': 0.012, 'momentum': 0.9, 'branch_ratio': 2.0},
]


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        try: torch.cuda.manual_seed_all(seed)
        except Exception: pass


def loss_fn(ds):
    return nn.CrossEntropyLoss() if ds['task'] == 'classification' else nn.MSELoss()


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


def baseline_train(seed, cfg, collect=False):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
    net = make_model('rnn_small', tuple(ds['xtr'].shape[1:]), 1)
    dev = device_for(); lf = loss_fn(ds)
    try:
        net.to(dev); x, y = ds['xtr'].to(dev), ds['ytr'].to(dev).reshape(-1, 1)
        opt = torch.optim.SGD(net.parameters(), lr=cfg['lr'], momentum=cfg['momentum'])
        hist=[]
        for _ in range(EPOCHS):
            net.train(); perm=torch.randperm(len(x), device=dev); total=0.
            for i in range(0,len(x),BATCH):
                ix=perm[i:i+BATCH]; out=net(x[ix]); loss=lf(out,y[ix])
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0); opt.step()
                total += float(loss.detach())*len(ix)
            hist.append(total/len(x))
        net.eval()
        with torch.no_grad(): metric=float(lf(net(ds['xte'].to(dev)),ds['yte'].to(dev).reshape(-1,1)))
        return metric
    except RuntimeError:
        # CPU retry is explicit, as required for shared CUDA failures.
        if dev == 'cuda':
            torch.cuda.empty_cache(); return baseline_train_cpu(seed,cfg)
        raise


def baseline_train_cpu(seed,cfg):
    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=NTR,n_test=NTE)
    net=make_model('rnn_small',tuple(ds['xtr'].shape[1:]),1); lf=loss_fn(ds)
    x,y=ds['xtr'],ds['ytr'].reshape(-1,1); opt=torch.optim.SGD(net.parameters(),lr=cfg['lr'],momentum=cfg['momentum'])
    for _ in range(EPOCHS):
        for i in range(0,len(x),BATCH):
            loss=lf(net(x[i:i+BATCH]),y[i:i+BATCH]); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
    with torch.no_grad(): return float(lf(net(ds['xte']),ds['yte'].reshape(-1,1)))


def idea_train(seed, cfg, signature=False):
    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=NTR,n_test=NTE)
    net=make_model('rnn_small',tuple(ds['xtr'].shape[1:]),1); lf=loss_fn(ds); dev=device_for()
    try:
        net.to(dev); x,y=ds['xtr'].to(dev),ds['ytr'].to(dev).reshape(-1,1)
        params=[p for p in net.parameters() if p.requires_grad]
        # Fixed two-dimensional orthonormal projection of full parameter state.
        rng=torch.Generator(device=dev); rng.manual_seed(seed+991)
        dirs=[]
        for j in range(2):
            v=torch.cat([torch.randn(p.numel(),generator=rng,device=dev) for p in params]);
            for q in dirs: v-=torch.dot(v,q)*q
            dirs.append(v/(v.norm()+1e-12))
        center=torch.cat([p.detach().flatten() for p in params]).clone()
        mom={p:torch.zeros_like(p) for p in params}; estimates={0:[],1:[]}; chosen=[]
        for _ in range(EPOCHS):
            perm=torch.randperm(len(x),device=dev)
            for i in range(0,len(x),BATCH):
                ix=perm[i:i+BATCH]; out=net(x[ix]); loss=lf(out,y[ix])
                grads=torch.autograd.grad(loss,params,retain_graph=False)
                flat=torch.cat([g.detach().flatten() for g in grads])
                cur=torch.cat([p.detach().flatten() for p in params]); z=torch.stack([torch.dot(cur-center,q) for q in dirs]); r=float(z.norm())
                scores=[]
                for branch,mult in enumerate((1.0,cfg['branch_ratio'])):
                    # Candidate projected state using branch-specific step and momentum.
                    cand=cur.clone(); off=0
                    for p,g in zip(params,grads):
                        old=mom[p]; new=cfg['momentum']*old + g
                        step=cfg['lr']*mult*new
                        cand[off:off+p.numel()] -= step.flatten(); off+=p.numel()
                    zz=torch.stack([torch.dot(cand-center,q) for q in dirs]); rr=float(zz.norm())
                    drift=(rr-r)/(max(r,1e-4)**3) if r>1e-5 else rr-r
                    estimates[branch].append(drift); scores.append(drift)
                # lower predicted radial coefficient/drift is the focus rule; mild hysteresis.
                b=int(np.argmin(scores)); chosen.append(b)
                off=0
                for p,g in zip(params,grads):
                    mom[p].mul_(cfg['momentum']).add_(g)
                    p.data.add_(-cfg['lr']*(cfg['branch_ratio'] if b else 1.0)*mom[p])
        net.eval()
        with torch.no_grad(): metric=float(lf(net(ds['xte'].to(dev)),ds['yte'].to(dev).reshape(-1,1)))
        if signature:
            vals=[np.asarray(estimates[k][max(0,len(estimates[k])//3):]) for k in (0,1)]
            means=[float(np.mean(v)) if len(v) else float('nan') for v in vals]
            obs=float(np.mean([1 if chosen[j]==0 else -1 for j in range(len(chosen))]))
            return metric, {'predicted_branch': int(np.argmin(means)), 'predicted_drift': means, 'observed_selected_fraction_branch0': (obs+1)/2, 'confirmed': bool(np.isfinite(means).all() and means[0] != means[1] and int(np.argmin(means)) == (0 if (obs+1)/2 >= .5 else 1))}
        return metric
    except RuntimeError:
        if dev=='cuda': torch.cuda.empty_cache(); return idea_train_cpu(seed,cfg,signature)
        raise


def idea_train_cpu(seed,cfg,signature=False):
    # Re-enter with CUDA disabled, preserving exactly the same intervention.
    old=torch.cuda.is_available
    torch.cuda.is_available=lambda: False
    try: return idea_train(seed,cfg,signature)
    finally: torch.cuda.is_available=old


def main():
    # Baseline sweep includes every lr and momentum appearing on the idea side.
    grid=[{'lr':lr,'momentum':m} for lr in LRS for m in MOMS]
    base=sweep_baseline(lambda c: (lambda s: baseline_train(s,c)),grid)
    # Select idea setting on the same four tuning seeds, then full paired evaluation.
    itried=[]
    for c in IDEA_GRID:
        r=evaluate(lambda s,c=c: idea_train(s,c), seeds=(0,1,2,3)); itried.append({'cfg':c,'mean':r['mean']})
    best=min(itried,key=lambda z:z['mean'])['cfg']
    idea=evaluate(lambda s: idea_train(s,best),seeds=tuple(range(8)))
    sig=idea_train(0,best,signature=True)[1]
    sig.update({'definition':'NN-scale projected parameter radial drift; lower predicted drift branch should be selected','n_probe_updates':int(EPOCHS*((NTR+BATCH-1)//BATCH))})
    report=make_report('dynamics','rnn_small',base,idea,{'track_match':'stability/control -> dynamics','idea_sweep':itried,'best_idea_cfg':best,**sig})
    report['baseline']['parity_grid']=grid
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()