import sys, json, math, random
from pathlib import Path
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, train_model, evaluate, sweep_baseline, make_report

OUT = Path('bench_report.json')
SEEDS = tuple(range(8))
# Same learning-rate union is used by baseline and idea.
LR_GRID = [1e-3, 3e-3, 1e-2]
EPOCHS = 14
NTRAIN, NTEST = 900, 300
BATCH = 128


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        try: torch.cuda.manual_seed_all(seed)
        except Exception: pass


def baseline_metric(cfg, seed):
    seed_all(seed)
    d = get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
    try:
        _, metric, _ = train_model(make_model('rnn_small', d['input_shape'], d['out_dim']),
                                   d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH,
                                   weight_decay=cfg['weight_decay'], log=lambda *_: None)
        return float(metric)
    except (RuntimeError, torch.cuda.OutOfMemoryError):
        return _baseline_cpu(cfg, seed)


def _baseline_cpu(cfg, seed):
    # train_model already has fallback; this path is only a defensive retry.
    seed_all(seed); d = get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
    old = torch.cuda.is_available
    try:
        net = make_model('rnn_small', d['input_shape'], d['out_dim']).cpu()
        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        for _ in range(EPOCHS):
            p = torch.randperm(len(d['xtr']))
            for j in range(0, len(p), BATCH):
                ix=p[j:j+BATCH]; loss=((net(d['xtr'][ix])-d['ytr'][ix])**2).mean()
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): return float(((net(d['xte'])-d['yte'])**2).mean())
    finally:
        pass


def integral_regressor(x):
    """Omega = dt sum Phi(theta, omega, u), with Phi=[theta,omega,u,sin(theta)]."""
    z=x.reshape(x.shape[0], 8, 3)
    phi=torch.stack((z[:,:,0], z[:,:,1], z[:,:,2], torch.sin(z[:,:,0])), dim=-1)
    return phi.mean(dim=1)


def gram_score(omega, eps=0.035):
    # Computable conservative certificate q=lambda_min(Ghat)-sum(2||O||e+e^2).
    g=omega.T @ omega
    lam=torch.linalg.eigvalsh(g)[0]
    err=2*torch.linalg.matrix_norm(omega, ord=2)*eps + eps*eps
    return lam - err


def idea_train(cfg, seed, return_info=False):
    seed_all(seed)
    d=get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        net=make_model('rnn_small', d['input_shape'], d['out_dim']).to(device)
        xtr,ytr=d['xtr'].to(device),d['ytr'].to(device)
        xte,yte=d['xte'].to(device),d['yte'].to(device)
        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
        rng=torch.Generator(device=device); rng.manual_seed(seed+10000)
        # fixed-size history; greedy replacement maximizes current minimum eigenvalue
        history=[]; activated=[]; losses=[]; qvals=[]; activation_epoch=None
        for ep in range(EPOCHS):
            net.train(); perm=torch.randperm(len(xtr),generator=rng,device=device); ep_loss=0.
            for j in range(0,len(perm),BATCH):
                ix=perm[j:j+BATCH]; xb,yb=xtr[ix],ytr[ix]
                om=integral_regressor(xb); q=gram_score(om, cfg['eps'])
                qvals.append(float(q.detach().cpu()))
                exciting=bool(q.item()>cfg['gamma'])
                if exciting:
                    if activation_epoch is None: activation_epoch=ep
                    activated.append(1)
                    # Replay selected history plus current batch, not arbitrary old data.
                    cand=(float(torch.linalg.eigvalsh(om.T@om)[0].detach().cpu()), xb.detach(), yb.detach(), om.detach())
                    history.append(cand); history.sort(key=lambda a:a[0],reverse=True); history=history[:cfg['replay']]
                    batches=[(xb,yb)] + [(h[1],h[2]) for h in history[:-1]]
                    opt.zero_grad(); loss=sum(((net(a)-b)**2).mean() for a,b in batches)/len(batches)
                    loss.backward(); opt.step()
                else:
                    activated.append(0)
                    # Conservative policy: freeze adapter update until finite excitation.
                    loss=torch.zeros((),device=device)
                ep_loss += float(loss.detach().cpu())*len(ix)
            losses.append(ep_loss/len(xtr))
        net.eval()
        with torch.no_grad(): metric=float(((net(xte)-yte)**2).mean().cpu())
        info={'metric':metric,'activation_epoch':activation_epoch,
              'activation_rate':float(np.mean(activated)),'mean_q':float(np.mean(qvals)),
              'positive_q_rate':float(np.mean(np.asarray(qvals)>cfg['gamma'])),
              'loss_first':losses[0],'loss_last':losses[-1],
              'post_activation_loss_drop': (float(losses[activation_epoch]-losses[-1]) if activation_epoch is not None else 0.0),
              'model_params':sum(p.numel() for p in net.parameters())}
        return info if return_info else metric
    except (RuntimeError, torch.cuda.OutOfMemoryError):
        # Shared GPU can fail; repeat entirely on CPU.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        old=torch.cuda.is_available
        # identical loop through a temporary CPU-only recursive implementation
        if device.type=='cuda':
            torch.cuda.is_available=lambda: False
            try: return idea_train(cfg,seed,return_info)
            finally: torch.cuda.is_available=old
        raise


def main():
    base_grid=[{'lr':lr,'weight_decay':wd} for lr in LR_GRID for wd in [0.0,1e-4]]
    # Baseline decisive Adam knob (lr and weight decay) is swept. Idea uses same union.
    base=sweep_baseline(lambda cfg: lambda s: baseline_metric(cfg,s), base_grid, seeds=(0,1,2,3))
    idea_cfgs=[{'lr':lr,'weight_decay':base['best_cfg']['weight_decay'], 'gamma':g, 'eps':0.035, 'replay':4}
               for lr in LR_GRID for g in [0.0]]
    # Evaluate the three idea settings on all paired seeds; choose by the same 4-seed tuning split.
    idea_trials=[]
    for cfg in idea_cfgs:
        r=evaluate(lambda s,cfg=cfg: idea_train(cfg,s), seeds=(0,1,2,3))
        idea_trials.append({'cfg':cfg,'mean':r['mean']})
    best=min(idea_trials,key=lambda z:z['mean'])['cfg']
    idea=evaluate(lambda s: idea_train(best,s), seeds=SEEDS)
    # Signature is measured from trained models, not the algebraic toy: aggregate per-seed behavior.
    sig=[]
    for s in SEEDS:
        sig.append(idea_train(best,s,True))
    signature={'prediction':'parameter updates become active after q>gamma and loss then decreases',
      'predicted_vs_observed':{'predicted_positive_q_activation':True,
        'observed_positive_q_rate_mean':float(np.mean([z['positive_q_rate'] for z in sig])),
        'observed_activation_rate_mean':float(np.mean([z['activation_rate'] for z in sig])),
        'observed_post_activation_loss_drop_mean':float(np.mean([z['post_activation_loss_drop'] for z in sig])),
        'activation_epoch_values':[z['activation_epoch'] for z in sig]},
      'confirmed':bool(np.mean([z['positive_q_rate'] for z in sig])>0 and np.mean([z['post_activation_loss_drop'] for z in sig])>0)}
    # Report idea sweep alongside canonical make_report output.
    rep=make_report('dynamics','rnn_small',base,idea,signature)
    rep['idea_sweep']=idea_trials; rep['protocol']={'seeds':list(SEEDS),'n_train':NTRAIN,'n_test':NTEST,'epochs':EPOCHS,
      'structural_match':'controlled pendulum rollout / latent dynamics', 'custom_track':None}
    OUT.write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()