import sys, json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from scipy.linalg import expm, eigvals
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, sweep_baseline, evaluate, make_report

OUT = Path('bench_report.json')
SEEDS = tuple(range(8))
# Union parity: baseline and idea both evaluated at every lr in this grid.
GRID = [{'lr': 0.001}, {'lr': 0.003}, {'lr': 0.006}]
EPOCHS = 18
BATCH = 128


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


def math_check():
    A = np.array([[-2., 5.], [0., -.5]])
    d = 1.25
    T = .4
    g0 = np.log(max(abs(eigvals(expm(A*T))))) / T
    gd = np.log(max(abs(eigvals(expm((A-d*np.eye(2))*T))))) / T
    return {'predicted_damping_shift': d, 'observed_shift': float(g0-gd),
            'abs_error': float(abs((g0-gd)-d)), 'passed': bool(abs(g0-gd-d)<1e-10)}


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


def train(seed, lr, periodic=False, amp=0.15, period=8, return_model=False):
    set_seed(seed)
    d = get_dataset('dynamics', seed, n_train=400, n_test=400)
    dev = device()
    try:
        model = make_model('rnn_small', d['input_shape'], d['out_dim']).to(dev)
        opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
        lossfn = nn.MSELoss()
        x, y = d['xtr'].to(dev), d['ytr'].to(dev)
        model.train()
        step = 0
        for ep in range(EPOCHS):
            # fixed order per seed; identical minibatches between systems
            perm = torch.arange(len(x), device=dev)
            for start in range(0, len(x), BATCH):
                ix = perm[start:start+BATCH]
                phase = 2*math.pi*(step % period)/period
                factor = 1.0 + amp*math.sin(phase) if periodic else 1.0
                for pg in opt.param_groups: pg['lr'] = lr*factor
                opt.zero_grad(set_to_none=True)
                pred = model(x[ix]); loss = lossfn(pred, y[ix])
                loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
                opt.step(); step += 1
        model.eval()
        with torch.no_grad():
            pred = model(d['xte'].to(dev)); metric = lossfn(pred, d['yte'].to(dev)).item()
        if return_model: return metric, model, d, dev
        return metric
    except RuntimeError:
        # Explicit CPU fallback for constrained/shared CUDA environments.
        if dev.type == 'cuda':
            torch.cuda.empty_cache()
            return train_cpu(seed, lr, periodic, amp, period, return_model)
        raise


def train_cpu(seed, lr, periodic=False, amp=.15, period=8, return_model=False):
    old = torch.cuda.is_available
    # same implementation forced onto CPU, avoiding recursive GPU selection
    set_seed(seed); d=get_dataset('dynamics',seed,n_train=400,n_test=400)
    model=make_model('rnn_small',d['input_shape'],d['out_dim'])
    opt=torch.optim.AdamW(model.parameters(),lr=lr,weight_decay=1e-4); lossfn=nn.MSELoss(); x,y=d['xtr'],d['ytr']; step=0
    model.train()
    for ep in range(EPOCHS):
        for start in range(0,len(x),BATCH):
            phase=2*math.pi*(step%period)/period; f=1+amp*math.sin(phase) if periodic else 1
            opt.param_groups[0]['lr']=lr*f; opt.zero_grad(); loss=lossfn(model(x[start:start+BATCH]),y[start:start+BATCH]); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step(); step+=1
    model.eval()
    with torch.no_grad(): metric=lossfn(model(d['xte']),d['yte']).item()
    return (metric,model,d,torch.device('cpu')) if return_model else metric


def signature():
    # Measured on trained models: local output Jacobian products over one schedule.
    vals=[]
    for s in (0,1):
        mb, base, d, dev = train(s,.003,False,return_model=True)
        mi, idea, _, _ = train(s,.003,True,return_model=True)
        x=d['xte'][:1].to(dev).requires_grad_(True)
        def local_jac(m):
            z=x.detach().clone().requires_grad_(True); out=m(z)[0,0]
            g=torch.autograd.grad(out,z,create_graph=False)[0].detach().flatten().cpu().numpy()
            # Scalar observable Jacobian norm is a conservative NN-scale tangent proxy.
            return float(np.linalg.norm(g))
        vals.append({'seed':s,'baseline_jacobian_norm':local_jac(base),'idea_jacobian_norm':local_jac(idea),'baseline_mse':mb,'idea_mse':mi})
    b=np.mean([v['baseline_jacobian_norm'] for v in vals]); i=np.mean([v['idea_jacobian_norm'] for v in vals])
    return {'quantity':'trained-model local input Jacobian norm','predicted':'periodic modulation reduces perturbation growth','baseline_mean':float(b),'idea_mean':float(i),'relative_change':float((i-b)/max(b,1e-12)),'confirmed':bool(i<b)}


def main():
    sanity=math_check()
    base=sweep_baseline(lambda cfg: lambda seed: train(seed,cfg['lr'],False), GRID)
    # Idea at all union lrs; select best on the same four-seed selection protocol.
    tried=[]
    for cfg in GRID:
        r=evaluate(lambda seed,cfg=cfg: train(seed,cfg['lr'],True), seeds=(0,1,2,3))
        tried.append({'cfg':cfg,'mean':r['mean']})
    best_cfg=min(tried,key=lambda z:z['mean'])['cfg']
    idea=evaluate(lambda seed: train(seed,best_cfg['lr'],True), seeds=SEEDS)
    report=make_report('dynamics','rnn_small',base,idea,{'math_sanity':sanity,'trained_jacobian':signature(),'idea_cfg':best_cfg,'idea_sweep':tried})
    report['selection_note']='Baseline and idea share rnn_small, data, epochs, batch, AdamW weight decay, and union learning-rate grid; only periodic LR modulation differs.'
    OUT.write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__': main()