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

SEEDS = tuple(range(8))
LRS = [0.001, 0.003, 0.01]
EPOCHS, BATCH = 18, 128
NU_MULTS = [0.0, 0.05, 0.15]


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 get_data(seed):
    return get_dataset('tabular', seed, n_train=400, n_test=400)


def hessian_scale(model, ds, dev):
    """Power iteration on the minibatch Hessian, giving an actual curvature scale."""
    model.zero_grad(set_to_none=True)
    x, y = ds['xtr'][:128].to(dev), ds['ytr'][:128].to(dev)
    loss = nn.MSELoss()(model(x), y)
    gs = torch.autograd.grad(loss, tuple(model.parameters()), create_graph=True)
    vs = [torch.randn_like(p) for p in model.parameters()]
    norm = torch.sqrt(sum((v*v).sum() for v in vs))
    vs = [v / norm for v in vs]
    val = 1e-4
    for _ in range(4):
        dot = sum((g*v).sum() for g, v in zip(gs, vs))
        hv = torch.autograd.grad(dot, tuple(model.parameters()), retain_graph=True)
        norm = torch.sqrt(sum((h*h).sum() for h in hv))
        val = float(norm.detach().cpu())
        vs = [h / (norm + 1e-12) for h in hv]
    model.zero_grad(set_to_none=True)
    return max(val, 1e-5)


def baseline_one(cfg, seed):
    seed_all(seed); ds = get_data(seed)
    model = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
    _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH,
                               weight_decay=cfg['weight_decay'], log=lambda *a: None)
    return float('inf') if metric is None else metric


def shifted_one(cfg, seed, collect=False):
    seed_all(seed); ds = get_data(seed); model = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
    try:
        dev = 'cuda' if torch.cuda.is_available() else 'cpu'
        model = model.to(dev)
        initial = [p.detach().clone() for p in model.parameters()]
        scale = hessian_scale(model, ds, dev); nu = cfg['nu_mult'] * scale
        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        x, y = ds['xtr'].to(dev), ds['ytr'].to(dev); lossf = nn.MSELoss(); hist = []
        for _ in range(EPOCHS):
            for ix in torch.randperm(len(x), device=dev).split(BATCH):
                opt.zero_grad(set_to_none=True); loss = lossf(model(x[ix]), y[ix]); loss.backward()
                # Apply the negative quadratic term to displacement delta=p-p0.
                with torch.no_grad():
                    for p, p0 in zip(model.parameters(), initial):
                        if p.grad is not None: p.grad.add_(-nu * (p - p0))
                torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0); opt.step()
            hist.append(float(loss.detach().cpu()))
        with torch.no_grad():
            test = float(lossf(model(ds['xte'].to(dev)), ds['yte'].to(dev)).cpu())
            disp = float(torch.sqrt(sum(((p-p0)**2).sum() for p,p0 in zip(model.parameters(), initial))).cpu())
        if collect: return test, {'model': model, 'ds': ds, 'scale': scale, 'nu': nu, 'disp': disp, 'history': hist}
        return test
    except Exception:
        try: torch.cuda.empty_cache()
        except Exception: pass
        return float('inf')


def main():
    base_grid = [{'lr': lr, 'weight_decay': wd} for lr in LRS for wd in [0.0, 1e-4]]
    baseline = sweep_baseline(lambda c: lambda s: baseline_one(c, s), base_grid, seeds=(0,1,2,3))
    # Same lr union and baseline's selected central regularization; only nu differs.
    wd = baseline['best_cfg']['weight_decay']
    idea_grid = [{'lr': lr, 'weight_decay': wd, 'nu_mult': nu} for lr in LRS for nu in NU_MULTS]
    tried = []
    for cfg in idea_grid:
        r = evaluate(lambda s, c=cfg: shifted_one(c, s), seeds=(0,1,2,3))
        tried.append({'cfg': cfg, 'mean': r['mean']})
    best_cfg = min(idea_grid, key=lambda c: next(q['mean'] for q in tried if q['cfg'] == c))
    idea = evaluate(lambda s: shifted_one(best_cfg, s), seeds=SEEDS)
    # Signature is measured on trained models: positive-shift run versus zero-shift run.
    sig_cfg = next(c for c in idea_grid if c['lr'] == best_cfg['lr'] and c['nu_mult'] == 0.15)
    a = shifted_one(sig_cfg, 0, collect=True); z = shifted_one({**sig_cfg, 'nu_mult': 0.0}, 0, collect=True)
    predicted = 1.0 + sig_cfg['lr'] * a[1]['nu']
    observed = (a[1]['disp'] / max(z[1]['disp'], 1e-12))
    signature = {'predicted': {'one_step_displacement_factor': predicted, 'nu': a[1]['nu'], 'curvature_scale': a[1]['scale']},
                 'observed': {'trained_positive_shift_disp': a[1]['disp'], 'trained_zero_shift_disp': z[1]['disp'], 'whole_run_ratio': observed},
                 'confirmed': bool(np.isfinite(observed) and abs(observed-predicted) / max(abs(predicted),1e-9) < 0.25)}
    rep = make_report('tabular', 'mlp_tiny', baseline, idea, signature)
    rep['idea_sweep'] = {'grid': tried, 'best_cfg': best_cfg, 'shared_lr_union': LRS}
    Path('bench_report.json').write_text(json.dumps(rep, indent=2)); print(json.dumps(rep, indent=2))

if __name__ == '__main__': main()