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

SEEDS = tuple(range(8))
# Shared union: every idea learning rate is also evaluated for baseline.
LR_GRID = [0.0015, 0.003, 0.006]
EPOCHS = 15
NTRAIN, NTEST = 1200, 400
LAMBDA_D = 1e-3
TAU = 0.8


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 baseline_fn(cfg):
    def run(seed):
        seed_all(seed)
        ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
        net = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
        _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=cfg["lr"], batch=128)
        return float(metric)
    return run


class ActuatorRNN(nn.Module):
    """Same rnn_small GRU/head, with d(u)=alpha*tanh(u/tau) before it."""
    def __init__(self, base):
        super().__init__()
        self.rnn = base.rnn
        self.head = base.head
        self.alpha = nn.Parameter(torch.zeros(1))
        self.tau = TAU

    def disturbance(self, x):
        z = x.view(x.shape[0], -1, 3).clone()
        u = z[:, :, 2]
        d = self.alpha * torch.tanh(u / self.tau)
        z[:, :, 2] = u + d
        return z.reshape(x.shape)

    def forward(self, x):
        seq = self.disturbance(x).view(x.shape[0], -1, 3)
        try:
            _, h = self.rnn(seq)
        except RuntimeError:
            old = torch.backends.cudnn.enabled
            torch.backends.cudnn.enabled = False
            try:
                _, h = self.rnn(seq)
            finally:
                torch.backends.cudnn.enabled = old
        return self.head(h[-1])


def train_idea(seed, cfg, return_model=False):
    seed_all(seed)
    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
    base = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
    net = ActuatorRNN(base)
    # This is a custom loop because the intervention explicitly changes the loss.
    errors = []
    for device in (["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"]):
        try:
            net = net.to(device)
            xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device)
            opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"])
            for _ in range(EPOCHS):
                net.train()
                perm = torch.randperm(len(xtr), device=device)
                for i in range(0, len(xtr), 128):
                    ix = perm[i:i+128]
                    pred = net(xtr[ix])
                    d = net.disturbance(xtr[ix]).view(len(ix), -1, 3)[:, :, 2] - xtr[ix].view(len(ix), -1, 3)[:, :, 2]
                    loss = ((pred - ytr[ix]) ** 2).mean() + LAMBDA_D * (d ** 2).mean()
                    opt.zero_grad(); loss.backward(); opt.step()
            net.eval()
            with torch.no_grad():
                metric = float(((net(ds["xte"].to(device)) - ds["yte"].to(device)) ** 2).mean())
            return (metric, net, ds) if return_model else metric
        except RuntimeError as e:
            errors.append(str(e))
            if device == "cpu":
                raise
            # Recreate clean CPU tensors/model after any CUDA failure.
            base = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
            net = ActuatorRNN(base)
    raise RuntimeError("training failed: " + repr(errors))


def idea_fn(cfg):
    return lambda seed: train_idea(seed, cfg)


def mechanism_signature():
    # Measurements are from trained benchmark models, not an analytic toy graph.
    vals = []
    corrs = []
    for seed in SEEDS:
        metric, net, ds = train_idea(seed, {"lr": 0.003}, True)
        with torch.no_grad():
            dev = next(net.parameters()).device
            x = ds["xte"].to(dev)
            u = x.view(len(x), -1, 3)[:, :, 2]
            d = net.alpha.detach().cpu().item() * torch.tanh(u / TAU)
            vals.append(abs(net.alpha.detach().cpu().item()) / TAU)
            # In a null-distortion track, residual/action correlation should be small.
            pred = net(x).detach().cpu().numpy().reshape(-1)
            residual = ds["yte"].numpy() - pred
            signal = d.detach().cpu().numpy().mean(axis=1)
            corrs.append(float(np.corrcoef(residual, signal)[0, 1]) if np.std(signal) > 1e-10 else 0.0)
    observed = float(np.mean(vals))
    return {"predicted_max_jacobian": 0.0, "observed_mean_abs_alpha_over_tau": observed,
            "observed_residual_disturbance_corr_abs": float(np.mean(np.abs(corrs))),
            "n_trained_models": 8, "confirmed": bool(observed < 0.05),
            "interpretation": "The matched bench has no latent actuator error; the simplicity prior should identify zero disturbance."}


def main():
    base = sweep_baseline(baseline_fn, [{"lr": x} for x in LR_GRID])
    # Same 3-point grid for the idea; best selected on the same four sweep seeds.
    idea_trials = []
    for cfg in [{"lr": x} for x in LR_GRID]:
        r = evaluate(idea_fn(cfg), 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(idea_fn(best_cfg), seeds=SEEDS)
    report = make_report("dynamics", "rnn_small", base, idea,
                         {"idea_lr_sweep": idea_trials, "chosen_idea_cfg": best_cfg,
                          **mechanism_signature()})
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()