# Databricks notebook source
# MAGIC %md
# MAGIC # Link Prediction demo
# MAGIC
# MAGIC In this notebook we provide a tutorial for performing link prediction. That is, given a graph dataset, our task is to learn a Graph Neural Network model for predicting the existence of edges between node pairs. This tutorial assumes some theoretical knowledge on GNNs and some basic skill in Graph manupilation in Python. See, for example, [graphsage_theory](/docs/theory_docs/graphsage_theory.md) and [demo_graph_structure](/notebooks/graph_mdodels/demo_graph_structure.py).
# MAGIC
# MAGIC The following aspects will be developed in this notebook:
# MAGIC - Load the FinCEN graph - positive and negative example
# MAGIC - Customized NN layers
# MAGIC  - Graph Convolutional layers
# MAGIC  - GraphSAGE layers
# MAGIC - Score function
# MAGIC  - Dot product
# MAGIC  - Multilayer perceptron
# MAGIC - Model training
# MAGIC - Performance evaluation

# COMMAND ----------

# import required libraries

import itertools
import os

import dgl
import dgl.function as fn
import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch as th
import torch.nn as nn
import torch.nn.functional as F
from dgl.data import DGLDataset
from dgl.data.utils import load_graphs
from dgl.nn import SAGEConv
from matplotlib import pyplot
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
    roc_curve,
)

# COMMAND ----------

# set seed for reproducibility

th.manual_seed(4)

# COMMAND ----------

# MAGIC %md
# MAGIC ## Load the FinCEN graph - positive and negative example
# MAGIC From the original bank transactions FinCEN dataset an undirected graph having both nodes and edges features has been generated (see how [here](/notebooks/
# MAGIC graph_models/demo_graph_structure.py)). The node feature is a 135-dimensional _dummy_ vector which identify the country of the bank branch represented by the node. The edge feature is a 2-dimensional vector representing the amount of the transaction and the number of sub-transactions carried through that edge; both features are standardized. In order to train a link prediction model, it is required the creation of the _positive_ and _negative_ examples associated to observed graph; both subgraphs have been previously generated [here](/notebooks/
# MAGIC graph_models/positive_and_negative_graph.py).

# COMMAND ----------

# load the FinCEN graph

fincen_graph = load_graphs("/dbfs/mnt/public/clean/fincen-graph/fincen_dgl")[0][0]

fincen_graph  # short summary of the graph

# COMMAND ----------

# load train_fincen_graph - used to train the GNN

train_fincen_graph = load_graphs(
    "/dbfs/mnt/public/clean/fincen-graph/train_fincen_graph"
)[0][0]


# COMMAND ----------

# load positive examples for train and test set

train_pos_fincen_graph = load_graphs(
    "/dbfs/mnt/public/clean/fincen-graph/train_pos_fincen_graph"
)[0][0]

test_pos_fincen_graph = load_graphs(
    "/dbfs/mnt/public/clean/fincen-graph/test_pos_fincen_graph"
)[0][0]


# COMMAND ----------

# load negative examples for train and test set

train_neg_fincen_graph = load_graphs(
    "/dbfs/mnt/public/clean/fincen-graph/train_neg_fincen_graph"
)[0][0]

test_neg_fincen_graph = load_graphs(
    "/dbfs/mnt/public/clean/fincen-graph/test_neg_fincen_graph"
)[0][0]


# COMMAND ----------

# MAGIC %md
# MAGIC ## Customized NN layers
# MAGIC
# MAGIC We now provide 2 different layers, where the embeddings for each nodes are generated using different _message_ functions. The Convolutional layers are built entirely by us and apply first a concatenation of edge and node features from the neighbours, forwarded through a linear transformation. In the second example we exploit the built-in SAGEconv layers, ad-hoc modified in order to allow the passing of edges features in addition to the node features. Although both methods result quite similar, the Convolutional layers only use information about the neighbours while SAGEconv update the embedding of a node by weigthing information from the neighbours and from the node itself.

# COMMAND ----------

