import sys, json, math, time
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, train_model, sweep_baseline, make_report

SEEDS = tuple(range(8))
# Union is shared by baseline and idea; baseline sweep covers every idea LR.
GRID = [
    {"lr": 1e-3, "epochs": 12},
    {"lr": 3e-3, "epochs": 12},
    {"lr": 1e-2, "epochs": 12},
]
HIDDEN = 32
SWEEPS = 3
CLUSTERS = 4
ETA = 0.55
R_BINS = 4
A_MAX = 8.0

class IterativeGraphNet(nn.Module):
    """Small recurrent message-passing solver over the 8-step pendulum window."""
    def __init__(self, mode="full", eta=ETA, clusters=CLUSTERS, sweeps=SWEEPS):
        super().__init__()
        self.mode, self.eta = mode, eta
        self.clusters, self.sweeps = clusters, sweeps
        self.inp = nn.Linear(3, HIDDEN)
        self.msg = nn.Linear(HIDDEN, HIDDEN, bias=False)
        self.self_proj = nn.Linear(HIDDEN, HIDDEN)
        self.head = nn.Linear(HIDDEN, 1)

    def one_update(self, h, x, ids):
        # Gather all neighbors from the pre-update tensor, then scatter together.
        old = h
        left = torch.roll(old, 1, dims=1)
        right = torch.roll(old, -1, dims=1)
        neigh = 0.5 * (left + right)
        z = torch.tanh(self.self_proj(old) + self.msg(neigh) + self.inp(x))
        out = old.clone()
        out[:, ids, :] = (1.0 - self.eta) * old[:, ids, :] + self.eta * z[:, ids, :]
        return out

    def histograms(self, h):
        # Observable local mismatch: disagreement with ring-neighbor states.
        d = 0.5 * (torch.roll(h, 1, 1) - h).norm(dim=-1) + \
            0.5 * (torch.roll(h, -1, 1) - h).norm(dim=-1)
        omega = torch.clamp((d / (d.detach().mean(dim=1, keepdim=True) + 1e-6) * 4).long(), 0, int(A_MAX))
        q = torch.clamp((R_BINS * omega / A_MAX).long(), 0, R_BINS - 1)
        hs = []
        for a in range(self.clusters):
            ids = list(range(a * 8 // self.clusters, (a + 1) * 8 // self.clusters))
            hs.append(torch.stack([(q[:, ids] == r).float().mean(1) for r in range(R_BINS)], 1))
        return torch.stack(hs, 1), d

    def forward(self, x):
        # x is [batch, 24], interpreted as 8 nodes with 3 features.
        x = x.view(x.shape[0], 8, 3)
        h = torch.tanh(self.inp(x))
        if self.mode == "full":
            for _ in range(self.sweeps):
                h = self.one_update(h, x, list(range(8)))
        else:
            # Equal node-update budget: clusters*sweeps decisions.
            for _ in range(self.sweeps * self.clusters):
                hist, _ = self.histograms(h.detach())
                action = hist[:, :, -1].argmax(1)
                # A batch has one action per example; grouping preserves synchronous updates.
                nxt = h.clone()
                for a in range(self.clusters):
                    mask = action == a
                    if mask.any():
                        ids = list(range(a * 8 // self.clusters, (a + 1) * 8 // self.clusters))
                        nxt[mask] = self.one_update(h[mask], x[mask], ids)
                h = nxt
        return self.head(h[:, -1, :])

def train_idea(seed, cfg):
    torch.manual_seed(seed); np.random.seed(seed)
    ds = get_dataset("dynamics", seed, n_train=400, n_test=100)
    model = IterativeGraphNet("cluster")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        model = model.to(device)
        opt = torch.optim.Adam(model.parameters(), lr=cfg["lr"])
        lossf = nn.MSELoss(); x, y = ds["xtr"].to(device), ds["ytr"].to(device)
        for _ in range(cfg["epochs"]):
            model.train(); p = torch.randperm(len(x), device=device)
            for i in range(0, len(x), 128):
                ix = p[i:i+128]; loss = lossf(model(x[ix]), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
        model.eval()
        with torch.no_grad():
            metric = float(lossf(model(ds["xte"].to(device)), ds["yte"].to(device)))
        return metric, model.cpu(), ds
    except (RuntimeError, torch.cuda.CudaError):
        # Required robust fallback; recreate model to avoid partial CUDA state.
        model = IterativeGraphNet("cluster")
        opt = torch.optim.Adam(model.parameters(), lr=cfg["lr"])
        lossf = nn.MSELoss(); x, y = ds["xtr"], ds["ytr"]
        for _ in range(cfg["epochs"]):
            p = torch.randperm(len(x))
            for i in range(0, len(x), 128):
                ix = p[i:i+128]; loss = lossf(model(x[ix]), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): metric = float(lossf(model(ds["xte"]), ds["yte"]))
        return metric, model, ds

def baseline_factory(cfg):
    def run(seed):
        torch.manual_seed(seed); np.random.seed(seed)
        ds = get_dataset("dynamics", seed, n_train=400, n_test=100)
        model = IterativeGraphNet("full")
        _, metric, _ = train_model(model, ds, epochs=cfg["epochs"], lr=cfg["lr"], batch=128, log=lambda *_: None)
        return metric
    return run

def idea_eval(cfg):
    vals = []
    for s in SEEDS:
        v, _, _ = train_idea(s, cfg); vals.append(v)
    return {"mean": float(np.mean(vals)), "std": float(np.std(vals)), "per_seed": vals, "n": len(vals)}

def mechanism_signature(base_cfg, idea_cfg):
    # Re-test eta^2 ordering prediction on trained models, not a toy identity.
    b, bm, ds = train_idea(0, idea_cfg)
    bm.eval(); x = ds["xte"][:32].view(32, 8, 3)
    rows = []
    with torch.no_grad():
        for eta in (0.2, 0.4, 0.6, 0.8):
            bm.eta = eta
            h = torch.tanh(bm.inp(x)); old = h.clone()
            ids = [0, 1]
            # synchronous cluster update
            sync = bm.one_update(old, x, ids)
            # deliberately sequential ordering on the same trained update rule
            seq = old.clone()
            for i in ids: seq = bm.one_update(seq, x, [i])
            err = float((sync[:, ids] - seq[:, ids]).abs().max())
            rows.append({"eta": eta, "observed_linf": err, "eta2": eta * eta})
    ratios = [r["observed_linf"] / max(r["eta2"], 1e-9) for r in rows]
    # Prediction is quantitative scaling, assessed by normalized variation.
    confirmed = bool(max(ratios) < 1.5 * max(min(ratios), 1e-9))
    return {"prediction": "trained-model synchronous-vs-sequential discrepancy scales approximately eta^2", "rows": rows, "eta2_ratio_range": [float(min(ratios)), float(max(ratios))], "confirmed": confirmed}

def main():
    t0 = time.time()
    base = sweep_baseline(baseline_factory, GRID)
    idea_trials = [(cfg, idea_eval(cfg)) for cfg in GRID]
    idea_cfg, idea = min(idea_trials, key=lambda z: z[1]["mean"])
    rep = make_report("dynamics", "rnn_small", base, idea, {
        "track_structure": "controlled pendulum multi-step dynamics; recurrent state updates",
        "scheduler": "largest high-bin hidden-state disagreement histogram",
        "clusters": CLUSTERS, "histogram_bins": R_BINS,
        "signature": mechanism_signature(base["best_cfg"], base["best_cfg"])
    })
    rep["runtime_sec"] = time.time() - t0
    rep["protocol_notes"] = {"paired_seeds": list(SEEDS), "dataset_sizes": [400, 100], "idea_grid": [{"cfg": c, "mean": r["mean"]} for c, r in idea_trials], "selected_idea_cfg": idea_cfg, "baseline_grid_union": GRID}
    Path("bench_report.json").write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))

if __name__ == "__main__": main()