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

SEEDS=tuple(range(8)); EPOCHS=10; BATCH=128
LR_GRID=[1e-3,3e-3,1e-2]
SIGMA_GRID=[0.01,0.03,0.1]

class MatchedRNN(nn.Module):
    """Matched recurrent baseline/idea; noise is added after every GRUCell step."""
    def __init__(self, out_dim, sigma=0.0):
        super().__init__(); self.sigma=float(sigma)
        self.inp=nn.Linear(3,64); self.cell=nn.GRUCell(64,64); self.head=nn.Linear(64,out_dim)
    def forward(self,x,return_hidden=False):
        seq=x.reshape(x.shape[0],-1,3)
        h=torch.tanh(self.inp(seq[:,0])); states=[]
        for t in range(seq.shape[1]):
            z=self.inp(seq[:,t])
            h=self.cell(z,h)
            if self.training and self.sigma>0: h=h+self.sigma*torch.randn_like(h)
            states.append(h)
        hs=torch.stack(states,1)
        out=self.head(h)
        return (out,hs) if return_hidden else out

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 train_one(seed,cfg,with_noise=False,signature=False):
    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
    sigma=cfg.get('sigma',0.0) if with_noise else 0.0
    net=MatchedRNN(ds['out_dim'],sigma)
    trained,metric,_=train_model(net,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=BATCH,log=lambda *a,**k:None)
    if trained is None: return float('inf') if not signature else {'metric':float('inf')}
    if not signature: return float(metric)
    trained.eval(); dev=next(trained.parameters()).device; x=ds['xte'].to(dev)
    with torch.no_grad(): _,h=trained(x,return_hidden=True)
    return {'metric':float(metric),'hidden_variance':float(h.var().item()),
            'temporal_energy':float(((h[:,1:]-h[:,:-1])**2).mean().item()),'sigma':float(sigma)}

def main():
    grid=[{'lr':lr,'epochs':EPOCHS,'sigma':0.0} for lr in LR_GRID]
    base=sweep_baseline(lambda c:lambda s:train_one(s,c,False),grid,seeds=(0,1,2,3))
    best_lr=base['best_cfg']['lr']
    idea_cfgs=[{'lr':best_lr,'epochs':EPOCHS,'sigma':s} for s in SIGMA_GRID]
    runs=[]
    for c in idea_cfgs:
        runs.append({'cfg':c,'result':evaluate(lambda s,c=c:train_one(s,c,True),seeds=SEEDS)})
    best_cfg,best=min(((q['cfg'],q['result']) for q in runs),key=lambda z:z[1]['mean'])
    sig0=train_one(0,best_cfg,True,True)
    sig_base=train_one(0,{'lr':best_cfg['lr'],'epochs':EPOCHS,'sigma':0.0},False,True)
    sig_rows=[]
    for q in runs:
        s=q['cfg']['sigma']; z=train_one(0,q['cfg'],True,True)
        sig_rows.append({'sigma':s,'sigma2':s*s,'observed_temporal_energy':z['temporal_energy']})
    x=np.array([r['sigma2'] for r in sig_rows]); y=np.array([r['observed_temporal_energy'] for r in sig_rows])
    slope=float(np.dot(x,y)/np.dot(x,x)); r2=float(1-np.sum((y-slope*x)**2)/max(np.sum((y-y.mean())**2),1e-12))
    signature={'prediction':'trained recurrent hidden temporal energy is nonzero and increases with sigma^2',
      'trained_model_observed':sig_rows,'baseline_sigma0':sig_base,'best_sigma':sig0,
      'through_origin_slope':slope,'through_origin_r2':r2,
      'confirmed':bool(r2>0.8 and y[-1]>y[0]),
      'note':'measured from trained benchmark model hidden trajectories'}
    report=make_report('dynamics','rnn_small',base,best,extra={'mechanism_signature':signature,
      'idea_sweep':runs,'track_match':'stability/control and recurrent state evolution -> dynamics'})
    report['custom_track']=None
    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__': main()