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

SEEDS = tuple(range(8)); EPOCHS = 30; BATCH = 128
GRID = [{'lr': 1e-3, 'rho_cap': None}, {'lr': 3e-3, 'rho_cap': None}, {'lr': 5e-3, 'rho_cap': None}]
IDEA_GRID = [{'lr': 1e-3, 'rho_cap': .90}, {'lr': 3e-3, 'rho_cap': .90}, {'lr': 5e-3, 'rho_cap': .90}]

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 operator_gain(W):
    return float(np.linalg.svd(W, compute_uv=False)[0])

def project_gru(net, cap):
    with torch.no_grad():
        w = net.rnn.weight_hh_l0
        blocks = w.chunk(3, 0)
        gains = [float(torch.linalg.matrix_norm(b, 2).cpu()) for b in blocks]
        g = max(gains)
        if g > cap: w.mul_(cap / (g + 1e-12))
        return g

def train(seed, lr, cap=None, return_net=False):
    seed_all(seed); d = get_dataset('dynamics', seed, n_train=400, n_test=200)
    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net.to(device); x, y = d['xtr'].to(device), d['ytr'].to(device)
        opt = torch.optim.Adam(net.parameters(), lr=lr); lossf = nn.MSELoss(); hist=[]
        for _ in range(EPOCHS):
            net.train(); perm=torch.randperm(len(x), device=device); total=0.
            for i in range(0,len(x),BATCH):
                ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
                if cap is not None: project_gru(net, cap)
                total += float(loss.detach())*len(ix)
            hist.append(total/len(x))
        net.eval()
        with torch.no_grad(): metric=float(lossf(net(d['xte'].to(device)), d['yte'].to(device)).cpu())
        return (metric, net, d) if return_net else metric
    except RuntimeError:
        seed_all(seed); net=make_model('rnn_small', d['input_shape'], d['out_dim']); net.to('cpu')
        x,y=d['xtr'],d['ytr']; opt=torch.optim.Adam(net.parameters(),lr=lr); lossf=nn.MSELoss()
        for _ in range(EPOCHS):
            perm=torch.randperm(len(x))
            for i in range(0,len(x),BATCH):
                ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]),y[ix]); opt.zero_grad(); loss.backward(); opt.step()
                if cap is not None: project_gru(net,cap)
        with torch.no_grad(): metric=float(lossf(net(d['xte']),d['yte']))
        return (metric,net,d) if return_net else metric

def baseline_factory(cfg): return lambda seed: train(seed, cfg['lr'], None)
def idea_factory(cfg): return lambda seed: train(seed, cfg['lr'], cfg['rho_cap'])

def signature():
    rows=[]
    for seed in (0,1,2,3):
        _, net, d = train(seed, 3e-3, .90, True); dev=next(net.parameters()).device
        x=d['xte'][:32].to(dev); delta=torch.randn_like(x)*1e-3
        with torch.no_grad(): z1=net(x); z2=net(x+delta)
        observed=float(((z2-z1).norm(dim=1)/(delta.view(len(delta),-1).norm(dim=1)+1e-12)).mean().cpu())
        blocks=net.rnn.weight_hh_l0.detach().cpu().numpy().reshape(3,64,64)
        pred=max(operator_gain(b) for b in blocks)
        rows.append({'seed':seed,'predicted_gain_proxy':pred,'observed_output_gain':observed})
    maxpred=max(r['predicted_gain_proxy'] for r in rows)
    return {'prediction':'projected recurrent operator-gain proxy stays <= 0.90', 'predicted_max':maxpred,
            'observed':rows, 'confirmed': bool(maxpred <= .90+1e-4)}

def main():
    base=sweep_baseline(baseline_factory, GRID, seeds=(0,1,2,3))
    nearby={str(c['lr']): evaluate(idea_factory(c), seeds=SEEDS) for c in IDEA_GRID}
    best=min(nearby, key=lambda k: nearby[k]['mean']); idea=nearby[best]
    rep=make_report('dynamics','rnn_small',base,idea,{'signature':signature(), 'idea_sweep':nearby,
      'selected_idea_lr':float(best), 'structure':'GRU recurrent dynamics; spectral-gain projection is the only intervention'})
    rep['baseline_sweep_union_note']='Baseline evaluated all idea learning rates; baseline method has no additional decisive knob.'
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()