import sys, json, math
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

_LAST_NET = None
_LAST_X = None
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report

# Pullback monitor: replay the same additive latent perturbation for K replicas.
def pairwise_rms(z):
    d = z[:, None, :] - z[None, :, :]
    return torch.sqrt(torch.mean(torch.sum(d*d, dim=-1)) + 1e-12)

def pullback_penalty(net, x, steps=4, replicas=8, noise_scale=0.03, target_lambda=-0.08):
    # Use the trained GRU's recurrent cell on a short observed dynamics prefix.
    seq = x.view(x.shape[0], -1, 3)[:1]
    with torch.no_grad():
        _, h0 = net.rnn(seq)
    h = h0[-1].detach().repeat(replicas, 1)
    init = torch.randn_like(h)
    init = init / (init.norm(dim=1, keepdim=True) + 1e-8) * 0.5
    h = h + init
    # Same noise vector at every step for all replicas; noise is diagnostic-only.
    ds = []
    for r in range(steps):
        inp = torch.zeros(replicas, 3, device=h.device)
        h = net.rnn(inp.unsqueeze(1), h.unsqueeze(0))[1][-1]
        if noise_scale:
            eps = torch.randn(1, h.shape[1], device=h.device) * noise_scale
            h = h + eps
        ds.append(pairwise_rms(h))
    logs = torch.log(torch.stack(ds))
    # Penalize growth beyond the desired contraction envelope.
    t = torch.arange(1, steps + 1, device=h.device, dtype=h.dtype)
    excess = torch.relu(logs[1:] - logs[:-1] - target_lambda)
    return torch.mean(excess**2), float((logs[-1]-logs[0]).detach().cpu())

def train_one(seed, cfg, idea):
    global _LAST_NET, _LAST_X
    torch.manual_seed(seed); np.random.seed(seed)
    d = get_dataset('dynamics', seed, n_train=400, n_test=160)
    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net.to(device)
        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        lossf = nn.MSELoss(); xtr,ytr=d['xtr'].to(device),d['ytr'].to(device)
        for ep in range(cfg['epochs']):
            net.train(); perm=torch.randperm(len(xtr),device=device)
            for i in range(0,len(xtr),128):
                ix=perm[i:i+128]; pred=net(xtr[ix]); loss=lossf(pred,ytr[ix])
                if idea:
                    reg,_=pullback_penalty(net,xtr[ix],target_lambda=cfg['target_lambda'])
                    loss=loss+cfg['attr_weight']*reg
                opt.zero_grad(); loss.backward(); opt.step()
        net.eval()
        with torch.no_grad(): metric=float(lossf(net(d['xte'].to(device)),d['yte'].to(device)).cpu())
        _LAST_NET, _LAST_X = net, d['xte'][:1].to(device)
        return metric
    except RuntimeError:
        # CPU fallback, matching the harness safety requirement.
        net=net.cpu(); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
        xtr,ytr=d['xtr'],d['ytr']
        for ep in range(cfg['epochs']):
            perm=torch.randperm(len(xtr))
            for i in range(0,len(xtr),128):
                ix=perm[i:i+128]; loss=lossf(net(xtr[ix]),ytr[ix])
                if idea: loss=loss+cfg['attr_weight']*pullback_penalty(net,xtr[ix],target_lambda=cfg['target_lambda'])[0]
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): metric=float(lossf(net(d['xte']),d['yte']))
        _LAST_NET, _LAST_X = net, d['xte'][:1]
        return metric

def signature(seed=0, cfg=None):
    # Uses the final trained model retained by evaluate(), not an analytic toy.
    global _LAST_NET, _LAST_X
    if _LAST_NET is None:
        return {'error': 'no trained model', 'confirmed': False}
    net, x = _LAST_NET, _LAST_X
    vals=[]
    old = torch.backends.cudnn.enabled
    try:
        torch.backends.cudnn.enabled = False
        net.eval()
        with torch.no_grad():
            seq=x.view(1,-1,3); _,h0=net.rnn(seq); h=h0[-1].repeat(16,1)
            q=torch.randn_like(h); q=q/(q.norm(dim=1,keepdim=True)+1e-8)*.5; h+=q
            for _ in range(6):
                inp=torch.zeros(16,1,3,device=h.device)
                h=net.rnn(inp,h.unsqueeze(0))[1][-1]
                vals.append(float(pairwise_rms(h).cpu()))
    except RuntimeError:
        net=net.cpu(); x=x.cpu(); torch.backends.cudnn.enabled=False
        with torch.no_grad():
            seq=x.view(1,-1,3); _,h0=net.rnn(seq); h=h0[-1].repeat(16,1)
            q=torch.randn_like(h); q=q/(q.norm(dim=1,keepdim=True)+1e-8)*.5; h+=q
            for _ in range(6):
                h=net.rnn(torch.zeros(16,1,3),h.unsqueeze(0))[1][-1]
                vals.append(float(pairwise_rms(h)))
    finally:
        torch.backends.cudnn.enabled=old
    observed=float(np.polyfit(np.arange(6),np.log(np.maximum(vals,1e-12)),1)[0])
    return {'predicted_log_diameter_step': -0.08, 'observed_log_diameter_slope': observed,
            'final_diameter': vals[-1], 'confirmed': bool(observed < 0 and abs(observed+0.08)<0.15)}

def main():
    # Union parity: every idea lr is also in baseline sweep.
    grid=[{'lr':lr,'weight_decay':wd,'epochs':12,'attr_weight':0.0,'target_lambda':-0.08} for lr in (0.0015,0.003,0.006) for wd in (0.0,1e-4)]
    base=sweep_baseline(lambda c: lambda s: train_one(s,c,False),grid)
    idea_cfgs=[dict(c,attr_weight=w,target_lambda=t) for c in [base['best_cfg']] for w,t in ((0.001,-0.08),(0.01,-0.08),(0.03,-0.04))]
    idea_runs=[]
    for c in idea_cfgs:
        idea_runs.append((c,evaluate(lambda s,c=c: train_one(s,c,True))))
    best_cfg,best=min(idea_runs,key=lambda z:z[1]['mean'])
    rep=make_report('dynamics','rnn_small',base,best,{'prediction':'trained hidden pullback diameter contracts exponentially under shared replay','config':best_cfg,**signature(0,best_cfg)})
    rep['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
    rep['track_justification']='Dynamics is structurally matched: the idea diagnoses stability and contraction of recurrent latent trajectories.'
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()