# MAGIC %md
# MAGIC
# MAGIC ## Graph Convolutional layers
# MAGIC
# MAGIC In order to pass the messages from nodes and edges to the embddings we define 2 convolutional layers:
# MAGIC 1. GCNLayer_one
# MAGIC 2. GCNLayer_two
# MAGIC
# MAGIC For each node, GCNLayer_one concatenate the features of incoming edges and neighbouring nodes. Hence the messages is reduced to a single vector by taking the element-wise mean. The new embeddings are then generated by applying a linear transformation.
# MAGIC The second layer, GCNLayer_two, performs the same operation but receive as input the embedding generate by the first layer. Below we show how both layer are defined in Python using the library _DGl_ and _torch.nn_.

# COMMAND ----------

# GCNLayer_one: the first neural network layer


class GCNLayer_one(nn.Module):

    # INPUTS:
    # in_feats: dimension of the features (nodes + edges)
    # out_feats: dimension of the embeddings

    def __init__(self, in_feats, out_feats):

        super(GCNLayer_one, self).__init__()

        self.linear = nn.Linear(
            in_feats, out_feats
        )  # define the name of the linear transformation function

    # INPUTS:
    # g: the observed graph
    # feature_nodes: the torch vector with all the node features
    # feature_edges: the torch vector with all the edge features

    def forward(self, g, feature_nodes, feature_edges):
        with g.local_scope():

            g.ndata["h"] = feature_nodes  # store the node features in the  graph as "h"
            g.edata["e"] = feature_edges  # store the edge features in the graph as "e"

            g.update_all(
                fn.copy_e("e", "m_e"), fn.mean("m_e", "h_e")
            )  # take all the features attached to all the in-coming edges of a node and perform the mean; call the output vector "h_e"

            g.update_all(
                fn.copy_u("h", "m_n"), fn.mean("m_n", "h_n")
            )  # take all the features attached to all the neighbouring nodes of a node and perform the mean; call the output vector "h_n"

            g.ndata["h"] = th.cat(
                (g.ndata.pop("h_n"), g.ndata.pop("h_e")), 1
            )  # ceate a new embedding for each node by concatenating the vectors "h_e" and "h_n"; call the embedding "h"

            h_n = g.ndata["h"]  # store the new embedding matric as "h_n"

            h_n = self.linear(h_n)  # perform the linear transformation

            return F.normalize(h_n)  # normalize the embedding so to have length 1


# COMMAND ----------

# GCNLayer_two: the second neural network layer


class GCNLayer_two(nn.Module):

    # INPUTS:
    # in_feats: dimension of the input embeddings
    # out_feats: dimension of the output embeddings

    def __init__(self, in_feats, out_feats):

        super(GCNLayer_two, self).__init__()

        self.linear = nn.Linear(
            in_feats, out_feats
        )  # define the name of the linear transformation function

    def forward(self, g, embeddings):

        with g.local_scope():

            g.ndata["emb"] = embeddings  # store the input embedding as  "emb"

            g.update_all(
                fn.copy_u("emb", "m_emb"), fn.mean("m_emb", "h_emb")
            )  # take the embeddings of the neighbouring nodes of a node and perform the mean; call the output vector "h_emb"

            h_final = g.ndata["h_emb"]  # store the new embedding matrix as "h_final"

            h_final = self.linear(h_final)  # perform the linear transformation

            return F.normalize(h_final)  # normalize the embedding so to have length 1


# COMMAND ----------

# MAGIC %md
# MAGIC Once the layers have been designed we define the Neural Network layer.

# COMMAND ----------

# define Neural Network consisting of 2 layers


class conv_network(nn.Module):

    # INPUTS:
    # in_feats: dimension of the initial embedding
    # h_feats: dimension of the output embedding

    def __init__(self, in_feats, h_feats):

        super(conv_network, self).__init__()

        self.layer1 = GCNLayer_one(in_feats, h_feats)  # first NN layer

        self.layer2 = GCNLayer_two(h_feats, h_feats)  # second NN layer

    def forward(self, g, feature_nodes, feature_edges):

        x = self.layer1(g, feature_nodes, feature_edges)  # first embedding

        x = F.relu(
            x
        )  # non-linearity relu function: if x[i]>0 --> x[i] = x[i]; else x[i] = 0

        x = self.layer2(g, x)  # the second embedding

        return x


