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

EPOCHS, BATCH = 12, 128
GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 6e-3}]


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 controlled_forward(net, x, stabilized=False):
    seq = x.reshape(x.shape[0], 8, 3)
    gru, head = net.rnn, net.head
    h = torch.zeros(1, x.shape[0], gru.hidden_size, device=x.device, dtype=x.dtype)
    for k in range(8):
        raw, _ = gru(seq[:, k:k+1, :], h)
        proposed = raw[:, 0, :]
        h = (h[:, 0, :] + (0.5 if stabilized else 1.0) *
             (proposed - h[:, 0, :])).unsqueeze(0)
    return head(h[0])


def train_idea(ds, lr, seed):
    seed_all(seed)
    net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        net.to(device)
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
        for _ in range(EPOCHS):
            net.train(); perm = torch.randperm(len(x), device=device)
            for i in range(0, len(x), BATCH):
                ix = perm[i:i+BATCH]
                loss = nn.functional.mse_loss(controlled_forward(net, x[ix], True), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
        net.eval()
        with torch.no_grad():
            m = float(((controlled_forward(net, ds['xte'].to(device), True) - ds['yte'].to(device)) ** 2).mean())
        return net, m
    except RuntimeError:
        if device.type != 'cuda': raise
        torch.cuda.empty_cache()
        net = make_model('rnn_small', ds['input_shape'], ds['out_dim']).cpu()
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        x, y = ds['xtr'], ds['ytr']
        for _ in range(EPOCHS):
            perm = torch.randperm(len(x))
            for i in range(0, len(x), BATCH):
                ix = perm[i:i+BATCH]
                loss = nn.functional.mse_loss(controlled_forward(net, x[ix], True), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            m = float(((controlled_forward(net, ds['xte'], True)-ds['yte'])**2).mean())
        return net, m


def baseline_factory(cfg):
    def run(seed):
        ds = get_dataset('dynamics', int(seed), 400, 200)
        seed_all(int(seed)); net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
        _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None)
        return metric
    return run


def idea_factory(cfg):
    return lambda seed: train_idea(get_dataset('dynamics', int(seed), 400, 200), cfg['lr'], int(seed))[1]


def train_one_baseline(ds, lr, seed):
    seed_all(seed); net=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
    net,m,_=train_model(net,ds,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
    return net,m


def trained_exponent(net, x, stabilized):
    """Exact JVP estimate on a trained GRU; called on CPU to avoid cuDNN JVP limits."""
    net = net.cpu().eval(); x = x.cpu(); gru = net.rnn
    seq = x[:1].reshape(1, 8, 3)
    q = torch.randn(gru.hidden_size); q /= q.norm(); total = 0.0
    h = torch.zeros(gru.hidden_size)
    for k in range(8):
        u = seq[:, k:k+1, :]
        def fmap(v):
            raw, _ = gru(u, v.reshape(1,1,-1)); r = raw[0,0]
            return v + 0.5*(r-v) if stabilized else r
        _, v = torch.autograd.functional.jvp(fmap, h, q)
        n = float(v.norm().detach()); total += math.log(max(n,1e-12)); q=v.detach()/max(n,1e-12)
        with torch.no_grad():
            raw,_=gru(u,h.reshape(1,1,-1)); r=raw[0,0]
            h=(h+0.5*(r-h) if stabilized else r).detach()
    return total/8.0


def signature(seed=0):
    ds=get_dataset('dynamics',seed,400,200)
    bnet,bm=train_one_baseline(ds,3e-3,seed+1000)
    inet,im=train_idea(ds,3e-3,seed+1000)
    eb=trained_exponent(bnet,ds['xte'],False); ei=trained_exponent(inet,ds['xte'],True)
    return {'trained_baseline_mse':bm,'trained_idea_mse':im,
            'baseline_discrete_exponent':eb,'controlled_discrete_exponent':ei,
            'predicted_sign_preserved':bool(eb<0),'observed_sign_preserved':bool(ei<0),
            'confirmed':bool((eb<0)==(ei<0))}


def main():
    base=sweep_baseline(baseline_factory, GRID, seeds=(0,1,2,3))
    vals=[idea_factory(base['best_cfg'])(s) for s in DEFAULT_SEEDS]
    idea={'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)}
    sig=signature(0)
    idea_sweep=[]
    for cfg in GRID:
        vals4=[idea_factory(cfg)(s) for s in (0,1,2,3)]
        idea_sweep.append({'cfg':cfg,'mean':float(np.mean(vals4))})
    rep=make_report('dynamics','rnn_small',base,idea,extra={'mechanism_signature':sig,'idea_sweep':idea_sweep,'track_justification':'Controlled pendulum rollout is the built-in stability/control/Lyapunov match.','protocol':'8 paired seeds; baseline and idea share lr union; standard test MSE.'})
    out={'bench_report':rep}; print(json.dumps(out,indent=2)); open('bench_report.json','w').write(json.dumps(out,indent=2))

if __name__=='__main__': main()