import json, math, sys, time
import numpy as np
import torch
from torch import nn

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

SEEDS = tuple(range(8))
LR_GRID = [0.0015, 0.003, 0.006]
EPOCHS = 15
BATCH = 128
WIDTH = 32
RADIUS = 2


def dense_attention(h, q, k, v):
    qq, kk, vv = q(h), k(h), v(h)
    scores = qq @ kk.transpose(-1, -2) / math.sqrt(qq.shape[-1])
    return torch.softmax(scores, dim=-1) @ vv


def gated_local_global(h, q, k, v, gate, radius=RADIUS):
    b, n, _ = h.shape
    qq, kk, vv = q(h), k(h), v(h)
    offsets = list(range(-radius, radius + 1))
    positions = torch.arange(n, device=h.device)
    dst = torch.stack([(positions + off).clamp(0, n - 1) for off in offsets], dim=1)
    valid = torch.stack([(positions + off >= 0) & (positions + off < n)
                         for off in offsets], dim=1)
    kn = kk[:, dst]
    vn = vv[:, dst]
    scores = (qq[:, :, None, :] * kn).sum(-1) / math.sqrt(qq.shape[-1])
    scores = scores.masked_fill(~valid[None, :, :], -torch.inf)
    local = (torch.softmax(scores, dim=-1)[..., None] * vn).sum(dim=2)
    phi_q = torch.nn.functional.elu(qq) + 1.0
    phi_k = torch.nn.functional.elu(kk) + 1.0
    stats = phi_k.transpose(1, 2) @ vv
    norm = phi_k.sum(dim=1)
    glob = (phi_q @ stats) / ((phi_q @ norm[:, :, None]).squeeze(-1)[:, :, None] + 1e-6)
    g = torch.sigmoid(gate(h))
    return g * local + (1.0 - g) * glob, g


class SharedTransformer(nn.Module):
    def __init__(self, win, out_dim, idea=False):
        super().__init__()
        self.idea = idea
        self.inp = nn.Linear(1, WIDTH)
        self.pos = nn.Parameter(torch.randn(1, win, WIDTH) * 0.02)
        self.q = nn.Linear(WIDTH, WIDTH, bias=False)
        self.k = nn.Linear(WIDTH, WIDTH, bias=False)
        self.v = nn.Linear(WIDTH, WIDTH, bias=False)
        self.proj = nn.Linear(WIDTH, WIDTH)
        self.norm = nn.LayerNorm(WIDTH)
        self.ff = nn.Sequential(nn.Linear(WIDTH, 64), nn.ReLU(), nn.Linear(64, WIDTH))
        self.head = nn.Linear(win * WIDTH, out_dim)
        if idea:
            self.gate = nn.Linear(WIDTH, 1)

    def forward(self, x, return_gate=False):
        h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
        if self.idea:
            att, g = gated_local_global(h, self.q, self.k, self.v, self.gate)
        else:
            att = dense_attention(h, self.q, self.k, self.v)
            g = None
        h = self.norm(h + self.proj(att))
        h = h + self.ff(h)
        out = self.head(h.reshape(h.shape[0], -1))
        return (out, g) if return_gate else out


def run_one(seed, lr, idea, collect=False):
    torch.manual_seed(seed)
    np.random.seed(seed)
    ds = get_dataset("sequence", seed, n_train=400, n_test=400)
    win = int(ds["xtr"].shape[1])
    model = SharedTransformer(win, int(ds["ytr"].shape[-1]), idea=idea)
    trained, metric, history = train_model(model, ds, epochs=EPOCHS, lr=lr,
                                           batch=BATCH, weight_decay=0.0, log=lambda *_: None)
    if trained is None:
        return float("nan") if not collect else {"metric": float("nan")}
    if not collect:
        return float(metric)
    device = next(trained.parameters()).device
    x = ds["xte"].to(device)
    with torch.no_grad():
        if idea:
            pred, gates = trained(x, return_gate=True)
            gate_mean = float(gates.mean().cpu())
            gate_std = float(gates.std().cpu())
        else:
            pred = trained(x)
            gate_mean = None
            gate_std = None
    # Measured forward latency on the trained model, synchronized when CUDA is used.
    trained.eval()
    for _ in range(3):
        with torch.no_grad(): trained(x[:128])
    if device.type == "cuda": torch.cuda.synchronize(device)
    t0 = time.perf_counter()
    for _ in range(10):
        with torch.no_grad(): trained(x[:128])
    if device.type == "cuda": torch.cuda.synchronize(device)
    latency_ms = (time.perf_counter() - t0) * 1000.0 / 10.0
    return {"metric": float(metric), "gate_mean": gate_mean, "gate_std": gate_std,
            "latency_ms": latency_ms, "n_nodes": win, "edges_per_node": 2 * RADIUS + 1}


def baseline_factory(cfg):
    return lambda seed: run_one(seed, float(cfg["lr"]), False)


def main():
    # Baseline sweep uses exactly the union of all idea learning rates.
    grid = [{"lr": lr} for lr in LR_GRID]
    base = sweep_baseline(baseline_factory, grid, seeds=SEEDS)
    best_lr = float(base["best_cfg"]["lr"])
    idea_cfgs = [{"lr": lr} for lr in LR_GRID]
    idea_runs = {str(c["lr"]): evaluate(lambda seed, lr=c["lr"]: run_one(seed, lr, True), seeds=SEEDS)
                 for c in idea_cfgs}
    best_key = min(idea_runs, key=lambda k: idea_runs[k]["mean"])
    idea = idea_runs[best_key]
    sig = [run_one(s, float(best_key), True, collect=True) for s in SEEDS]
    base_sig = [run_one(s, best_lr, False, collect=True) for s in SEEDS]
    signature = {
        "prediction": "bounded-degree local-global attention scales with O(E*r + N*r^2), while dense attention scales with O(N^2*r)",
        "observed": {
            "sequence_length": int(np.mean([z["n_nodes"] for z in sig])),
            "local_edges_per_node": int(np.mean([z["edges_per_node"] for z in sig])),
            "idea_latency_ms_mean": float(np.mean([z["latency_ms"] for z in sig])),
            "baseline_latency_ms_mean": float(np.mean([z["latency_ms"] for z in base_sig])),
            "idea_gate_mean": float(np.mean([z["gate_mean"] for z in sig])),
            "idea_gate_std_mean": float(np.mean([z["gate_std"] for z in sig]))
        },
        "confirmed": False
    }
    # At this fixed short sequence length, observed latency confirms the
    # mechanism only if the hybrid is faster; otherwise report honestly.
    signature["confirmed"] = signature["observed"]["idea_latency_ms_mean"] < signature["observed"]["baseline_latency_ms_mean"]
    report = make_report("sequence", "transformer_tiny", base, idea, signature)
    report["idea_sweep"] = {k: v for k, v in idea_runs.items()}
    report["protocol"] = {"seeds": list(SEEDS), "epochs": EPOCHS, "batch": BATCH,
                           "lr_union": LR_GRID, "best_idea_lr": float(best_key)}
    with open("bench_report.json", "w") as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()