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

TRACK, MODEL = 'dynamics', 'rnn_small'
SEEDS, SWEEP_SEEDS = tuple(range(8)), (0, 1, 2, 3)
EPOCHS, BATCH = 10, 128
# This is the complete shared union: baseline is evaluated at every idea lr/decay.
GRID = [{'lr': lr, 'weight_decay': wd} for lr in (1e-3, 3e-3, 1e-2)
        for wd in (0.0, 1e-4)]


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 transfer_stats(net, x, device):
    """Empirical local linear transfer T from hidden trajectories of trained GRU."""
    net.eval()
    with torch.no_grad():
        seq = x.to(device).view(x.shape[0], -1, 3)
        h = torch.zeros(1, seq.shape[0], net.rnn.hidden_size, device=device)
        hs = []
        for t in range(seq.shape[1]):
            _, h = net.rnn(seq[:, t:t+1], h); hs.append(h[-1])
        A = torch.stack(hs[:-1], 1).reshape(-1, net.rnn.hidden_size)
        B = torch.stack(hs[1:], 1).reshape(-1, net.rnn.hidden_size)
        eye = torch.eye(A.shape[1], device=device)
        T = torch.linalg.solve(A.T @ A + 1e-3 * eye, A.T @ B)
        ev, vec = torch.linalg.eig(T)
        order = torch.argsort(ev.abs(), descending=True); ev, vec = ev[order], vec[:, order]
        mags = ev.abs().real; l0 = max(float(mags[0]), 1e-8)
        l1 = float(mags[1]) if len(mags) > 1 else 0.0
        gap = (l0-l1)/l0
        xi = 1.0/max(math.log(l0/max(l1, 1e-8)), 1e-8)
        V = vec[:, :min(3, vec.shape[1])].real
        z = A @ V; p = (z*z).mean(0); p = p/(p.sum()+1e-8)
        entropy = float(-(p*torch.log(p+1e-8)).sum())
        return gap, entropy, xi


def run(seed, cfg, scheduler, return_state=False):
    seed_all(seed)
    ds = get_dataset(TRACK, seed, n_train=4000, n_test=1000)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net = make_model(MODEL, ds['input_shape'], 1).to(device)
        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        lossf = nn.MSELoss(); xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device)
        current_lr = float(cfg['lr']); stats=[]
        for _ in range(EPOCHS):
            net.train(); perm=torch.randperm(len(xtr), device=device)
            for i in range(0, len(xtr), BATCH):
                idx=perm[i:i+BATCH]; loss=lossf(net(xtr[idx]), ytr[idx])
                opt.zero_grad(set_to_none=True); loss.backward()
                torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0); opt.step()
            if scheduler:
                gap, entropy, xi=transfer_stats(net, xtr[:min(512,len(xtr))], device)
                stats.append((gap, entropy, xi))
                current_lr=float(np.clip(current_lr*np.clip(gap/.25,.5,1.25), 1e-5, cfg['lr']))
                if entropy>.70 and gap<.12: current_lr=max(.5*current_lr,1e-5)
                for pg in opt.param_groups: pg['lr']=current_lr
        net.eval()
        with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
        state={'gap':stats[-1][0], 'entropy':stats[-1][1], 'xi':stats[-1][2],
               'lr_final':current_lr, 'trigger_count':sum(g<.12 and h>.70 for g,h,_ in stats)} if stats else {}
        return (metric, state) if return_state else metric
    except RuntimeError:
        if device == 'cuda':
            torch.cuda.empty_cache(); old=torch.cuda.is_available; torch.cuda.is_available=lambda:False
            try: return run(seed,cfg,scheduler,return_state)
            finally: torch.cuda.is_available=old
        raise


def main():
    base=sweep_baseline(lambda c: lambda s: run(s,c,False), GRID, seeds=SWEEP_SEEDS)
    # Three idea settings, all present in baseline sweep; report the best idea.
    idea_candidates=[]
    for cfg in GRID:
        r=evaluate(lambda s, c=cfg: run(s,c,True), seeds=SEEDS)
        idea_candidates.append({'cfg':cfg, 'result':r})
    best=min(idea_candidates, key=lambda z:z['result']['mean'])
    idea=best['result']; sig=[]
    for s in SEEDS: _, st=run(s,best['cfg'],True,return_state=True); sig.append(st)
    gaps=np.array([x['gap'] for x in sig]); ents=np.array([x['entropy'] for x in sig])
    corr=float(np.corrcoef(gaps,ents)[0,1]) if np.std(gaps)>0 and np.std(ents)>0 else 0.0
    signature={'prediction':'low normalized transfer gap with high projection entropy marks a pseudo-transition',
      'predicted_thresholds':{'gap_lt':0.12,'entropy_gt':0.70},
      'observed_mean_gap':float(gaps.mean()),'observed_mean_entropy':float(ents.mean()),
      'gap_entropy_correlation':corr,'trigger_count':int(sum(x['trigger_count'] for x in sig)),
      'confirmed':bool(np.any(gaps<.12) and np.any(ents>.70))}
    base['idea_grid']= [{'cfg':z['cfg'],'mean':z['result']['mean']} for z in idea_candidates]
    report=make_report(TRACK,MODEL,base,idea,signature)
    report['idea']['best_cfg']=best['cfg']
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()