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

SEEDS=tuple(range(8))
GRID=[{'lr':1e-3,'penalty':0.0},{'lr':3e-3,'penalty':0.0},{'lr':6e-3,'penalty':0.0},
      {'lr':1e-3,'penalty':0.01},{'lr':3e-3,'penalty':0.01},{'lr':6e-3,'penalty':0.01},
      {'lr':1e-3,'penalty':0.05},{'lr':3e-3,'penalty':0.05},{'lr':6e-3,'penalty':0.05}]
EPOCHS=14
sig_rows=[]

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 baseline_fn(cfg):
    def run(seed):
        seed_all(seed)
        ds=get_dataset('dynamics',seed,n_train=800,n_test=300)
        model=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
        net, metric, _=train_model(model,ds,epochs=EPOCHS,lr=cfg['lr'],batch=128)
        # Behavior signature measured on this trained baseline model.
        dev=next(net.parameters()).device
        with torch.no_grad():
            xt=ds['xte'].to(dev)
            pred=net(xt); theta=xt[:,21:22]
            viol=(pred.abs()>0.94*theta.abs()).float().mean().item()
            ratio=(pred.abs()/(theta.abs()+1e-3)).mean().item()
        sig_rows.append({'kind':'baseline','seed':seed,'lr':cfg['lr'],
                         'mse':float(metric),'contraction_violation_rate':viol,
                         'mean_pred_abs_over_theta':ratio})
        return float(metric)
    return run

def idea_train(seed, cfg):
    seed_all(seed)
    ds=get_dataset('dynamics',seed,n_train=800,n_test=300)
    model=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
    for device in (['cuda'] if torch.cuda.is_available() else []) + ['cpu']:
        try:
            net=model.to(device); x=ds['xtr'].to(device); y=ds['ytr'].to(device)
            opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
            rho=0.94
            for _ in range(EPOCHS):
                net.train(); perm=torch.randperm(len(x),device=device)
                for i in range(0,len(x),128):
                    ix=perm[i:i+128]; pred=net(x[ix]); target=y[ix]
                    theta=x[ix,21:22]
                    lyap=torch.relu(pred.abs()-rho*theta.abs())**2
                    loss=((pred-target)**2).mean()+cfg['penalty']*lyap.mean()
                    opt.zero_grad(); loss.backward(); opt.step()
            net.eval()
            with torch.no_grad():
                pred=net(ds['xte'].to(device)); target=ds['yte'].to(device)
                mse=((pred-target)**2).mean().item()
                theta=ds['xte'][:,21:22].to(device)
                ratio=(pred.abs()/(theta.abs()+1e-3)).mean().item()
                violation=(pred.abs()>rho*theta.abs()).float().mean().item()
            sig_rows.append({'kind':'idea','seed':seed,'lr':cfg['lr'],'penalty':cfg['penalty'],
                             'mse':mse,'mean_pred_abs_over_theta':ratio,
                             'contraction_violation_rate':violation})
            return float(mse)
        except RuntimeError:
            if device=='cuda': torch.cuda.empty_cache(); continue
            raise
    raise RuntimeError('training failed')

def idea_fn(cfg): return lambda seed: idea_train(seed,cfg)

def main():
    base=sweep_baseline(baseline_fn,GRID)
    blr=base['best_cfg']['lr']
    idea_grid=[{'lr':blr,'penalty':0.01},{'lr':blr,'penalty':0.05},
               {'lr':(blr/3 if blr>1.1e-3 else 3e-3),'penalty':0.01}]
    tried=[]; best=None
    for cfg in idea_grid:
        r=evaluate(idea_fn(cfg),SEEDS); tried.append({'cfg':cfg,'mean':r['mean'],'full':r})
        if best is None or r['mean']<best['result']['mean']:
            best={'cfg':cfg,'result':r}
    chosen=best['cfg']
    ir=[r for r in sig_rows if r.get('kind')=='idea' and r.get('penalty')==chosen['penalty'] and r.get('lr')==chosen['lr']]
    br=[r for r in sig_rows if r.get('kind')=='baseline' and r.get('lr')==blr]
    iv=float(np.mean([r['contraction_violation_rate'] for r in ir]))
    bv=float(np.mean([r['contraction_violation_rate'] for r in br]))
    signature={'prediction':'Lyapunov penalty lowers observed contraction-violation rate',
               'predicted_direction':'lower','baseline_violation_rate':bv,
               'idea_violation_rate':iv,
               'baseline_mean_abs_prediction_over_theta':float(np.mean([r['mean_pred_abs_over_theta'] for r in br])),
               'idea_mean_abs_prediction_over_theta':float(np.mean([r['mean_pred_abs_over_theta'] for r in ir])),
               'rho':0.94,'confirmed':bool(iv < bv)}
    report=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},best['result'],
                       {'mechanism_signature':signature,'idea_sweep':tried,
                        'protocol_notes':'Dynamics is the structurally matched control/stability track; identical rnn_small systems and datasets, only training loss differs.'})
    report['idea_selected_cfg']=chosen
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()