import json, random, sys
import numpy as np
import torch
from torch import nn

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, sweep_baseline, make_report, count_params

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
EPOCHS = 15
BATCH = 128
LR_GRID = [1e-3, 3e-3, 6e-3]
ALPHA_GRID = [0.25, 0.5, 0.75]


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 device():
    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def proximal_train(model, ds, lr, alpha, epochs=EPOCHS, batch=BATCH, return_stats=False):
    """Blockwise stochastic proximal response T=w-lr*g, then w<-w+alpha(T-w).
    Blocks are parameter tensors; gradients are computed jointly but applied blockwise,
    making the intervention the optimizer update rather than the network architecture.
    """
    try:
        dev = device(); model = model.to(dev)
        xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev)
    except Exception:
        dev = torch.device('cpu'); model = model.to(dev)
        xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev)
    loss_fn = nn.MSELoss()
    params = [p for p in model.parameters() if p.requires_grad]
    rng = np.random.default_rng(12345)
    update_norms = []; grad_snapshots = []
    n = len(xtr)
    model.train()
    for ep in range(epochs):
        order = rng.permutation(n)
        for start in range(0, n, batch):
            ix = torch.as_tensor(order[start:start+batch], device=dev)
            model.zero_grad(set_to_none=True)
            pred = model(xtr[ix])
            loss = loss_fn(pred, ytr[ix])
            loss.backward()
            norms = []
            with torch.no_grad():
                for p in params:  # each tensor is one proximal block
                    if p.grad is None: continue
                    # T_hat = p - lr*g; relaxed response is p - alpha*lr*g.
                    step = alpha * lr * p.grad
                    p.sub_(step)
                    norms.append(float(step.norm().detach().cpu()))
            if norms:
                update_norms.append(float(np.sqrt(np.sum(np.square(norms)))))
                grad_snapshots.append(float(np.mean(np.square(norms))))
    model.eval()
    with torch.no_grad():
        metric = float(loss_fn(model(ds['xte'].to(dev)), ds['yte'].to(dev)).cpu())
    stats = {'update_rms': float(np.sqrt(np.mean(np.square(update_norms)))) if update_norms else 0.0,
             'update_std': float(np.std(update_norms)) if update_norms else 0.0,
             'n_updates': len(update_norms)}
    return model, metric, stats


def baseline_fn(cfg):
    def run(seed):
        seed_all(seed)
        ds = get_dataset('dynamics', seed, n_train=4000, n_test=1000)
        model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
        _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None)
        return float(metric)
    return run


def idea_fn(cfg, retain=None):
    def run(seed):
        seed_all(seed)
        ds = get_dataset('dynamics', seed, n_train=4000, n_test=1000)
        model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
        net, metric, stats = proximal_train(model, ds, cfg['lr'], cfg['alpha'])
        if retain is not None: retain[seed] = (net, ds, stats)
        return float(metric)
    return run


def main():
    # Baseline includes every lr used by the idea; alpha is the idea-only relaxation knob.
    base = sweep_baseline(baseline_fn, [{'lr': x} for x in LR_GRID], seeds=SWEEP_SEEDS)
    idea_grid = [{'lr': lr, 'alpha': a} for lr, a in zip(LR_GRID, ALPHA_GRID)]
    idea_sweep = []
    for cfg in idea_grid:
        vals = [idea_fn(cfg)(s) for s in SWEEP_SEEDS]
        idea_sweep.append({'cfg': cfg, 'mean': float(np.mean(vals)), 'per_seed': vals})
    best_cfg = min(idea_sweep, key=lambda z: z['mean'])['cfg']

    retained = {}
    idea_vals = [idea_fn(best_cfg, retained)(s) for s in SEEDS]
    idea_res = {'mean': float(np.mean(idea_vals)), 'std': float(np.std(idea_vals)),
                'per_seed': idea_vals, 'n': 8, 'chosen_cfg': best_cfg,
                'sweep': idea_sweep}

    base_vals = []
    for s in SEEDS:
        seed_all(s); ds = get_dataset('dynamics', s, n_train=4000, n_test=1000)
        m = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
        _, metric, _ = train_model(m, ds, epochs=EPOCHS, lr=base['best_cfg']['lr'], batch=BATCH, log=lambda *_: None)
        base_vals.append(float(metric))
    base['full'] = {'mean': float(np.mean(base_vals)), 'std': float(np.std(base_vals)),
                    'per_seed': base_vals, 'n': 8}

    # Signature is measured on trained benchmark models: response/update noise and damping.
    rms = [v[2]['update_rms'] for v in retained.values()]
    std = [v[2]['update_std'] for v in retained.values()]
    alpha = best_cfg['alpha']
    # For the relaxed recursion, observed update magnitude should scale approximately alpha.
    sig = {'prediction': 'relaxation damps stochastic block-response updates approximately linearly in alpha',
           'predicted_alpha': alpha, 'observed_update_rms_mean': float(np.mean(rms)),
           'observed_update_std_mean': float(np.mean(std)),
           'predicted_vs_observed_relation': 'measured from trained rnn_small updates; alpha scaling was not quantitatively tested',
           'confirmed': False}
    report = make_report('dynamics', 'rnn_small', base, idea_res, {'mechanism_signature': sig})
    report['mechanism_signature'] = sig
    report['custom_track'] = None
    with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()