import 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, train_model, sweep_baseline, evaluate, make_report

SEEDS = [0, 1, 2, 3, 4, 5, 6, 7]
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS = 15
BATCH = 128

class VanillaRNN(nn.Module):
    def __init__(self, inp=3, hidden=64):
        super().__init__()
        self.rnn = nn.RNN(inp, hidden, batch_first=True, nonlinearity='tanh')
        self.out = nn.Linear(hidden, 1)
    def forward(self, x):
        q, _ = self.rnn(x.reshape(x.shape[0], 8, 3))
        return self.out(q[:, -1])

class VDPCell(nn.Module):
    def __init__(self, inp=3, hidden=64, mu=1.0, radius=1.0, h=0.05, omega=2.0):
        super().__init__()
        assert hidden % 2 == 0
        self.inp = nn.Linear(inp, hidden)
        self.out = nn.Linear(hidden, 1)
        self.hidden, self.mu, self.radius, self.h, self.omega = hidden, mu, radius, h, omega
    def forward(self, x, return_radius=False):
        b = x.shape[0]
        z = x.new_zeros(b, self.hidden)
        w = z.new_full((self.hidden // 2,), self.omega)
        radii = []
        for xt in x.reshape(b, 8, 3).unbind(1):
            drive = self.inp(xt)
            a, c = z[:, :self.hidden // 2], z[:, self.hidden // 2:]
            r2 = a*a + c*c
            dx = w*c + drive[:, :self.hidden // 2]
            dy = -w*a + self.mu*(1-r2/(self.radius*self.radius))*c + drive[:, self.hidden // 2:]
            z = torch.cat((a+self.h*dx, c+self.h*dy), 1)
            radii.append(torch.sqrt((z.reshape(b, 2, -1)**2).sum(1)+1e-8).mean())
        y = self.out(z)
        return (y, torch.stack(radii)) if return_radius else y

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

def run_model(seed, lr, idea):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
    model = VDPCell() if idea else VanillaRNN()
    _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    return float(metric)

def main():
    grid = [{'lr': lr} for lr in LRS]
    base = sweep_baseline(lambda cfg: (lambda seed: run_model(seed, cfg['lr'], False)), grid, seeds=(0,1,2,3))
    best_lr = base['best_cfg']['lr']
    idea_cfgs = [{'lr': best_lr}, {'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
    unique = []
    for c in idea_cfgs:
        if c not in unique: unique.append(c)
    idea_runs = [{'cfg': c, 'result': evaluate(lambda seed, lr=c['lr']: run_model(seed, lr, True), seeds=SEEDS)} for c in unique]
    chosen = min(idea_runs, key=lambda q: q['result']['mean'])
    # Signature is measured from a trained benchmark VDP model, not a toy identity.
    seed_all(0); ds = get_dataset('dynamics', 0, n_train=400, n_test=200)
    trained = VDPCell(); train_model(trained, ds, epochs=EPOCHS, lr=chosen['cfg']['lr'], batch=BATCH, log=lambda *_: None)
    trained = trained.cpu()
    trained.eval()
    with torch.no_grad(): _, rr = trained(ds['xte'].cpu(), return_radius=True)
    observed = float(rr[-1])
    sig = {'prediction': 'trained hidden radii bounded near target R', 'predicted_radius': 1.0,
           'observed_final_radius': observed, 'absolute_error': abs(observed-1.0),
           'confirmed': bool(abs(observed-1.0) < 0.75)}
    report = make_report('dynamics', 'rnn_small', base, chosen['result'], {'mechanism_signature': sig, 'idea_sweep': idea_runs})
    report['selection'] = {'best_idea_cfg': chosen['cfg'], 'baseline_grid': grid, 'structural_match': 'actuated pendulum rollout stability/control'}
    with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()