import sys, json, random
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import get_dataset, sweep_baseline, evaluate, make_report

NMOD = 4
EPOCHS = 10
BATCH = 128
SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))


def graph_laplacian():
    A = np.array([[0,.8,.2,.1],[.8,0,.5,.15],[.2,.5,0,.7],[.1,.15,.7,0]], dtype=float)
    return A, np.diag(A.sum(1)) - A


def grounded_gap(L, pins, strength):
    p = np.zeros(len(L)); p[list(pins)] = strength
    return float(np.linalg.eigvalsh(L + np.diag(p))[0])


def greedy_pins(L, m=2, strength=.5):
    chosen = []
    for _ in range(m):
        candidates = [(grounded_gap(L, chosen + [i], strength), i)
                      for i in range(len(L)) if i not in chosen]
        chosen.append(max(candidates, key=lambda z: (z[0], -z[1]))[1])
    return chosen, grounded_gap(L, chosen, strength)


class Module(nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = nn.GRU(3, 32, batch_first=True)
        self.head = nn.Linear(32, 1)

    def forward(self, x):
        _, h = self.rnn(x.view(x.shape[0], 8, 3))
        return self.head(h[-1]), h[-1]


def train_system(seed, lr, consensus, pin_strength=0.0, return_signature=False):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        if device == "cuda":
            try:
                torch.cuda.set_device(0)
                torch.backends.cudnn.enabled = False
            except Exception:
                device = "cpu"
    except Exception:
        device = "cpu"
    ds = get_dataset("dynamics", seed, n_train=400, n_test=200)
    xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device).reshape(-1, 1)
    xte, yte = ds["xte"].to(device), ds["yte"].to(device).reshape(-1, 1)
    nets = nn.ModuleList([Module() for _ in range(NMOD)]).to(device)
    opt = torch.optim.Adam(nets.parameters(), lr=lr)
    _, L = graph_laplacian()
    pins, gap = greedy_pins(L, 2, pin_strength if pin_strength else .5)
    pinmask = torch.zeros(NMOD, device=device)
    if pin_strength > 0: pinmask[pins] = 1.
    history = []
    for ep in range(EPOCHS):
        perm = torch.randperm(len(xtr), device=device)
        for start in range(0, len(xtr), BATCH):
            idx = perm[start:start+BATCH]
            preds, hs = [], []
            for net in nets:
                q, h = net(xtr[idx]); preds.append(q); hs.append(h)
            pred = torch.stack(preds)
            h = torch.stack(hs)
            loss = ((pred - ytr[idx].unsqueeze(0)) ** 2).mean()
            # L2 graph-consensus regularization, plus stronger corrective
            # teacher/anchor loss on spectral-greedy pinned modules.
            diff = h[:, None] - h[None, :]
            loss = loss + consensus * sum(L[i,j] * (h[i]-h[j]).pow(2).mean()
                                          for i in range(NMOD) for j in range(NMOD)) / 2
            if pin_strength > 0:
                mean = h.mean(0, keepdim=True).detach()
                loss = loss + pin_strength * ((h - mean).pow(2).mean(1) * pinmask[:, None]).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            ev = torch.stack([net(xte)[0] for net in nets])
            history.append(float(((ev-yte.unsqueeze(0))**2).mean()))
    with torch.no_grad():
        out = torch.stack([net(xte)[0] for net in nets])
        mse = float(((out-yte.unsqueeze(0))**2).mean())
        module_mean = out.mean(0, keepdim=True)
        disagreement = float(.5 * ((out-module_mean)**2).mean())
        # Measured trained-model behavior, not an analytical identity.
        train_h = torch.stack([net(xte)[1] for net in nets])
        measured_dis = float(.5 * ((train_h-train_h.mean(0,keepdim=True))**2).mean())
    if return_signature:
        return mse, {"pins": pins, "lambda_min_grounded": gap,
                      "test_representation_disagreement": measured_dis,
                      "task_prediction_disagreement": disagreement}
    return mse


def make_fn(cfg):
    return lambda seed: train_system(seed, cfg["lr"], cfg["consensus"], 0.0)


def main():
    # Search-space parity: every idea lr/consensus pair is also evaluated by
    # the baseline sweep. The baseline method's central consensus knob is swept.
    grid = [{"lr": lr, "consensus": c} for lr in [1e-3, 3e-3, 1e-2]
            for c in [0.0, 0.01, 0.05]]
    base = sweep_baseline(make_fn, grid, seeds=SWEEP_SEEDS)
    best = base["best_cfg"]
    idea_cfgs = [best,
                 {"lr": best["lr"], "consensus": best["consensus"] + .01},
                 {"lr": best["lr"], "consensus": max(0.0, best["consensus"] - .01)}]
    # Keep the idea's three settings on the same union grid where possible.
    all_grid = list(grid)
    for cfg in idea_cfgs:
        if cfg not in all_grid: all_grid.append(cfg)
    base = sweep_baseline(make_fn, all_grid, seeds=SWEEP_SEEDS)
    idea_trials = []
    for cfg in idea_cfgs:
        r = evaluate(lambda s, cfg=cfg: train_system(s, cfg["lr"], cfg["consensus"], .5), seeds=SEEDS)
        idea_trials.append({"cfg": cfg, "result": r})
    chosen = min(idea_trials, key=lambda z: z["result"]["mean"])
    idea = chosen["result"]
    best_idea_cfg = chosen["cfg"]
    sig_mse, sig = train_system(0, best_idea_cfg["lr"], best_idea_cfg["consensus"], .5, True)
    _, L = graph_laplacian()
    random_gap = float(np.mean([grounded_gap(L, p, .5) for p in ([0,1],[1,2],[2,3],[0,3])]))
    sig.update({"baseline_random_gap": random_gap,
                "observed_vs_predicted": "task-independent disagreement is measured on trained representations; no exact decay slope is identifiable from final snapshots",
                "confirmed": False})
    report = make_report("dynamics", "rnn_small", base, idea, sig)
    report["idea_sweep"] = idea_trials
    report["baseline_sweep_union"] = all_grid
    with open("bench_report.json", "w") as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()