#!/usr/bin/env python3
"""
Train a Variational Graph Autoencoder (VGAE) on a chromatin contact graph.

Inputs
------
  PyTorch Geometric Data object saved by build_graph.py.

Outputs (under --outdir)
------------------------
  model.pt      trained VGAE state_dict (includes BatchNorm running statistics)
  emb.npy       node embeddings — mu vector, shape [num_nodes, latent_dim]
  metrics.json  val/test AUC & AP plus all hyperparameters
"""

import argparse
import json
import os
import sys

import numpy as np
import torch
from torch_geometric.nn.models import VGAE
from torch_geometric.transforms import RandomLinkSplit
from torch_geometric.utils import (
    negative_sampling,
    remove_self_loops,
    to_undirected,
)
from sklearn.metrics import average_precision_score, roc_auc_score

sys.path.insert(0, os.path.dirname(__file__))
from model import Encoder


@torch.no_grad()
def _eval_linkpred(z, pos_edges, neg_edges):
    """Return (AUROC, AP) for link prediction."""
    def _sigmoid(x):
        return 1.0 / (1.0 + torch.exp(-x))

    def _score(edges):
        src, dst = edges
        return _sigmoid((z[src] * z[dst]).sum(dim=1)).cpu().numpy()

    y_true = np.concatenate([np.ones(pos_edges.size(1)),
                              np.zeros(neg_edges.size(1))])
    y_pred = np.concatenate([_score(pos_edges), _score(neg_edges)])
    return roc_auc_score(y_true, y_pred), average_precision_score(y_true, y_pred)


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--graph",    required=True,
                    help="Path to Data .pt file from build_graph.py")
    ap.add_argument("--epochs",   type=int,   default=300)
    ap.add_argument("--patience", type=int,   default=20,
                    help="Early-stopping patience (val-AUC epochs without improvement)")
    ap.add_argument("--lr",       type=float, default=1e-3)
    ap.add_argument("--hidden",   type=int,   default=64)
    ap.add_argument("--latent",   type=int,   default=32)
    ap.add_argument("--dropout",  type=float, default=0.2)
    ap.add_argument("--seed",     type=int,   default=42)
    ap.add_argument("--outdir",   default="results")
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    os.makedirs(args.outdir, exist_ok=True)

    # ---- Load and clean graph ----
    data = torch.load(args.graph, weights_only=False)
    ei, _ = remove_self_loops(data.edge_index)
    data.edge_index = to_undirected(ei, num_nodes=data.num_nodes)
    x = data.x.float()
    print(f"Graph: {data.num_nodes} nodes  "
          f"{data.edge_index.shape[1]} edges  "
          f"{x.shape[1]} node features")

    # ---- Edge splits for link-prediction evaluation ----
    splitter = RandomLinkSplit(
        num_val=0.1, num_test=0.1,
        is_undirected=True,
        add_negative_train_samples=False,
        split_labels=False,
    )
    train_data, val_data, test_data = splitter(data)
    train_data.pos_edge_index = train_data.edge_index

    for split in (val_data, test_data):
        split.pos_edge_index = split.edge_index
        split.neg_edge_index = negative_sampling(
            edge_index=split.edge_index,
            num_nodes=data.num_nodes,
            num_neg_samples=split.edge_index.size(1),
            method="sparse",
        )

    # ---- Model ----
    enc   = Encoder(in_dim=x.size(1), hidden=args.hidden,
                    latent=args.latent, dropout=args.dropout)
    model = VGAE(enc)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    # ---- Training loop with early stopping ----
    best_val_auc = -1.0
    best_state   = None
    no_improve   = 0
    epochs_ran   = 0

    for epoch in range(1, args.epochs + 1):
        model.train()
        optimizer.zero_grad()
        z    = model.encode(x, train_data.edge_index)
        loss = (model.recon_loss(z, train_data.pos_edge_index)
                + (1.0 / data.num_nodes) * model.kl_loss())
        loss.backward()
        optimizer.step()

        model.eval()
        with torch.no_grad():
            z_full = model.encode(x, data.edge_index)
            val_auc, val_ap = _eval_linkpred(
                z_full, val_data.pos_edge_index, val_data.neg_edge_index
            )

        if val_auc > best_val_auc:
            best_val_auc = val_auc
            best_state   = {k: v.cpu().clone()
                            for k, v in model.state_dict().items()}
            no_improve   = 0
        else:
            no_improve += 1

        epochs_ran = epoch
        if epoch % 10 == 0 or epoch == 1:
            print(f"[{epoch:03d}/{args.epochs}] "
                  f"loss={loss.item():.4f}  "
                  f"val AUC={val_auc:.4f}  AP={val_ap:.4f}")

        if no_improve >= args.patience:
            print(f"Early stopping at epoch {epoch} "
                  f"(no val-AUC improvement for {args.patience} epochs)")
            break

    # ---- Restore best checkpoint and compute test metrics ----
    model.load_state_dict(best_state)
    model.eval()
    with torch.no_grad():
        z_final = model.encode(x, data.edge_index)
        test_auc, test_ap = _eval_linkpred(
            z_final, test_data.pos_edge_index, test_data.neg_edge_index
        )

    # ---- Save outputs ----
    model_path = os.path.join(args.outdir, "model.pt")
    torch.save(best_state, model_path)

    emb_path = os.path.join(args.outdir, "emb.npy")
    np.save(emb_path, z_final.cpu().numpy())

    metrics = {
        "val_auc":    float(best_val_auc),
        "test_auc":   float(test_auc),
        "test_ap":    float(test_ap),
        "epochs_ran": epochs_ran,
        "epochs_max": args.epochs,
        "patience":   args.patience,
        "hidden":     args.hidden,
        "latent":     args.latent,
        "dropout":    args.dropout,
        "lr":         args.lr,
        "seed":       args.seed,
    }
    with open(os.path.join(args.outdir, "metrics.json"), "w") as f:
        json.dump(metrics, f, indent=2)

    print(f"\nSaved model       → {model_path}")
    print(f"Saved embeddings  → {emb_path}  shape={z_final.shape}")
    print(f"Test  AUC={test_auc:.4f}  AP={test_ap:.4f}")


if __name__ == "__main__":
    main()