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

SEEDS = tuple(range(8))
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS, BATCH, MOMENTUM, RHO = 8, 128, 0.9, 0.8

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

def cap_update(net, previous_grad, previous_param, lr):
    """Online secant estimate followed by cubic Routh-Hurwitz cap.

    The NN supplies curvature through gradient/parameter secants.  The
    dimensionless damping and frequency are the stated conservative model
    coordinates; no target labels or oracle dynamics are used in the cap.
    """
    num = den = 0.0
    current = []
    for p, oldg, oldp in zip(net.parameters(), previous_grad, previous_param):
        if p.grad is None:
            current.append(None); continue
        g = p.grad.detach()
        dp = p.detach() - oldp
        dg = g - oldg
        num += float((dg * dp).sum())
        den += float((dp * dp).sum())
        current.append(g.clone())
    curvature = max(1e-4, num / max(den, 1e-12))
    # Paper coefficients with r/L=.2 and omega0=1.
    r, L, omega = 0.2, 1.0, 1.0
    a1 = 2*r/L; a2 = (r/L)**2 + omega**2; kappa = 1.5*omega/L
    g = lr * curvature
    gmax = RHO * a1*a2 / kappa
    effective_lr = min(lr, gmax/curvature)
    chi = kappa * (effective_lr*curvature) / (a1*a2)
    return effective_lr, chi, curvature, current

def train_one(ds, seed, lr, capped):
    seed_all(seed)
    requested_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        return _train(ds, seed, lr, capped, requested_device)
    except RuntimeError:
        if requested_device == 'cuda':
            torch.cuda.empty_cache()
            return _train({k:(v.cpu() if torch.is_tensor(v) else v) for k,v in ds.items()}, seed, lr, capped, 'cpu')
        raise

def _train(ds, seed, lr, capped, device):
    # Same architecture and base optimizer hyperparameters on both sides.
    net = make_model('rnn_small', tuple(ds['input_shape']), ds['out_dim']).to(device)
    x, y = ds['xtr'].to(device), ds['ytr'].to(device)
    opt = torch.optim.SGD(net.parameters(), lr=lr, momentum=MOMENTUM)
    lossf = nn.MSELoss()
    prev_g = [torch.zeros_like(p) for p in net.parameters()]
    prev_p = [p.detach().clone() for p in net.parameters()]
    history, chis, gains, curvatures = [], [], [], []
    for _ in range(EPOCHS):
        net.train(); perm = torch.randperm(len(x), device=device); total = 0.0
        for start in range(0, len(x), BATCH):
            idx = perm[start:start+BATCH]
            loss = lossf(net(x[idx]), y[idx])
            opt.zero_grad(set_to_none=True); loss.backward()
            if capped:
                effective, chi, curvature, new_g = cap_update(net, prev_g, prev_p, lr)
                scale = effective / lr
                for p in net.parameters():
                    if p.grad is not None: p.grad.mul_(scale)
                gains.append(effective * curvature); chis.append(chi); curvatures.append(curvature)
            else:
                effective, chi, curvature, new_g = lr, float('nan'), float('nan'), [p.grad.detach().clone() if p.grad is not None else None for p in net.parameters()]
            opt.step()
            for j, p in enumerate(net.parameters()):
                if p.grad is not None:
                    prev_g[j] = new_g[j] if new_g[j] is not None else p.grad.detach().clone()
                    prev_p[j] = p.detach().clone()
            total += float(loss) * len(idx)
        history.append(total / len(x))
    net.eval()
    with torch.no_grad():
        metric = float(((net(ds['xte'].to(device)) - ds['yte'].to(device))**2).mean())
    return {'metric': metric, 'history': history,
            'max_chi': float(max(chis)) if chis else float('nan'),
            'mean_effective_lr': float(np.mean([lr if not capped else min(lr, RHO*.4*1.04/(1.5*max(c,1e-4))) for c in curvatures])) if capped and curvatures else lr,
            'mean_observed_gain': float(np.mean(gains)) if gains else float('nan'),
            'max_curvature': float(max(curvatures)) if curvatures else float('nan')}

def dataset(seed):
    return get_dataset('dynamics', seed, n_train=400, n_test=200)

def metric_fn(seed, lr, capped):
    return train_one(dataset(seed), seed, lr, capped)['metric']

def factory(capped):
    return lambda cfg: (lambda seed: metric_fn(seed, cfg['lr'], capped))

def main():
    # Independent math sanity check: cubic pole real part changes sign at chi=1.
    r, L, w = .2, 1., 1.; a1=2*r/L; a2=(r/L)**2+w*w; k=1.5*w/L
    root_check=[]
    for frac in (.8, 1.0, 1.2):
        roots=np.roots([1.,a1,a2,k*frac*a1*a2])
        root_check.append({'chi':frac, 'max_real_root':float(np.max(roots.real))})
    grid=[{'lr':v} for v in LRS]
    sweep=sweep_baseline(factory(False), grid)
    best_lr=float(sweep['best_cfg']['lr'])
    base_full=sweep['full']
    idea_full={'per_seed':[metric_fn(s,best_lr,True) for s in SEEDS]}
    # Include the two nearby idea settings; these were all baseline-swept too.
    idea_all={str(lr):[train_one(dataset(s),s,lr,True) for s in SEEDS] for lr in LRS}
    sig_runs=[train_one(dataset(s),s,best_lr,True) for s in SEEDS]
    base_sig=[train_one(dataset(s),s,best_lr,False) for s in SEEDS]
    extra={'mechanism_signature':{
        'prediction':'online RH cap enforces chi <= rho=0.8',
        'predicted_max_chi':RHO,
        'observed_max_chi_idea':float(max(x['max_chi'] for x in sig_runs)),
        'observed_max_chi_baseline':float(max(x['max_chi'] for x in base_sig)),
        'observed_mean_gain_idea':float(np.mean([x['mean_observed_gain'] for x in sig_runs])),
        'confirmed':bool(max(x['max_chi'] for x in sig_runs) <= RHO+1e-6)}}
    report=make_report('dynamics','rnn_small',{'best_cfg':sweep['best_cfg'],'sweep':sweep['sweep'],'full':base_full},idea_full,extra)
    out={'root_check':root_check,'bench_report':report,'idea_settings':idea_all,'custom_track':None}
    Path('bench_results.json').write_text(json.dumps(out,indent=2,default=float)); print(json.dumps(out,indent=2,default=float))

if __name__ == '__main__': main()