# COMMAND ----------

# MAGIC %md
# MAGIC ## GraphSAGE layers
# MAGIC
# MAGIC The second Neural Network module we design is a modification of the GraphSAGE model. At each iteration the new embedding of a node is generated by combining information from the neighbour nodes (and incident edges) with the embedding prevously generated.

# COMMAND ----------

# customized SAGEconv layer: it uses also the edge features

# build a two-layer GraphSAGE model
class GraphSAGE(nn.Module):

    # INPUTS:
    # in_feats: dimension of the initial embedding
    # h_feats: dimension of the output embedding
    def __init__(self, in_feats, h_feats):

        super(GraphSAGE, self).__init__()

        self.conv1 = SAGEConv(in_feats, h_feats, "mean")  # first GraphSAGE layer

        self.conv2 = SAGEConv(h_feats, h_feats, "mean")  # second GraphSAGE layer

    def forward(self, g, feature_nodes, feature_edges):

        g.ndata["h_n"] = feature_nodes  # store the node features in the  graph as "h"

        g.edata["e"] = feature_edges  # store the edge features in the graph as "e"

        g.update_all(
            fn.copy_e("e", "m_e"), fn.mean("m_e", "h_e")
        )  # take all the features attached to all the in-coming edges of a node and perform the mean; call the output vector "h_e"

        g.ndata["h"] = th.cat(
            (g.ndata.pop("h_n"), g.ndata.pop("h_e")), 1
        )  # ceate a new embedding for each node by concatenating the vectors "h_e" and "h_n"; call the embedding "h"

        h_n = g.ndata["h"]  # store the new embedding matric as "h_n"

        h = self.conv1(g, h_n)  # perform the first  convolutional layer

        h = F.relu(h)  # apply non-linearity

        h = self.conv2(g, h)  # perform the second convolutional layer

        return h


# COMMAND ----------

# MAGIC %md
# MAGIC
# MAGIC ## The score function
# MAGIC
# MAGIC The Neural Networks defined above allow us to learn the embeddings of the nodes. In order to fullfill the link prediction task, we use the embeddings for producing an edge score. Intuitively, the higher the score the more likely is the connection between the nodes. We propose two score functions: a simple dot producut and a multilayer perceptron.

# COMMAND ----------

# compute the edge score


class Dot_Predictor(nn.Module):
    def forward(self, g, h):
        with g.local_scope():

            g.ndata["h"] = h  # store the embdeddings as node features

            # Compute a new edge feature named 'score' by a dot-product between the
            # source node embedding 'h' and destination node embedding 'h'.

            g.apply_edges(fn.u_dot_v("h", "h", "score"))

            # u_dot_v returns a 1-element vector for each edge so you need to squeeze it.

            return g.edata["score"][:, 0]


# COMMAND ----------


class MLP_Predictor(nn.Module):

    # INPUTS:
    # h_heats: size of the embeddgins
    def __init__(self, h_feats):

        super().__init__()

        self.W1 = nn.Linear(h_feats * 2, h_feats)  # first layer

        self.W2 = nn.Linear(h_feats, 1)  # second layer - note output dimension is 1

    def apply_edges(self, edges):

        # concatenate features from source and destination node associated to each edge
        h = th.cat([edges.src["h"], edges.dst["h"]], 1)

        # return the score after applying layer 1 and 2 (and non linearity in btw)
        return {"score": self.W2(F.relu(self.W1(h))).squeeze(1)}

    def forward(self, g, h):
        with g.local_scope():

            g.ndata["h"] = h

            g.apply_edges(self.apply_edges)  # create the new node feature

            return g.edata["score"]


# COMMAND ----------

