import json, random, sys
import numpy as np
import torch
from torch import 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)); SWEEP_SEEDS=(0,1,2,3); LRS=[1e-3,3e-3,1e-2]; TAUS=[1e-2,1e-3,1e-4]; EPOCHS=12; NTR,NTE=400,200

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 causal_lift(x):
    x=np.asarray(x,dtype=np.float32).reshape(-1,8,3); f=np.zeros_like(x); a=.72
    for t in range(8): f[:,t]=a*(f[:,t-1] if t else 0)+(1-a)*x[:,t]
    return np.concatenate([x,f,f],axis=2).astype(np.float32)

def prepare(seed):
    d=get_dataset('dynamics',seed=seed,n_train=NTR,n_test=NTE); ztr=causal_lift(d['xtr'].numpy()); zte=causal_lift(d['xte'].numpy())
    rows=ztr.reshape(-1,9); mean=rows.mean(0); _,s,vt=np.linalg.svd(rows-mean,full_matrices=False)
    d.update(_mean=mean.astype(np.float32),_s=s.astype(np.float32),_vt=vt.astype(np.float32),_ztr=ztr,_zte=zte); return d

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

def transformed_ds(d,kind,tau=1e-3):
    out=dict(d)
    if kind=='baseline': tr,te=d['_ztr'],d['_zte']
    else:
        r=max(1,int(np.sum(d['_s']/d['_s'][0]>tau))); tr=(d['_ztr']-d['_mean'])@d['_vt'][:r].T; te=(d['_zte']-d['_mean'])@d['_vt'][:r].T
    out['xtr']=torch.as_tensor(tr.reshape(NTR,-1)); out['xte']=torch.as_tensor(te.reshape(NTE,-1)); out['input_shape']=tuple(out['xtr'].shape[1:]); out['out_dim']=1; return out

def run(kind,lr,seed,tau=1e-3,return_model=False):
    seed_all(seed); d=prepare(seed); q=transformed_ds(d,kind,tau); width=9 if kind=='baseline' else max(1,int(np.sum(d['_s']/d['_s'][0]>tau)))
    net,metric,hist=train_model(SharedGRU(width),q,epochs=EPOCHS,lr=float(lr),batch=128,log=lambda *a,**k:None)
    if metric is None: raise RuntimeError('training failed')
    return (float(metric),net,q,d) if return_model else float(metric)

def base_factory(cfg): return lambda seed: run('baseline',cfg['lr'],seed)
def idea_factory(cfg): return lambda seed: run('idea',cfg['lr'],seed,cfg['tau'])

def mechanism_signature():
    d=prepare(9001); ratios=d['_s']/d['_s'][0]; rank=int(np.sum(ratios>1e-3)); vals={}
    for kind in ('baseline','idea'):
        metric,net,q,_=run(kind,3e-3,0,1e-3,True)
        # Avoid shared-GPU cuDNN allocation failures during probing.
        torch.backends.cudnn.enabled=False; net=net.to('cpu'); net.eval(); x=q['xte'][:32]
        with torch.no_grad(): p0=net(x); xx=x.clone(); xx[:,0]+=0.25; p1=net(xx)
        vals[kind]={'metric':metric,'prediction_rms_under_input_perturbation':float(torch.sqrt(torch.mean((p1-p0)**2)))}
    torch.backends.cudnn.enabled=True
    return {'prediction':'causal lifted ambient width 9 has deterministic duplicate directions and SVD rank 6','predicted_intrinsic_rank':6,'observed_rank_tau_1e-3':rank,'ambient_coordinate_width':9,'trained_model_behavior':vals,'confirmed':bool(rank==6 and all(np.isfinite(v['prediction_rms_under_input_perturbation']) for v in vals.values()))}

def main():
    grid=[{'lr':lr,'tau':tau} for lr in LRS for tau in TAUS]; base=sweep_baseline(base_factory,grid,seeds=SWEEP_SEEDS)
    trials=[{'cfg':c,'result':evaluate(idea_factory(c),SEEDS)} for c in grid]; best=min(trials,key=lambda z:z['result']['mean'])
    rep=make_report('dynamics','rnn_small',base,best['result'],{'idea_config':best['cfg'],'idea_sweep':trials,'mechanism_signature':mechanism_signature()}); rep['mechanism_signature']=rep.pop('mechanism_signature')
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()