import json
import random
from pathlib import Path
import numpy as np
import torch
from torch import nn
import sys
sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report

EPOCHS = 10
BATCH = 128
WEIGHT_DECAY = 0.0


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 sign_transform(x):
    # Pendulum equations are equivariant under (theta, omega, u) -> -(theta, omega, u).
    return -x


def output_transform(y):
    return -y


def equivariance_penalty(net, x):
    pred = net(x)
    pred_transformed = net(sign_transform(x))
    return ((pred_transformed - output_transform(pred)) ** 2).mean()


def train_idea(net, ds, epochs, lr, lam):
    # Same Adam/batch/epochs as bench.train_model; only the proposed loss is added.
    devices = [("cuda", False), ("cuda", True), ("cpu", False)] if torch.cuda.is_available() else [("cpu", False)]
    errors = []
    for devname, no_cudnn in devices:
        try:
            device = torch.device(devname)
            old_cudnn = torch.backends.cudnn.enabled
            if no_cudnn:
                torch.backends.cudnn.enabled = False
            net = net.to(device)
            opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=WEIGHT_DECAY)
            xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device)
            for _ in range(epochs):
                net.train()
                perm = torch.randperm(len(xtr), device=device)
                for i in range(0, len(xtr), BATCH):
                    idx = perm[i:i+BATCH]
                    pred = net(xtr[idx])
                    loss = ((pred - ytr[idx]) ** 2).mean() + lam * equivariance_penalty(net, xtr[idx])
                    opt.zero_grad(set_to_none=True)
                    loss.backward()
                    opt.step()
            net.eval()
            with torch.no_grad():
                metric = float(((net(ds["xte"].to(device)) - ds["yte"].to(device)) ** 2).mean().cpu())
            torch.backends.cudnn.enabled = old_cudnn
            return net, metric
        except RuntimeError as exc:
            errors.append(str(exc)[:160])
            if devname == "cuda":
                continue
            raise
    raise RuntimeError("training failed: " + " | ".join(errors))


def fit_baseline(cfg, seed):
    seed_all(seed)
    ds = get_dataset("dynamics", seed, n_train=400, n_test=400)
    net = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
    _, metric, _ = train_model(net, ds, epochs=cfg["epochs"], lr=cfg["lr"], batch=BATCH, weight_decay=WEIGHT_DECAY, log=lambda *_: None)
    return metric


def fit_idea(cfg, seed, capture=False):
    seed_all(seed)
    ds = get_dataset("dynamics", seed, n_train=400, n_test=400)
    net = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
    trained, metric = train_idea(net, ds, cfg["epochs"], cfg["lr"], cfg["lam"])
    if not capture:
        return metric
    device = next(trained.parameters()).device
    x = ds["xte"].to(device)
    with torch.no_grad():
        p = trained(x)
        pt = trained(-x)
        eq = float(((pt + p) ** 2).mean().cpu())
        # Behavioural signature: compare observed transformed prediction to the
        # transformed prediction expected from the trained model's original output.
        pred_scale = float(p.abs().mean().cpu())
    return metric, {"mean_prediction_abs": pred_scale, "trained_model_c2_penalty": eq}


def math_check():
    x = torch.randn(64, 24)
    y = torch.randn(64, 1)
    return {"c2_input_composition_max_abs": float((sign_transform(sign_transform(x)) - x).abs().max()),
            "c2_output_composition_max_abs": float((output_transform(output_transform(y)) - y).abs().max()),
            "exact_transform_identity": bool(torch.equal(sign_transform(sign_transform(x)), x))}


def main():
    # Union parity: every lr and lambda considered for the idea is also run by
    # baseline; baseline lambda is the standard no-penalty value 0.
    lrs = [1e-3, 3e-3, 1e-2]
    base_grid = [{"lr": lr, "epochs": EPOCHS, "lam": 0.0} for lr in lrs]
    idea_grid = [{"lr": lr, "epochs": EPOCHS, "lam": lam} for lr in lrs for lam in [0.05, 0.2, 1.0]]
    def base_fn(cfg):
        return lambda seed: fit_baseline(cfg, seed)
    base = sweep_baseline(base_fn, base_grid)
    # Evaluate all idea settings on sweep seeds for selection, then full 8 seeds.
    idea_trials = []
    for cfg in idea_grid:
        r = evaluate(lambda seed, c=cfg: fit_idea(c, seed), seeds=(0, 1, 2, 3))
        idea_trials.append({"cfg": cfg, "mean": r["mean"]})
    best_cfg = min(idea_trials, key=lambda z: z["mean"])["cfg"]
    idea = evaluate(lambda seed: fit_idea(best_cfg, seed), seeds=tuple(range(8)))
    sig_metric, sig = fit_idea(best_cfg, 0, capture=True)
    sig.update({"prediction_metric_seed0": sig_metric,
                "expected_c2_penalty_direction": "lower is more equivariant",
                "confirmed": sig["trained_model_c2_penalty"] < 1e-3})
    report = make_report("dynamics", "rnn_small", {"best_cfg": base["best_cfg"], "sweep": base["sweep"], "full": base["full"]}, idea, {"mechanism_signature": sig, "math_check": math_check(), "idea_sweep": idea_trials, "structural_match": "controlled pendulum rollout has sign-equivariant local dynamics"})
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()