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

SEEDS = tuple(range(8))
SWEEP_SEEDS = (0, 1, 2, 3)
LR_GRID = [1e-3, 3e-3, 1e-2]
EPOCHS = 6
BATCH = 128
M = 4
ALPHA = 0.10
LAMBDA = 0.002
BETA = 0.90


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 train_one(seed, lr, regularized, collect=False):
    seed_all(seed)
    ds = get_dataset("dynamics", seed, n_train=800, n_test=300)
    net = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        return _train(net, ds, lr, regularized, dev, collect)
    except Exception:
        if dev == "cuda":
            torch.cuda.empty_cache()
            return _train(net.cpu(), ds, lr, regularized, "cpu", collect)
        raise


def _train(net, ds, lr, regularized, dev, collect=False):
    net = net.to(dev)
    x, y = ds["xtr"].to(dev), ds["ytr"].to(dev)
    xt, yt = ds["xte"].to(dev), ds["yte"].to(dev)
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    lossf = nn.MSELoss()
    queue = []
    eps = 0.0
    violations, ratios, grad_norms, losses = [], [], [], []
    for _ in range(EPOCHS):
        net.train()
        perm = torch.randperm(len(x), device=dev)
        for i in range(0, len(x), BATCH):
            idx = perm[i:i + BATCH]
            pred = net(x[idx])
            task = lossf(pred, y[idx])
            grads = torch.autograd.grad(task, tuple(net.parameters()),
                                        create_graph=regularized, retain_graph=True)
            flat = torch.cat([g.reshape(-1) for g in grads])
            # Gradient certificate q_k and V_k = ||q_k||^2/2.
            v_now = 0.5 * (flat * flat).sum()
            penalty = torch.zeros((), device=dev)
            if regularized and len(queue) >= M:
                v_old = queue[-M]
                delta = v_now - v_old + ALPHA * v_old
                eps = BETA * eps + (1.0 - BETA) * float(delta.detach().abs().cpu())
                residual = delta - float(eps)
                # Normalization and clipping prevent certificate scale domination.
                penalty = LAMBDA * torch.relu(residual).clamp(max=10.0).pow(2) / (1.0 + v_old)
                violations.append(float((residual.detach() > 0).cpu()))
                ratios.append(float((v_now.detach() / (v_old + 1e-12)).cpu()))
            total = task + penalty if regularized else task
            opt.zero_grad(set_to_none=True)
            total.backward()
            opt.step()
            queue.append(v_now.detach())
            grad_norms.append(float(flat.detach().norm().cpu()))
            losses.append(float(task.detach().cpu()))
    net.eval()
    with torch.no_grad():
        metric = float(((net(xt) - yt) ** 2).mean().cpu())
    if not collect:
        return metric
    return {
        "metric": metric,
        "model": net,
        "signature": {
            "violation_rate": float(np.mean(violations)) if violations else 0.0,
            "mean_v_ratio": float(np.mean(ratios)) if ratios else float("nan"),
            "tail_grad_std": float(np.std(grad_norms[-100:])) if grad_norms else float("nan"),
            "tail_loss_std": float(np.std(losses[-100:])) if losses else float("nan"),
        },
    }


def metric_fn(cfg, regularized):
    lr = float(cfg["lr"])
    return lambda seed: train_one(seed, lr, regularized, False)


def main():
    grid = [{"lr": lr} for lr in LR_GRID]
    base = sweep_baseline(lambda cfg: metric_fn(cfg, False), grid, seeds=SWEEP_SEEDS)
    # The same union of lrs is used for the idea-side three-configuration sweep.
    idea_trials = []
    for cfg in grid:
        r = evaluate(metric_fn(cfg, True), seeds=SWEEP_SEEDS)
        idea_trials.append({"cfg": cfg, "mean": r["mean"]})
    best_idea_cfg = min(idea_trials, key=lambda z: z["mean"])["cfg"]
    idea = evaluate(metric_fn(best_idea_cfg, True), seeds=SEEDS)
    report = make_report(
        "dynamics", "rnn_small", base, idea,
        extra={
            "prediction": "finite-horizon gradient energy should contract on average and unstable training should have positive residual violations",
            "trained_model_measurements": {
                "idea_cfg": best_idea_cfg,
                "idea_signature_per_seed": [train_one(s, best_idea_cfg["lr"], True, True)["signature"] for s in SEEDS],
            },
            "confirmed": False,
            "idea_sweep": idea_trials,
            "certificate": {"M": M, "alpha": ALPHA, "lambda": LAMBDA, "beta": BETA},
        },
    )
    # Add a compact quantitative signature summary, measured from trained models.
    sigs = report["mechanism_signature"]["trained_model_measurements"]["idea_signature_per_seed"]
    report["mechanism_signature"]["observed_mean_v_ratio"] = float(np.nanmean([s["mean_v_ratio"] for s in sigs]))
    report["mechanism_signature"]["observed_violation_rate"] = float(np.mean([s["violation_rate"] for s in sigs]))
    ratio = report["mechanism_signature"]["observed_mean_v_ratio"]
    report["mechanism_signature"]["confirmed"] = bool(np.isfinite(ratio) and ratio < 1.0)
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()