import sys, json, 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, make_report

SEEDS = tuple(range(8))
# Includes every idea learning rate, and sweeps the baseline's weight decay knob too.
GRID = [
    {'lr': 1e-3, 'weight_decay': 0.0},
    {'lr': 1e-3, 'weight_decay': 1e-4},
    {'lr': 3e-3, 'weight_decay': 0.0},
    {'lr': 3e-3, 'weight_decay': 1e-4},
    {'lr': 6e-3, 'weight_decay': 0.0},
    {'lr': 6e-3, 'weight_decay': 1e-4},
]
# Three nearby settings, including the baseline's eventual best lr.
IDEA_GRID = [
    {'lr': 1e-3, 'weight_decay': 0.0, 'lam': 1e-3},
    {'lr': 3e-3, 'weight_decay': 0.0, 'lam': 1e-3},
    {'lr': 6e-3, 'weight_decay': 0.0, 'lam': 1e-3},
]
EPOCHS, BATCH = 6, 128


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 forward_with_hidden(net, x):
    seq = x.view(x.shape[0], -1, 3)
    h, _ = net.rnn(seq)
    return h, net.head(h[:, -1])


def cross_diffusion_monitor(h):
    """Empirical OU monitor on a two-dimensional learned latent subspace.

    A local drift is fitted from hidden increments and residual covariance gives
    D=residual_cov/(2 dt). The regularizer discourages cross diffusion relative
    to diagonal diffusion; the task loss remains the primary benchmark metric.
    """
    if h.shape[0] < 4 or h.shape[1] < 2:
        return h.new_zeros(()), h.new_zeros(())
    z = h[:, :-1, :2].reshape(-1, 2)
    dz = (h[:, 1:, :2] - h[:, :-1, :2]).reshape(-1, 2)
    # Regression is detached: the monitor estimates latent dynamics, while the
    # covariance penalty remains differentiable through residual increments.
    with torch.no_grad():
        xtx = z.T @ z + 1e-4 * torch.eye(2, device=z.device)
        coef = torch.linalg.solve(xtx, z.T @ dz)
    resid = dz - z @ coef.detach()
    cov = (resid.T @ resid) / max(1, resid.shape[0] - 1)
    d = cov / 2.0
    dxy = d[0, 1]
    scale = d.diag().mean().detach() + 1e-5
    penalty = (dxy / scale).pow(2)
    return penalty, dxy.detach()


def run_one(seed, cfg, idea, n_train=400, n_test=100, diagnostics=False):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=n_train, n_test=n_test)
    net = make_model('rnn_small', ds['input_shape'], ds['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'])
        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
        for _ in range(EPOCHS):
            net.train()
            perm = torch.randperm(len(x), device=device)
            for i in range(0, len(x), BATCH):
                ind = perm[i:i+BATCH]
                h, pred = forward_with_hidden(net, x[ind])
                loss = ((pred-y[ind])**2).mean()
                if idea:
                    reg, _ = cross_diffusion_monitor(h)
                    loss = loss + cfg['lam'] * reg
                opt.zero_grad(); loss.backward(); opt.step()
        net.eval()
        with torch.no_grad():
            hte, out = forward_with_hidden(net, ds['xte'].to(device))
            metric = float(((out-ds['yte'].to(device))**2).mean().cpu())
        if not diagnostics:
            return metric
        with torch.no_grad():
            reg, dxy = cross_diffusion_monitor(hte)
            out = out[:, 0].detach().cpu().numpy()
            hh = hte[:, :, :2].reshape(-1, 2).detach().cpu().numpy()
            # Behavior-derived output-spectrum proxy and reverse/forward gap.
            centered = out - out.mean()
            spectrum0 = float(np.mean(centered**2))
            a = hh[:-1]; b = hh[1:]
            dz = b-a
            coef = np.linalg.solve(a.T@a + 1e-4*np.eye(2), a.T@dz)
            r = dz-a@coef
            cov = np.cov(r, rowvar=False) + 1e-4*np.eye(2)
            inv = np.linalg.inv(cov)
            # Gaussian transition log-likelihood forward vs reversed ordering.
            rf = dz
            rr = a-b-b@coef
            gap = float(np.mean(-0.5*np.einsum('ij,jk,ik->i',rf,inv,rf) + 0.5*np.einsum('ij,jk,ik->i',rr,inv,rr)))
            return metric, {'dxy': float(dxy.cpu()), 'output_power': spectrum0,
                            'reverse_forward_gap': gap, 'monitor': float(reg.cpu())}
    except RuntimeError:
        # Required robust CUDA fallback: recreate and execute on CPU.
        if device != 'cuda': raise
        torch.cuda.empty_cache()
        return run_one_cpu(seed, cfg, idea, n_train, n_test, diagnostics)


