"""Node2Vec model for the edge_ML_expected_ge5 graph.

Default run is a smoke check (builds the model, runs a few optimizer steps).
Pass --epochs N to train, which writes embeddings to --out.
"""

import argparse
import time

import torch
from torch_geometric.data import Data
from torch_geometric.data.data import DataEdgeAttr, DataTensorAttr
from torch_geometric.data.storage import BaseStorage, EdgeStorage, GlobalStorage
from torch_geometric.nn import Node2Vec

from paths import EMB_PATH, GRAPH_PATH

GRAPH = GRAPH_PATH


def load_graph(path: str = GRAPH) -> Data:
    torch.serialization.add_safe_globals(
        [Data, DataEdgeAttr, DataTensorAttr, BaseStorage, EdgeStorage, GlobalStorage]
    )
    return torch.load(path, weights_only=True)


def build_model(data: Data, args: argparse.Namespace, device: torch.device) -> Node2Vec:
    return Node2Vec(
        data.edge_index,
        embedding_dim=args.embedding_dim,
        walk_length=args.walk_length,
        context_size=args.context_size,
        walks_per_node=args.walks_per_node,
        num_negative_samples=args.num_negative_samples,
        p=args.p,
        q=args.q,
        num_nodes=data.num_nodes,
        sparse=True,  # pairs with SparseAdam; the embedding table is the only param
    ).to(device)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--embedding-dim", type=int, default=128)
    ap.add_argument("--walk-length", type=int, default=20)
    ap.add_argument("--context-size", type=int, default=10)
    ap.add_argument("--walks-per-node", type=int, default=10)
    ap.add_argument("--num-negative-samples", type=int, default=1)
    ap.add_argument("--p", type=float, default=1.0, help="return parameter")
    ap.add_argument("--q", type=float, default=1.0, help="in-out parameter")
    ap.add_argument("--batch-size", type=int, default=128)
    ap.add_argument("--lr", type=float, default=0.01)
    ap.add_argument("--num-workers", type=int, default=4)
    ap.add_argument("--epochs", type=int, default=0, help="0 = smoke check only")
    ap.add_argument("--steps", type=int, default=5, help="steps for the smoke check")
    ap.add_argument("--out", default=EMB_PATH)
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    data = load_graph()
    model = build_model(data, args, device)

    print(f"graph      : {data.num_nodes:,} nodes, {data.edge_index.size(1) // 2:,} undirected edges")
    print(f"device     : {device}")
    print(f"model      : {model}")
    print(f"parameters : {sum(p.numel() for p in model.parameters()):,} "
          f"({data.num_nodes:,} x {args.embedding_dim})")
    print(f"walks      : length={args.walk_length} context={args.context_size} "
          f"per_node={args.walks_per_node} p={args.p} q={args.q}")

    loader = model.loader(batch_size=args.batch_size, shuffle=True,
                          num_workers=args.num_workers)
    optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=args.lr)
    print(f"loader     : {len(loader):,} batches/epoch of {args.batch_size} seed nodes")

    def run_epoch(max_steps: int | None = None) -> float:
        model.train()
        total, n = 0.0, 0
        for i, (pos_rw, neg_rw) in enumerate(loader):
            optimizer.zero_grad()
            loss = model.loss(pos_rw.to(device), neg_rw.to(device))
            loss.backward()
            optimizer.step()
            total, n = total + loss.item(), n + 1
            if max_steps is not None and i + 1 >= max_steps:
                break
        return total / max(n, 1)

    if args.epochs == 0:
        t0 = time.perf_counter()
        loss = run_epoch(max_steps=args.steps)
        print(f"\nsmoke check: {args.steps} steps, mean loss {loss:.4f}, "
              f"{time.perf_counter() - t0:.1f}s")
        z = model()
        print(f"embeddings : {tuple(z.shape)} {z.dtype} on {z.device}")
        print("model built and training step verified; pass --epochs N to train")
        return

    for epoch in range(1, args.epochs + 1):
        t0 = time.perf_counter()
        loss = run_epoch()
        print(f"epoch {epoch:>3}/{args.epochs}  loss {loss:.4f}  "
              f"{time.perf_counter() - t0:.1f}s")

    model.eval()
    with torch.no_grad():
        z = model().cpu()
    torch.save({"embedding": z, "node_id": data.node_id,
                "args": vars(args)}, args.out)
    print(f"saved embeddings {tuple(z.shape)} -> {args.out}")


if __name__ == "__main__":
    main()