import sys, json
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 train_model, sweep_baseline, evaluate, make_report
from graph_topology_track import get_dataset

SEEDS = tuple(range(8))
TUNE_SEEDS = (0, 1, 2, 3)
LRS = [0.003, 0.01, 0.03]
GRID = np.linspace(0.0, 1.0, 9, dtype=np.float32)
M, STRIDE, N = 4, 2, 12


def gf2_rank(a):
    a = (np.asarray(a, dtype=np.uint8) & 1).copy()
    r = 0
    for c in range(a.shape[1]):
        piv = np.flatnonzero(a[r:, c])
        if len(piv) == 0: continue
        p = r + int(piv[0]); a[[r, p]] = a[[p, r]]
        for q in range(a.shape[0]):
            if q != r and a[q, c]: a[q] ^= a[r]
        r += 1
        if r == a.shape[0]: break
    return r


def betti(nv, edges):
    B = np.zeros((nv, len(edges)), dtype=np.uint8)
    for j, (u, v) in enumerate(edges): B[u, j] = B[v, j] = 1
    rank = gf2_rank(B)
    return int(nv - rank), int(len(edges) - rank)


def tokens(x, topo):
    out = []
    for row in x:
        A, h = row[:N*N].reshape(N, N), row[N*N:]
        edges = [(i, j) for i in range(N) for j in range(i + 1, N) if A[i, j] > .5]
        seq = []
        for st in range(0, len(GRID) - M, STRIDE):
            lo, hi = float(GRID[st]), float(GRID[st + M])
            ee = [(u, v) for u, v in edges if lo <= max(h[u], h[v]) <= hi]
            keep = {i for i, z in enumerate(h) if lo <= z <= hi}
            for u, v in ee: keep.update((u, v))
            vs = sorted(keep); rem = {v: i for i, v in enumerate(vs)}
            redges = [(rem[u], rem[v]) for u, v in ee]
            b0, b1 = betti(len(vs), redges)
            if topo:
                seq.append([b0, b1, np.log1p(len(vs)), np.log1p(len(redges))])
            else:
                deg = np.zeros(len(vs), dtype=np.float32)
                for u, v in redges: deg[u] += 1; deg[v] += 1
                seq.append([np.mean(deg) if len(deg) else 0., np.std(deg) if len(deg) else 0.,
                            np.log1p(len(vs)), np.log1p(len(redges))])
        out.append(seq)
    return np.asarray(out, dtype=np.float32)


class TokenTransformer(nn.Module):
    def __init__(self, n_tokens, width=32):
        super().__init__()
        self.proj = nn.Linear(4, width)
        self.pos = nn.Parameter(torch.zeros(1, n_tokens, width))
        nn.init.normal_(self.pos, std=.02)
        layer = nn.TransformerEncoderLayer(width, 4, 2 * width, batch_first=True, dropout=0.)
        self.enc = nn.TransformerEncoder(layer, 1)
        self.head = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, 2))
    def forward(self, x):
        z = self.proj(x) + self.pos[:, :x.shape[1]]
        return self.head(self.enc(z).mean(1))


def make_ds(raw, topo):
    return {"xtr": torch.tensor(tokens(raw["xtr"], topo)), "ytr": torch.tensor(raw["ytr"], dtype=torch.long),
            "xte": torch.tensor(tokens(raw["xte"], topo)), "yte": torch.tensor(raw["yte"], dtype=torch.long),
            "task": "classification", "metric": "err"}


def run_one(raw, topo, lr, seed):
    torch.manual_seed(seed + (10000 if topo else 0)); np.random.seed(seed)
    ds = make_ds(raw, topo)
    net, metric, _ = train_model(TokenTransformer(ds["xtr"].shape[1]), ds, epochs=18, lr=float(lr), batch=128, log=lambda *_: None)
    return float(metric), net, ds


def factory(topo, lr):
    def fn(seed):
        raw = get_dataset(int(seed), 400, 200)
        return run_one(raw, topo, lr, int(seed))[0]
    return fn


def main():
    grid = [{"lr": lr} for lr in LRS]
    baseline = sweep_baseline(lambda cfg: factory(False, cfg["lr"]), grid, seeds=TUNE_SEEDS)
    idea_trials = [{"cfg": cfg, "mean": evaluate(factory(True, cfg["lr"]), seeds=TUNE_SEEDS)["mean"]} for cfg in grid]
    idea_cfg = min(idea_trials, key=lambda z: z["mean"])["cfg"]
    idea_lr = float(idea_cfg["lr"])
    idea_full = evaluate(factory(True, idea_lr), seeds=SEEDS)
    pred_acc, b1_corr = [], []
    for seed in SEEDS:
        raw = get_dataset(seed, 400, 200)
        _, net, ds = run_one(raw, True, idea_lr, seed)
        dev = next(net.parameters()).device
        net.eval()
        with torch.no_grad(): pred = net(ds["xte"].to(dev)).argmax(1).cpu().numpy()
        observed_b1 = []
        for row in raw["xte"]:
            A = row[:N*N].reshape(N, N)
            edges = [(i, j) for i in range(N) for j in range(i+1, N) if A[i,j] > .5]
            observed_b1.append(betti(N, edges)[1])
        observed_b1 = np.asarray(observed_b1)
        pred_acc.append(float(np.mean(pred == raw["yte"])))
        b1_corr.append(float(np.corrcoef(pred, observed_b1)[0, 1]))
    signature = {"predicted": {"class0_betti1": 1.0, "class1_betti1": 2.0},
                 "observed": {"idea_accuracy_mean": float(np.mean(pred_acc)),
                              "prediction_observed_betti1_correlation_mean": float(np.mean(b1_corr))},
                 "confirmed": bool(np.mean(pred_acc) >= .75 and np.mean(b1_corr) > .5),
                 "note": "Measured from eight independently trained idea models on held-out graphs."}
    base_block = {"best_cfg": baseline["best_cfg"], "sweep": baseline["sweep"], "full": baseline["full"]}
    idea = {"best_cfg": idea_cfg, "sweep": idea_trials, "mean": idea_full["mean"],
            "std": idea_full["std"], "per_seed": idea_full["per_seed"], "n": idea_full["n"]}
    report = make_report("cycle_topology_graph", "token_transformer", base_block, idea, signature)
    report["custom_track"] = {"name": "cycle_topology_graph", "file": "graph_topology_track.py", "domain": "graph_topology"}
    report["protocol_note"] = "Eight paired seeds; baseline and idea share the same Transformer architecture and lr union; only token construction differs."
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()