def run_one_cpu(seed, cfg, idea, n_train, n_test, diagnostics):
    old = torch.cuda.is_available
    ds = get_dataset('dynamics', seed, n_train=n_train, n_test=n_test)
    net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
    x, y = ds['xtr'], ds['ytr']
    for _ in range(EPOCHS):
        perm = torch.randperm(len(x))
        for i in range(0, len(x), BATCH):
            ind=perm[i:i+BATCH]; h,p=forward_with_hidden(net,x[ind]); loss=((p-y[ind])**2).mean()
            if idea: loss=loss+cfg['lam']*cross_diffusion_monitor(h)[0]
            opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        h,o=forward_with_hidden(net,ds['xte']); metric=float(((o-ds['yte'])**2).mean())
        if not diagnostics: return metric
        reg,dxy=cross_diffusion_monitor(h); q=o[:,0].numpy(); z=h[:,:,:2].reshape(-1,2).numpy()
        
        z=h[:,:,:2].reshape(-1,2).numpy(); aa=z[:-1]; bb=z[1:]; dd=bb-aa
        cc=np.cov(dd-aa@np.linalg.solve(aa.T@aa+1e-4*np.eye(2),aa.T@dd),rowvar=False)+1e-4*np.eye(2); ii=np.linalg.inv(cc)
        gap=float(np.mean(-.5*np.einsum('ij,jk,ik->i',dd,ii,dd)+.5*np.einsum('ij,jk,ik->i',aa-bb-bb@np.linalg.solve(aa.T@aa+1e-4*np.eye(2),aa.T@dd),ii,aa-bb-bb@np.linalg.solve(aa.T@aa+1e-4*np.eye(2),aa.T@dd))))
        return metric, {'dxy':float(dxy),'output_power':float(np.mean((q-q.mean())**2)), 'reverse_forward_gap':gap, 'monitor':float(reg)}


def vals(fn): return [float(fn(s)) for s in SEEDS]


def main():
    def base_make(cfg): return lambda s: run_one(s, cfg, False)
    base = sweep_baseline(base_make, GRID)
    idea_runs=[]
    for cfg in IDEA_GRID:
        v=vals(lambda s, c=cfg: run_one(s,c,True))
        idea_runs.append({'cfg':cfg,'mean':float(np.mean(v)),'per_seed':v})
    best=min(idea_runs,key=lambda r:r['mean'])
    idea_res={'mean':best['mean'],'std':float(np.std(best['per_seed'])),'per_seed':best['per_seed'],'n':8,'cfg':best['cfg'], 'sweep':idea_runs}
    sig_base=[]; sig_idea=[]
    for s in SEEDS:
        _,b=run_one(s,base['best_cfg'],False,diagnostics=True)
        _,a=run_one(s,best['cfg'],True,diagnostics=True)
        sig_base.append(b); sig_idea.append(a)
    def meanfield(rows,k): return float(np.mean([r[k] for r in rows]))
    signature={'predicted': 'hidden cross diffusion can vary while observed output power is comparatively stable',
      'baseline_observed_output_power_std':float(np.std([r['output_power'] for r in sig_base])),
      'idea_observed_output_power_std':float(np.std([r['output_power'] for r in sig_idea])),
      'baseline_hidden_dxy_std':float(np.std([r['dxy'] for r in sig_base])),
      'idea_hidden_dxy_std':float(np.std([r['dxy'] for r in sig_idea])),
      'idea_mean_reverse_forward_gap':meanfield(sig_idea,'reverse_forward_gap'),
      'confirmed': bool(np.std([r['dxy'] for r in sig_idea]) > np.std([r['output_power'] for r in sig_idea]))}
    report=make_report('dynamics','rnn_small',base,idea_res,{'mechanism_signature':signature,
      'protocol_note':'Baseline and idea use identical rnn_small systems; only latent diffusion penalty differs.',
      'idea_grid':idea_runs})
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()