# MAGIC %md
# MAGIC ## Model training
# MAGIC
# MAGIC We now learn end evaluate the the graphSAGE GNN with MLP score function. For comparison, at the end of the notebook we include the code for performing link prediction using the self-built GCN layer with MLP score function.
# MAGIC
# MAGIC Due to that we deal with a binary classification task, in order to train the model we specify the cross entropy loss. Furthermore, we asses the performance of the model using the ROC curve,  _Area Under the Curve_ and F1 score.

# COMMAND ----------

# Define the cross entropy loss fuction
def compute_loss(pos_score, neg_score):
    scores = th.cat([pos_score, neg_score])
    labels = th.cat([th.ones(pos_score.shape[0]), th.zeros(neg_score.shape[0])])
    return F.binary_cross_entropy_with_logits(scores, labels)


# Define the AUC metric
def compute_auc(pos_score, neg_score):
    scores = th.cat([pos_score, neg_score]).numpy()
    labels = th.cat([th.ones(pos_score.shape[0]), th.zeros(neg_score.shape[0])]).numpy()
    return roc_auc_score(labels, scores)


# COMMAND ----------

# Register the Neural Network with the "outer"  input
model = GraphSAGE(
    train_fincen_graph.ndata["feat_nodes"].shape[1]
    + train_fincen_graph.edata["feat_edges"].shape[1],
    16,
)

# Register the edge score fucntion with the "outer" input
# pred = Dot_Predictor()

pred = MLP_Predictor(16)

# COMMAND ----------

# set up loss and optimizer -
optimizer = th.optim.Adam(
    itertools.chain(model.parameters(), pred.parameters()), lr=0.01
)


###### TRAIN the MODEL

for e in range(100):
    # forward
    h = model(
        train_fincen_graph,
        train_fincen_graph.ndata["feat_nodes"].float(),
        train_fincen_graph.edata["feat_edges"].float(),
    )
    pos_score = pred(train_pos_fincen_graph, h)
    neg_score = pred(train_neg_fincen_graph, h)
    loss = compute_loss(pos_score, neg_score)

    # backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if e % 5 == 0:
        print("In epoch {}, loss: {}".format(e, loss))

# COMMAND ----------

# MAGIC %md
# MAGIC ## Performance evaluation
# MAGIC
# MAGIC We are ready now to evaluate the performance of the model: first we compute the ROC curve so to assess the _recall_ of the classifier and find the optimal threshold which, in general, depends on the experiment´s domain. Next we comopute the _Area Under the Curve_ (AUC) metric in order to get a measure of the averall predictive power of the model, regardless of the chosen threshold.

# COMMAND ----------

# Plot ROC curve

np.set_printoptions(suppress=True)

with th.no_grad():

    pos_score = pred(test_pos_fincen_graph, h)

    neg_score = pred(test_neg_fincen_graph, h)

# true labels for the edges in the test set
y_test = np.append(np.ones(len(pos_score)), np.zeros(len(neg_score)))

# compute predicted probabilities of existence of the edges in the test set
sigmoid_function = nn.Sigmoid()

# the predicted score are non-normalized, we use the sigmoid function to constraint them in (0,1)
score_test = sigmoid_function(
    th.from_numpy(np.append(pos_score.numpy(), neg_score.numpy()))
)

# retrive false-positive-rate, true-positive-rate and threshold
fpr, tpr, thresholds = roc_curve(y_test, score_test)

# show plot
pyplot.plot(fpr, tpr, marker=".", label="Logistic")

# axis labels
pyplot.xlabel("False Positive Rate")

pyplot.ylabel("True Positive Rate")

# COMMAND ----------

# Compute AUC measure

with th.no_grad():

    pos_score = pred(test_pos_fincen_graph, h)

    neg_score = pred(test_neg_fincen_graph, h)

    print("AUC", compute_auc(pos_score, neg_score))

# COMMAND ----------

# MAGIC %md
# MAGIC
# MAGIC The AUC is 0.91. Assuming we are saisfied with a recall equal to 0.9, we select the corresponding threshold and compute the F1 score which results 0.84.

# COMMAND ----------

