import sys, json, random, math
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, train_model, evaluate, sweep_baseline, make_report

SEEDS = list(range(8))
# Union is shared: every idea lr is also a baseline sweep point.
LRS = [1e-3, 2e-3, 3e-3]
EPOCHS = 4
BATCH = 128
KAPPA = [0.05, 0.2, 0.8]
LAMBDA_STAR = -0.02
GAMMA_STAR = 0.05


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 baseline_one(cfg, seed):
    seed_all(seed)
    d = get_dataset('dynamics', seed=seed, n_train=400, n_test=400)
    m = make_model('rnn_small', d['input_shape'], d['out_dim'])
    _, metric, hist = train_model(m, d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None)
    return metric


def _gru_step(rnn, xt, h):
    gi = xt @ rnn.weight_ih_l0.T + rnn.bias_ih_l0
    gh = h @ rnn.weight_hh_l0.T + rnn.bias_hh_l0
    iz, ir, inn = gi.chunk(3, 1); hz, hr, hnn = gh.chunk(3, 1)
    z = torch.sigmoid(iz + hz); r = torch.sigmoid(ir + hr)
    n = torch.tanh(inn + r * hnn)
    return (1 - z) * n + z * h


def jacobian_penalty(net, xb, need_signature=False):
    """Differentiable two-vector QR estimator using JVPs, without materializing J."""
    rnn = net.rnn
    h = torch.zeros(xb.shape[0], rnn.hidden_size, device=xb.device)
    q = torch.zeros(xb.shape[0], rnn.hidden_size, 2, device=xb.device)
    q[:, 0, 0] = 1.; q[:, 1, 1] = 1.
    logs = []
    try:
        from torch.func import jvp
    except ImportError:
        raise RuntimeError('torch.func.jvp unavailable')
    for xt in xb.view(xb.shape[0], -1, 3).unbind(1):
        # One JVP per tangent column; batching remains identical to task training.
        f = lambda hh: _gru_step(rnn, xt, hh)
        ucols = []
        for col in range(2):
            _, uj = jvp(f, (h,), (q[:, :, col],))
            ucols.append(uj)
        u = torch.stack(ucols, dim=2)
        q, R = torch.linalg.qr(u, mode='reduced')
        logs.append(torch.log(torch.diagonal(R, dim1=1, dim2=2).abs().clamp_min(1e-7)))
        h = f(h)
    rates = torch.stack(logs, 1).mean((0, 1))
    lam1, lam2 = rates[0], rates[1]
    gap = lam1 - lam2
    penalty = torch.relu(torch.as_tensor(GAMMA_STAR, device=xb.device) - gap).square() + (lam1 - LAMBDA_STAR).square()
    if need_signature:
        return penalty, float(lam1.detach()), float(lam2.detach()), float(gap.detach())
    return penalty

def idea_train(seed, lr, kappa, return_net=False):
    seed_all(seed)
    d = get_dataset('dynamics', seed=seed, n_train=400, n_test=400)
    # Same rnn_small architecture and data as baseline; only objective differs.
    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net.to(device); opt = torch.optim.Adam(net.parameters(), lr=lr)
        x, y = d['xtr'].to(device), d['ytr'].to(device)
        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]; pred = net(x[idx])
                task = ((pred-y[idx])**2).mean()
                reg = jacobian_penalty(net, x[idx[:16]])
                loss = task + kappa * reg
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0); opt.step()
        net.eval()
        with torch.no_grad(): metric = float(((net(d['xte'].to(device))-d['yte'].to(device))**2).mean())
        return (metric, net, d) if return_net else metric
    except RuntimeError:
        # Small CPU fallback; retry with a fresh model and no CUDA.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        torch.set_default_device('cpu')
        return idea_train_cpu(seed, lr, kappa, return_net)


