import sys, math, json
from pathlib import Path
import numpy as np
import torch
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, sweep_baseline, make_report

TRACK='dynamics'; MODEL='rnn_small'; EPOCHS=8; BATCH=128
SEEDS=tuple(range(8)); SWEEP_SEEDS=tuple(range(4))
# The union of all idea and baseline settings is evaluated by the baseline sweep.
GRID=[{'lr':1e-3,'rounds':4},{'lr':3e-3,'rounds':4},{'lr':1e-2,'rounds':4}]

# Iterative refinement map: each round applies the same trained rnn to a shifted
# state window. The controller uses only the model's observed output trajectory.
def rollout(net, x, rounds=4, controller=False, eps=0.01, delta=0.02):
    device=next(net.parameters()).device
    q=x.to(device).clone(); preds=[]; changes=[]; transitions=[]; stable=0
    with torch.no_grad():
        for t in range(rounds):
            pred=net(q); val=pred[:,0]
            if preds:
                change=float(torch.mean(torch.abs(pred-preds[-1])).item())
                changes.append(change); transitions.append((changes[-2] if len(changes)>1 else change, change))
                stable=stable+1 if change < delta else 0
            preds.append(pred.detach())
            q=q.clone(); q[:,:-3]=q[:,3:]; q[:,-3]=val
            if controller and len(changes)>=2 and stable>=2:
                # Fit the scalar Koopman map for the observed disagreement D_t.
                d0=np.asarray([a for a,b in transitions],dtype=float)
                d1=np.asarray([b for a,b in transitions],dtype=float)
                X=np.c_[d0,np.ones_like(d0)]
                coef=np.linalg.lstsq(X,d1,rcond=None)[0]
                lam=abs(float(coef[0])); D=changes[-1]
                remaining=max(0,rounds-t-1)
                future=D*(lam**remaining) if 0 < lam < 1 else D
                # Conservative gate: only stop when predicted disagreement and
                # recent observed changes are both below tolerance.
                if future <= eps and change < delta:
                    break
    arr=torch.stack(preds,dim=1)[:, :, 0]
    return arr, len(preds), changes

def train_score(seed, lr, controller, rounds=4):
    torch.manual_seed(seed); np.random.seed(seed)
    ds=get_dataset(TRACK, seed, n_train=400, n_test=400)
    net=make_model(MODEL, ds['input_shape'], ds['out_dim'])
    net, _, _=train_model(net, ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    if net is None: return float('nan'), {'rounds':float('nan')}, None
    net.eval(); device=next(net.parameters()).device
    pred,n,ch=rollout(net,ds['xte'],rounds,controller=controller)
    y=ds['yte'][:,0].to(device)
    mse=float(torch.mean((pred[:,-1]-y)**2).item())
    return mse, {'rounds':n,'changes':ch}, net

def one(seed,cfg,controller): return train_score(seed,cfg['lr'],controller,cfg['rounds'])
def baseline_sweep():
    def factory(cfg):
        return lambda seed: one(int(seed),cfg,False)[0]
    return sweep_baseline(factory,GRID,seeds=SWEEP_SEEDS)

def idea_eval(cfg,seeds=SEEDS):
    vals=[]; aux=[]
    for s in seeds:
        m,a,_=one(s,cfg,True); vals.append(m); aux.append(a)
    return {'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)},aux

def signature(cfg):
    fitted=[]; observed=[]
    for s in SEEDS:
        _,_,net=train_score(s,cfg['lr'],False,cfg['rounds'])
        ds=get_dataset(TRACK,s,n_train=400,n_test=80); net.eval()
        _,_,ch=rollout(net,ds['xte'],cfg['rounds'],False)
        if len(ch)>=2:
            x=np.asarray(ch[:-1]); y=np.asarray(ch[1:])
            lam=float(np.linalg.lstsq(np.c_[x,np.ones_like(x)],y,rcond=None)[0][0])
            fitted.append(abs(lam)); observed.append(float(np.mean(y/np.maximum(x,1e-8))))
    pf=float(np.mean(fitted)); po=float(np.mean(observed))
    return {'quantity':'NN observed disagreement contraction','predicted_mean_ratio':pf,'observed_mean_ratio':po,'absolute_error':abs(pf-po),'n_trajectories':len(fitted),'confirmed':bool(abs(pf-po)<0.10)}

def main():
    sweep=baseline_sweep(); base_cfg=sweep['best_cfg']
    idea_grid=[]
    for cfg in GRID:
        r,_=idea_eval(cfg); idea_grid.append({'cfg':cfg,'mean':r['mean']})
    best_idea_cfg=min(idea_grid,key=lambda z:z['mean'])['cfg']
    idea_res,aux=idea_eval(best_idea_cfg,SEEDS)
    base_block={'sweep':sweep,'best_cfg':base_cfg,'full':sweep['full']}
    rep=make_report(TRACK,MODEL,base_block,idea_res,{'mechanism_signature':signature(best_idea_cfg),'protocol_note':'Matched dynamics task and freshly trained identical rnn_small architecture; only the inference stopping rule differs.'})
    rep['idea_sweep']=idea_grid; rep['idea_best_cfg']=best_idea_cfg; rep['idea_auxiliary_rounds']=aux
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()