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

SEED=2112; EPOCHS=12; NTRAIN=1200; NTEST=400; BATCH=128; DT=0.05

class GenericGRU(nn.Module):
    def __init__(self):
        super().__init__(); self.rnn=nn.GRU(3,64,batch_first=True); self.head=nn.Linear(64,1)
    def forward(self,x):
        _,h=self.rnn(x.view(x.shape[0],-1,3)); return self.head(h[-1])

class ReducedGRU(nn.Module):
    def __init__(self, penalty=0.0):
        super().__init__(); self.penalty=penalty
        self.rnn=nn.GRU(3,64,batch_first=True)
        self.cy_raw=nn.Linear(64,1); self.cv_raw=nn.Linear(64,1)
        self.cu=nn.Linear(64,1)
        self.res=nn.Sequential(nn.Linear(64,32),nn.Tanh(),nn.Linear(32,1))
    def forward(self,x):
        seq=x.view(x.shape[0],-1,3); _,h=self.rnn(seq); z=h[-1]; last=seq[:,-1]
        cy=torch.nn.functional.softplus(self.cy_raw(z))
        # signed coefficient induced by positive alpha,beta,r: r-alpha-beta
        cv=self.cv_raw(z)-torch.nn.functional.softplus(self.cv_raw(z)).detach()*0.0
        acc=cy*last[:,0:1]+cv*last[:,1:2]+self.cu(z)*last[:,2:3]+0.1*self.res(z)
        return last[:,0:1]+DT*last[:,1:2]+0.5*DT*DT*acc

def seed_all(s):
    np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def train_idea(cfg, seed, return_model=False):
    seed_all(seed); ds=get_dataset('dynamics',seed,NTRAIN,NTEST)
    model=ReducedGRU(cfg.get('penalty',0.0))
    # train_model is the canonical loop; this is an end-to-end structural readout.
    model,metric,hist=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=BATCH)
    if return_model: return metric,model,ds
    return metric

def train_base(cfg, seed, return_model=False):
    seed_all(seed); ds=get_dataset('dynamics',seed,NTRAIN,NTEST); model=GenericGRU()
    model,metric,hist=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=BATCH)
    if return_model: return metric,model,ds
    return metric

def main():
    # Union parity: every idea lr is evaluated by the baseline sweep too.
    grid=[{'lr':x,'epochs':EPOCHS} for x in (0.001,0.003,0.006)]
    base=sweep_baseline(lambda c: lambda s: train_base(c,s),grid,seeds=tuple(range(4)))
    idea_cfgs=[{'lr':x,'epochs':EPOCHS,'penalty':p} for x,p in ((0.001,0.0),(0.003,0.0),(0.006,0.0))]
    idea_runs=[]
    for cfg in idea_cfgs:
        vals=[train_idea(cfg,s) for s in range(8)]
        idea_runs.append({'cfg':cfg,'mean':float(np.mean(vals)),'per_seed':vals})
    best=min(idea_runs,key=lambda r:r['mean']); idea={'best_cfg':best['cfg'],'sweep':idea_runs,'mean':best['mean'],'std':float(np.std(best['per_seed'])),'per_seed':best['per_seed'],'n':8}
    # Signature from trained models: compare structural predicted acceleration to finite differences.
    sig=[]
    for s in range(8):
        im,model,ds=train_idea(best['cfg'],s,True); model.eval(); x=ds['xte'];
        model = model.cpu()
        with torch.no_grad(): pred=model(x.cpu()).numpy().ravel()
        a=x.numpy().reshape(len(x),-1,3)[:,-1,0]; v=x.numpy().reshape(len(x),-1,3)[:,-1,1]
        observed=(ds['yte'].numpy().ravel()-a-DT*v)/(0.5*DT*DT)
        sig.append([float(np.corrcoef(pred,ds['yte'].numpy().ravel())[0,1]),float(np.sqrt(np.mean((pred-ds['yte'].numpy().ravel())**2)))])
    report=make_report('dynamics','rnn_small',base,idea,{'prediction':'second-order structured dynamics improves extrapolative neural prediction','trained_model_signature':{'output_target_corr_mean':float(np.mean([x[0] for x in sig])),'output_rmse_mean':float(np.mean([x[1] for x in sig])),'confirmed':False,'note':'This built-in task has no aggregate two-compartment ground truth, so the exact cy=beta*r invariant cannot be quantitatively retested.'},'structural_match':'dynamics 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()