star_p = thresholds[np.where(tpr >= 0.90)][0]  # selected threshold probability

y_predicted = np.zeros(len(score_test))

y_predicted[score_test >= star_p] = 1.0

f1_score(y_test, y_predicted)

# COMMAND ----------

pyplot.subplot(1, 2, 1)

# plot the predicted probability for "positive" and "negative" edges
pyplot.hist(score_test[0 : len(pos_score)], color="green")

# axis labels
pyplot.xlabel("predicted probability")
pyplot.ylabel("count")
pyplot.title("positive edges")
######

pyplot.subplot(1, 2, 2)


pyplot.hist(score_test[len(pos_score) :], color="red")

# axis labels
pyplot.xlabel("predicted probability")
pyplot.title("negative edges")


# COMMAND ----------

# MAGIC %md
# MAGIC
# MAGIC Store the results for later use (visualization link prediction).

# COMMAND ----------

# DataFrame with prediction for the positive test set

src_pos = (
    test_pos_fincen_graph.edges()[0].numpy().reshape(len(pos_score), 1)
)  # src nodes in positive test set

dst_pos = (
    test_pos_fincen_graph.edges()[1].numpy().reshape(len(pos_score), 1)
)  # dst nodes in positive test set

label_pos = np.ones(len(pos_score)).reshape(
    len(pos_score), 1
)  # true label for the edges in the positive test set

score_pos = score_test[: len(pos_score)].reshape(
    len(pos_score), 1
)  # predicted probabilitiy of connection for the edges in the positive test set


result_pos = np.concatenate((src_pos, dst_pos, score_pos, label_pos), axis=1)

result_pos = pd.DataFrame(data=result_pos, columns=["src", "dst", "score", "label"])


# DataFrame with prediction for the negative test set

src_neg = (
    test_neg_fincen_graph.edges()[0].numpy().reshape(len(neg_score), 1)
)  # src nodes in negative test set

dst_neg = (
    test_neg_fincen_graph.edges()[1].numpy().reshape(len(neg_score), 1)
)  # dst nodes in negative test set

label_neg = np.zeros(len(neg_score)).reshape(
    len(neg_score), 1
)  # true label for the edges in the negative test set

score_neg = score_test[len(neg_score) :].reshape(
    len(neg_score), 1
)  # predicted probabilitiy of connection for the edges in the negative test set


result_neg = np.concatenate((src_neg, dst_neg, score_neg, label_neg), axis=1)

result_neg = pd.DataFrame(data=result_neg, columns=["src", "dst", "score", "label"])

# combine the 2 datasets
results_data = result_pos.append(result_neg, ignore_index=True)


# store results

results_data.to_csv(
    "/dbfs/mnt/public/clean/fincen-graph/results_link_prediction.csv", index=False
)


# COMMAND ----------

# MAGIC %md
# MAGIC **Code for training link prediction using custom Graph Convolutional layers and MLP score predictor **

# COMMAND ----------


# Register the Neural Network with the "outer"  input

model_2 = conv_network(
    fincen_graph.ndata["feat_nodes"].shape[1]
    + fincen_graph.edata["feat_edges"].shape[1],
    16,
)

# Register the edge score fucntion with the "outer" input
pred_2 = MLP_Predictor(16)

# set up loss and optimizer

optimizer_2 = th.optim.Adam(
    itertools.chain(model_2.parameters(), pred_2.parameters()), lr=0.01
)

# training

for e in range(300):

    # forward

    h_2 = model_2(
        train_fincen_graph,
        train_fincen_graph.ndata["feat_nodes"].float(),
        train_fincen_graph.edata["feat_edges"].float(),
    )

    pos_score_2 = pred_2(train_pos_fincen_graph, h_2)
    neg_score_2 = pred_2(train_neg_fincen_graph, h_2)
    loss_2 = compute_loss(pos_score_2, neg_score_2)

    # backward
    optimizer_2.zero_grad()
    loss_2.backward()
    optimizer_2.step()

    if e % 20 == 0:
        print("In epoch {}, loss: {}".format(e, loss_2))