def idea_train_cpu(seed, lr, kappa, return_net=False):
    seed_all(seed); d=get_dataset('dynamics', seed, 400, 400); net=make_model('rnn_small', d['input_shape'], 1)
    opt=torch.optim.Adam(net.parameters(),lr=lr); x,y=d['xtr'],d['ytr']
    for _ in range(EPOCHS):
        for i in range(0,len(x),BATCH):
            pred=net(x[i:i+BATCH]); task=((pred-y[i:i+BATCH])**2).mean(); loss=task+kappa*jacobian_penalty(net,x[i:i+BATCH]); opt.zero_grad(); loss.backward(); opt.step()
    metric=float(((net(d['xte'])-d['yte'])**2).mean()); return (metric,net,d) if return_net else metric


def contraction_signature(net, d, n=16, K=8):
    """Behavioral NN-scale signature: two trained-model JVP directions and QR."""
    dev=next(net.parameters()).device; r=net.rnn; x=d['xte'][:n].to(dev)
    h=torch.zeros(n,r.hidden_size,device=dev); q1=nn.functional.normalize(torch.randn_like(h),dim=1); q2=nn.functional.normalize(torch.randn_like(h),dim=1)
    angles=[]; gaps=[]
    from torch.func import jvp
    for xt in x.view(n,-1,3).unbind(1):
        f=lambda hh: _gru_step(r,xt,hh)
        _,u1=jvp(f,(h,),(q1,)); _,u2=jvp(f,(h,),(q2,))
        q1=nn.functional.normalize(u1,dim=1); q2=nn.functional.normalize(u2,dim=1)
        angles.append(torch.acos((q1*q2).sum(1).abs().clamp(0,1)).mean().item())
        # singular gap is measured on the two-vector QR factors of the trained dynamics
        U=torch.stack([u1,u2],2); _,R=torch.linalg.qr(U,mode='reduced')
        gaps.append(torch.log(torch.diagonal(R,dim1=1,dim2=2).abs().clamp_min(1e-7)).diff(dim=1).neg().mean().item())
        h=f(h).detach()
    obs=float(np.polyfit(np.arange(len(angles)),np.log(np.maximum(angles,1e-8)),1)[0]); pred=-float(np.mean(gaps))
    return {'observed_log_angle_slope':obs,'predicted_minus_gap':pred,'observed_gap':float(np.mean(gaps)),'relative_error':abs(obs-pred)/max(abs(pred),1e-8),'confirmed':bool(abs(obs-pred)<=0.30*max(abs(pred),1e-8))}


def main():
    baseline_grid=[{'lr':lr} for lr in LRS]
    base=sweep_baseline(lambda cfg: lambda seed: baseline_one(cfg,seed), baseline_grid, seeds=SEEDS)
    # Idea sweep: baseline-best lr plus two nearby rates, and three predeclared kappa values.
    # Every idea learning rate is in baseline_grid (search-space parity).
    idea_grid=[]
    for lr in LRS:
        for kap in KAPPA:
            vals=[idea_train(seed,lr,kap) for seed in SEEDS]
            idea_grid.append({'lr':lr,'kappa':kap,'per_seed':vals,'mean':float(np.mean(vals)),'std':float(np.std(vals,ddof=1))})
    best=min(idea_grid,key=lambda z:z['mean'])
    idea={'per_seed':best['per_seed'],'mean':best['mean'],'std':best['std'],'config':{'lr':best['lr'],'kappa':best['kappa'],'lambda_star':LAMBDA_STAR,'gamma_star':GAMMA_STAR,'epochs':EPOCHS},'sweep':[{k:v for k,v in z.items() if k!='per_seed'} for z in idea_grid]}
    m,d0=idea_train(SEEDS[0],best['lr'],best['kappa'],True)[1:]
    sig=contraction_signature(m,d0)
    rep=make_report('dynamics','rnn_small',base,idea,{'track_match':'dynamics stability/control','prediction':sig,'selection':{'kappa_candidates':KAPPA,'best_config':{'lr':best['lr'],'kappa':best['kappa']}}})
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()