import os, sys, 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, train_model, sweep_baseline, make_report
from bench.protocol import evaluate

# Characteristic-invariant BT monitor on a measured 2D local recurrent chart.
# The architecture remains exactly bench rnn_small; h[0:2] are perturbed and
# the other recurrent coordinates are held at zero for a cheap local monitor.
def monitor(model):
    rnn = model.rnn
    W = rnn.weight_hh_l0
    b = rnn.bias_ih_l0 + rnn.bias_hh_l0
    z = torch.zeros(2, device=W.device, dtype=W.dtype, requires_grad=True)
    h = torch.cat([z, torch.zeros(W.shape[1]-2, device=W.device, dtype=W.dtype)])
    # PyTorch GRU gate order is reset, update, new.
    gates = torch.mv(W, h) + b
    r, u, npre = gates.chunk(3, 0)
    rr, uu = torch.sigmoid(r), torch.sigmoid(u)
    nnv = torch.tanh(npre + rr * 0)  # zero external input; recurrent reset is in Wn below
    # Recompute candidate with reset applied to recurrent hidden contribution.
    whr, whu, whn = W.chunk(3, 0)
    bir, biu, bin_ = b.chunk(3, 0)
    rr = torch.sigmoid(torch.mv(whr, h) + bir)
    uu = torch.sigmoid(torch.mv(whu, h) + biu)
    nnv = torch.tanh(torch.mv(whn, rr*h) + bin_)
    hp = (1-uu)*nnv + uu*h
    out = hp[:2]
    rows=[]
    for i in range(2):
        rows.append(torch.autograd.grad(out[i], z, create_graph=True, retain_graph=True)[0])
    A = torch.stack(rows)
    # Invariants of the continuous vector field f=h'-h: J=A-I.
    J = A - torch.eye(2, device=A.device, dtype=A.dtype)
    e1 = torch.trace(J); e2 = torch.linalg.det(J)
    delta, tau = e2, e1
    q = torch.linalg.svd(J.detach()).Vh[-1]
    gd = torch.autograd.grad(delta, z, create_graph=True, retain_graph=True)[0]
    gt = torch.autograd.grad(tau, z, create_graph=True, retain_graph=True)[0]
    a = -0.5 * torch.dot(gd, q)
    bb = torch.dot(gt, q)
    return delta, tau, a, bb, J.detach()

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)

def train_idea(seed, cfg, return_model=False):
    seed_all(seed)
    ds=get_dataset('dynamics', seed, n_train=400, n_test=400)
    model=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        model.to(dev); x,y=ds['xtr'].to(dev),ds['ytr'].to(dev)
        opt=torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        lossf=nn.MSELoss()
        for ep in range(cfg['epochs']):
            model.train(); perm=torch.randperm(len(x),device=dev)
            for ix in perm.split(128):
                pred=model(x[ix]); loss=lossf(pred,y[ix])
                d,t,a,b,_=monitor(model)
                loss=loss + cfg['rho']*(d*d+t*t)
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
        model.eval()
        with torch.no_grad(): metric=float(lossf(model(ds['xte'].to(dev)),ds['yte'].to(dev)).cpu())
    except Exception:
        # Required robust fallback: restart fully on CPU after any CUDA failure.
        seed_all(seed); dev=torch.device('cpu'); model=make_model('rnn_small',ds['input_shape'],ds['out_dim']).to(dev)
        x,y=ds['xtr'],ds['ytr']; opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
        for ep in range(cfg['epochs']):
            for ix in torch.randperm(len(x)).split(128):
                loss=lossf(model(x[ix]),y[ix]); d,t,a,b,_=monitor(model); loss=loss+cfg['rho']*(d*d+t*t)
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): metric=float(lossf(model(ds['xte']),ds['yte']))
    if return_model: return metric, model, ds
    return metric

def train_base(seed,cfg, return_model=False):
    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=400)
    model=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
    raw=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,weight_decay=cfg['weight_decay'],log=lambda *_:None)
    metric = float(raw[1] if isinstance(raw,(tuple,list)) else raw.get('test', raw.get('metric')) if isinstance(raw,dict) else raw)
    return (metric,model,ds) if return_model else metric

def signature(cfg):
    vals=[]
    for s in range(8):
        m,model,ds=train_idea(s,cfg,True)
        d,t,a,b,J=monitor(model)
        eig=torch.linalg.eigvals(J).cpu().numpy()
        # observed recurrent persistence is spectral radius of the trained local map.
        vals.append((float(abs(d)),float(abs(t)),float(np.max(np.abs(eig+1))),float(np.max(eig.real))))
    v=np.asarray(vals)
    return {'prediction':'BT monitor drives delta,tau toward zero and the trained local dynamics toward a near-zero continuous eigenvalue','trained_mean_abs_delta':float(v[:,0].mean()),'trained_mean_abs_tau':float(v[:,1].mean()),'observed_mean_discrete_radius':float(v[:,2].mean()),'observed_mean_continuous_max_real':float(v[:,3].mean()),'confirmed':bool(v[:,0].mean()<0.08 and v[:,1].mean()<0.08 and v[:,2].mean()>0.85)}

def main():
    base_grid=[{'lr':x,'weight_decay':w,'epochs':8} for x in (0.001,0.003,0.01) for w in (0.0,1e-4)]
    # Same union of decisive Adam knobs is used for baseline and idea.
    base=sweep_baseline(lambda c: (lambda s: train_base(s,c)),base_grid)
    idea_grid=[dict(base['best_cfg'],rho=r) for r in (0.001,0.01,0.05)]
    # idea uses best baseline lr/wd and two a-priori regularizer strengths
    idea_runs=[(c,evaluate(lambda s,c=c: train_idea(s,c))) for c in idea_grid]
    best_i,best_r=min(idea_runs,key=lambda z:z[1]['mean'])
    rep=make_report('dynamics','rnn_small',base,best_r,{'idea_config':best_i,'idea_sweep':[{'cfg':c,'mean':r['mean']} for c,r in idea_runs],**signature(best_i)})
    rep['custom_track']=None
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()