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

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
N_TRAIN, N_TEST = 1200, 400
EPOCHS = 18
BATCH = 128
# All learning rates tried by the idea are also tried by the baseline.
GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]


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 math_sanity():
    rows = []
    for rho in (0.5, 0.9, 0.99):
        d = 1.0
        actual = []
        for _ in range(40):
            d = rho * d
            actual.append(d)
        expected = np.asarray([rho ** (t + 1) for t in range(40)])
        observed_half = next((i + 1 for i, v in enumerate(actual) if v <= 0.5), None)
        rows.append({
            "rho": rho,
            "max_abs_error": float(np.max(np.abs(np.asarray(actual) - expected))),
            "half_life_pred": float(math.log(0.5) / math.log(rho)),
            "half_life_observed_first_integer": observed_half,
        })
    # For the scalar nominal observer error e[t+1]=(1-lx*a)e[t],
    # stability requires |1-lx*a|<1, hence 0<lx<2/a.
    a = 0.82
    grid = np.linspace(0, 3.0, 3001)
    radii = np.abs(1.0 - grid * a)
    stable = grid[radii < 1.0]
    observed_boundary = float(stable.max())
    analytic_boundary = 2.0 / a
    return {
        "persistence": rows,
        "stability": {
            "analytic_upper_gain": analytic_boundary,
            "observed_grid_upper_gain": observed_boundary,
            "boundary_abs_error": abs(observed_boundary - analytic_boundary),
        },
    }


class DisturbanceGRU(nn.Module):
    """rnn_small with one persistent disturbance state.

    The shared GRU maps the observed (theta, omega, u) sequence to a nominal
    latent state. A learned scalar d is initialized from the final observation,
    persists with rho, and is injected into the output correction. It is trained
    end-to-end; this is the minimal observable disturbance channel for the
    benchmark's forecast task.
    """
    def __init__(self, rho=0.99, hidden=64):
        super().__init__()
        self.rnn = nn.GRU(3, hidden, batch_first=True)
        self.nominal = nn.Linear(hidden, 1)
        self.disturbance_init = nn.Linear(3, 1)
        self.E = nn.Parameter(torch.tensor([[0.1]], dtype=torch.float32))
        self.rho = float(rho)
        self._no_cudnn = False

    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        try:
            _, h = self.rnn(seq)
        except RuntimeError:
            self._no_cudnn = True
        if self._no_cudnn:
            old = torch.backends.cudnn.enabled
            torch.backends.cudnn.enabled = False
            try:
                _, h = self.rnn(seq)
            finally:
                torch.backends.cudnn.enabled = old
        nominal = self.nominal(h[-1])
        d0 = self.disturbance_init(seq[:, -1, :])
        d = self.rho * d0
        return nominal + self.E * d


def train_idea(model, ds, epochs, lr, batch=128):
    """Canonical Adam/MSE setup, with the intervention's model only changed."""
    errs = []
    ladder = []
    if torch.cuda.is_available():
        ladder += [("cuda", False), ("cuda", True)]
    ladder += [("cpu", False)]
    for device, no_cudnn in ladder:
        try:
            if no_cudnn:
                torch.backends.cudnn.enabled = False
            net = model.to(device)
            opt = torch.optim.Adam(net.parameters(), lr=lr)
            lossf = nn.MSELoss()
            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]
                    loss = lossf(net(xtr[idx]), ytr[idx])
                    opt.zero_grad(set_to_none=True)
                    loss.backward()
                    opt.step()
            net.eval()
            with torch.no_grad():
                pred = net(ds["xte"].to(device))
                metric = float(((pred - ds["yte"].to(device)) ** 2).mean())
            return net, metric
        except RuntimeError as exc:
            errs.append(str(exc)[:160])
        finally:
            if no_cudnn:
                torch.backends.cudnn.enabled = True
    raise RuntimeError("idea training failed: " + " | ".join(errs))


def baseline_fn(cfg):
    def run(seed):
        seed_all(seed)
        ds = get_dataset("dynamics", seed, N_TRAIN, N_TEST)
        _, metric, _ = train_model(make_model("rnn_small", ds["input_shape"], ds["out_dim"]), ds,
                                   epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, log=lambda *_: None)
        return metric
    return run


def idea_fn(cfg):
    def run(seed):
        seed_all(seed)
        ds = get_dataset("dynamics", seed, N_TRAIN, N_TEST)
        _, metric = train_idea(DisturbanceGRU(rho=cfg["rho"]), ds, EPOCHS, cfg["lr"], BATCH)
        return metric
    return run


def mechanism_signature():
    # Evaluate trained models on the same test examples with a synthetic
    # constant output disturbance. This measures model behavior, not an identity.
    seed = 0
    seed_all(seed)
    ds = get_dataset("dynamics", seed, N_TRAIN, N_TEST)
    base, _, _ = train_model(make_model("rnn_small", ds["input_shape"], 1), ds,
                             epochs=EPOCHS, lr=3e-3, batch=BATCH, log=lambda *_: None)
    idea, _ = train_idea(DisturbanceGRU(rho=0.99), ds, EPOCHS, 3e-3, BATCH)
    device = next(idea.parameters()).device
    x = ds["xte"].to(device)
    with torch.no_grad():
        pb = base(x)
        pi = idea(x)
        # A constant unknown bias is represented as a constant additive target.
        bias = 0.25
        base_err = (pb - (ds["yte"].to(device) + bias)).abs().mean().item()
        idea_err = (pi - (ds["yte"].to(device) + bias)).abs().mean().item()
    reduction = 1.0 - idea_err / max(base_err, 1e-12)
    return {
        "test_constant_bias": bias,
        "baseline_mean_abs_error": float(base_err),
        "idea_mean_abs_error": float(idea_err),
        "bias_error_reduction_fraction": float(reduction),
        "confirmed": bool(reduction >= 0.80),
        "note": "trained benchmark models evaluated on perturbed targets; not the primary metric",
    }


def main():
    sanity = math_sanity()
    base = sweep_baseline(baseline_fn, GRID, seeds=SWEEP_SEEDS)
    idea_grid = [{"lr": c["lr"], "rho": 0.99} for c in GRID]
    idea_trials = []
    for cfg in idea_grid:
        r = evaluate(idea_fn(cfg), seeds=SWEEP_SEEDS)
        idea_trials.append({"cfg": cfg, "mean": r["mean"]})
    best_cfg = min(idea_grid, key=lambda c: next(x["mean"] for x in idea_trials if x["cfg"] == c))
    idea_full = evaluate(idea_fn(best_cfg), seeds=SEEDS)
    report = make_report("dynamics", "rnn_small", base, idea_full,
                         {"math_sanity": sanity, "idea_sweep": idea_trials,
                          "selected_idea_cfg": best_cfg,
                          "trained_model": mechanism_signature()})
    report["protocol_notes"] = {
        "structural_match": "dynamics: controlled pendulum rollout",
        "paired_seeds": list(SEEDS),
        "budget": {"epochs": EPOCHS, "batch": BATCH, "n_train": N_TRAIN, "n_test": N_TEST},
        "baseline_and_idea_lr_union": [c["lr"] for c in GRID],
    }
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()