"""
Train the VGAE model on the pre-built PyG graph.

Run from the repository root:
    python -m src.gnn.trainer
"""

import os
import torch
import torch.nn.functional as F
from torch_geometric.utils import negative_sampling

from src.gnn.model import get_model


def train_gnn(pyg_data_path: str = "./storage_graph/pyg_data.pt", epochs: int = 100):
    print(f"Loading PyG data from {pyg_data_path}…")
    if not os.path.exists(pyg_data_path):
        print(f"File not found: {pyg_data_path}")
        return

    # Allowlist PyG classes for safe loading under PyTorch >= 2.6.
    try:
        from torch_geometric.data import Data
        from torch_geometric.data.data import DataEdgeAttr
        torch.serialization.add_safe_globals([Data, DataEdgeAttr])
    except (ImportError, AttributeError):
        pass

    checkpoint = torch.load(pyg_data_path, weights_only=False)
    data = checkpoint["pyg_data"]

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    data = data.to(device)

    in_channels = data.x.size(1)
    print(f"Initialising VGAE: in={in_channels}, hidden=256, out=128 — device={device}")

    model = get_model(in_channels=in_channels, hidden_channels=256, out_channels=128).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    print("Training with link-prediction objective (L_recon + L_KL)…")
    model.train()
    for epoch in range(epochs):
        optimizer.zero_grad()

        z = model.encode(data.x, data.edge_index)

        # Reconstruction loss on positive edges.
        pos_loss = -torch.log(model.decode(z, data.edge_index, sigmoid=True) + 1e-15).mean()

        # Reconstruction loss on sampled negative edges (|E⁻| = |E⁺|).
        neg_edge_index = negative_sampling(
            edge_index=data.edge_index,
            num_nodes=data.num_nodes,
            num_neg_samples=data.edge_index.size(1),
        )
        neg_loss = -torch.log(1 - model.decode(z, neg_edge_index, sigmoid=True) + 1e-15).mean()

        # KL regularisation term — requires encode() to have been called above.
        kl = model.kl_loss()

        loss = pos_loss + neg_loss + kl
        loss.backward()
        optimizer.step()

        if epoch % 10 == 0:
            print(f"Epoch {epoch:03d}: loss={loss.item():.4f}  "
                  f"(recon={pos_loss.item()+neg_loss.item():.4f}, kl={kl.item():.4f})")

    print("Training complete.")

    with torch.no_grad():
        model.eval()
        z_final = model.encode(data.x, data.edge_index).cpu()

    print("Saving model weights and structural embeddings…")
    checkpoint["structural_embeddings"] = z_final
    torch.save(checkpoint, pyg_data_path)
    torch.save(model.state_dict(), "./storage_graph/gnn_model.pth")
    print("Done.")


if __name__ == "__main__":
    train_gnn()