# Graph-tool is not available in the standard colab environment so we must first install it...
%%capture
!echo "deb http://downloads.skewed.de/apt bionic main" >> /etc/apt/sources.list
!apt-key adv --keyserver keyserver.ubuntu.com --recv-key 612DEFB798507F25
!apt-get update
!apt-get install python3-graph-tool python3-matplotlib python3-cairo
#python3-cairo from Ubuntu's reposity is linked with a different python version; we need to improvise
!apt purge python3-cairo
!apt install libcairo2-dev pkg-config python3-dev
!pip install --force-reinstall --target=$nb_path pycairo
!pip install --target=$nb_path zstandard

# Install node2vec to create embeddings
%%capture
# update gensim
!pip install -U gensim
!pip install node2vec

def create_predictions(df_all_pairs, df_s, test):
  """
  Train a LogisticRegression given the labels from df_all_pairs and the scores from df_s. 
  The model is evaluated in the test dataset

  df_all_pairs: Columns 0,1,label (label = presence of a link in the training data (the graph))
  df_s: Columns 0,1,...
  test: Columns 0,1, link (link = presence of a link in the test data (the actual graph))
  """
  # Use all columns except the first two as independent variables
  col_emb = list(df_s.columns)[2:]

  # Merge the lables to the scores to create a model
  df = pd.merge(df_all_pairs, df_s, how="left", validate="1:1")

  # Create model
  model = LogisticRegressionCV(cv=3, max_iter=1000, class_weight="balanced", n_jobs=-1, scoring="f1")
  model.fit(df[col_emb], df["link"])
  print(model.coef_, model.intercept_)
  # Predict labels
  df["pred"] = model.predict(df[col_emb])
  
  # Merge the predictions with the test dataset to evaluate the prediction
  df_test = pd.merge(test, df)
  
  #recall = ability of the classifier to find all the positive samples.
  #precision = ability of the classifier not to label as positive a sample that is negative.
  print(classification_report(df_test["label"],df_test["pred"]))

  # Plot ROC curve
  plot_roc_curve(model, df_test[col_emb], df_test["label"]) 
  plt.grid(True)

  # return predicted probabilities
  return model.predict_proba(df[col_emb])[:,0]

# create edgelist dataframe (to be able to merge with the networkx data)
edgelist = pd.DataFrame([vi for vi in product(range(g.num_vertices()), repeat=2)])

# map graph-tool labels to graphml node labels
gt_to_graphml_labels = dict((i, int(g.vp['_graphml_vertex_id'][i])) for i in range(g.num_vertices()))
edgelist[0] = edgelist[0].map(gt_to_graphml_labels)
edgelist[1] = edgelist[1].map(gt_to_graphml_labels)

# Calculate jaccard similarity, convert to dataframe (cols 0 and 1 are the node ids)
## calculate jaccard similarity and convert to an array
s_jaccard = gt.vertex_similarity(g, "jaccard")
jaccard = edgelist.copy()
jaccard['score'] = np.array([s_jaccard[i] for i in range(num_vertices)]).reshape(-1)
jaccard['score'] = jaccard['score'].fillna(0) #there are somehow missing values (self-loops, but not only)


## You can change the "jaccard" argument to creat enew one
gt.vertex_similarity?

# Put it into a nice format
df_katz = edgelist.copy()
df_katz["katz"] = ks.reshape(-1)

def calculate_dist(v1, v2):
  """
  Calculate the L1 distance
  """
  return np.abs(v1 - v2)

def create_pair_distance(embeddings, G):
  """
  Calculates the distance between embeddings using calculate_dist
  """
  # Create dictionary with embeddings
  df_clsvd = pd.DataFrame(embeddings)
  df_clsvd.index = G.nodes()
  df_clsvd = {k: row.values for k,row in df_clsvd.iterrows()}


  # Calculate all distances and convert to dictionary
  pairs = combinations(G.nodes(), 2)
  all_pairs = [calculate_dist(df_clsvd[n1], df_clsvd[n2]) for n1,n2 in pairs]

  # Concatenate labels with node pairs      
  all_pairs = pd.DataFrame(all_pairs)
  all_pairs =  pd.concat([pd.DataFrame(combinations(G.nodes(), 2)), all_pairs.reset_index(drop=True)], axis=1)
  all_pairs.columns = [0,1]+[f"_{i}" for i in range(2, len(all_pairs.columns))]

  # Add the reverse of the node pair
  all_pairs = pd.concat([all_pairs, all_pairs.rename(columns={0:1, 1:0})])
  
  return all_pairs

A = nx.to_scipy_sparse_matrix(G)
clusters_svd = TruncatedSVD(n_components=10)
clusters_svd.fit(A)
cl_svd = clusters_svd.components_.T

# clusters_svd = SpectralEmbedding(affinity = 'precomputed', n_components=5)
# clusters_svd.fit(A)
# cl_svd = clusters_svd.embedding_

df_emb = create_pair_distance(cl_svd, G)

# Create all distances
df_emb = create_pair_distance(node2vec, G)

# Columns to use in the prediction
col_pred = ["jaccard","katz","svd","node2vec","sbm","pa"]

#Link prediction is hard!

# Split edge set for training and testing
u, v = g.edges()

eids = np.arange(g.number_of_edges())
eids = np.random.permutation(eids)
test_size = int(len(eids) * 0.1)
train_size = g.number_of_edges() - test_size
test_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]]
train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]

# Find all negative edges and split them for training and testing
adj = sp.coo_matrix((np.ones(len(u)), (u.numpy(), v.numpy())))
adj_neg = 1 - adj.todense() - np.eye(g.number_of_nodes())
neg_u, neg_v = np.where(adj_neg != 0)

neg_eids = np.random.choice(len(neg_u), g.number_of_edges())
test_neg_u, test_neg_v = neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]]
train_neg_u, train_neg_v = neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]]

train_g = dgl.remove_edges(g, eids[:test_size])

from dgl.nn import SAGEConv

# ----------- 2. create model -------------- #
# build a two-layer GraphSAGE model
class GraphSAGE(nn.Module):
    def __init__(self, in_feats, h_feats):
        super(GraphSAGE, self).__init__()
        self.conv1 = SAGEConv(in_feats, h_feats, 'mean')
        self.conv2 = SAGEConv(h_feats, h_feats, 'mean')

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = F.relu(h)
        h = self.conv2(g, h)
        return h

train_pos_g = dgl.graph((train_pos_u, train_pos_v), num_nodes=g.number_of_nodes())
train_neg_g = dgl.graph((train_neg_u, train_neg_v), num_nodes=g.number_of_nodes())

test_pos_g = dgl.graph((test_pos_u, test_pos_v), num_nodes=g.number_of_nodes())
test_neg_g = dgl.graph((test_neg_u, test_neg_v), num_nodes=g.number_of_nodes())

import dgl.function as fn

class DotPredictor(nn.Module):
    def forward(self, g, h):
        with g.local_scope():
            g.ndata['h'] = h
            # Compute a new edge feature named 'score' by a dot-product between the
            # source node feature 'h' and destination node feature '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]

class MLPPredictor(nn.Module):
    def __init__(self, h_feats):
        super().__init__()
        self.W1 = nn.Linear(h_feats * 2, h_feats)
        self.W2 = nn.Linear(h_feats, 1)

    def apply_edges(self, edges):
        """
        Computes a scalar score for each edge of the given graph.

        Parameters
        ----------
        edges :
            Has three members ``src``, ``dst`` and ``data``, each of
            which is a dictionary representing the features of the
            source nodes, the destination nodes, and the edges
            themselves.

        Returns
        -------
        dict
            A dictionary of new edge features.
        """
        h = torch.cat([edges.src['h'], edges.dst['h']], 1)
        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)
            return g.edata['score']

model = GraphSAGE(g.ndata['feat'].shape[1], 16)
# You can replace DotPredictor with MLPPredictor.
#pred = MLPPredictor(16)
pred = DotPredictor()

def compute_loss(pos_score, neg_score):
    scores = torch.cat([pos_score, neg_score])
    labels = torch.cat([torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])])
    return F.binary_cross_entropy_with_logits(scores, labels)

def compute_auc(pos_score, neg_score):
    scores = torch.cat([pos_score, neg_score]).numpy()
    labels = torch.cat(
        [torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]).numpy()
    return roc_auc_score(labels, scores)

# ----------- 3. set up loss and optimizer -------------- #
# in this case, loss will in training loop
optimizer = torch.optim.Adam(itertools.chain(model.parameters(), pred.parameters()), lr=0.01)

# ----------- 4. training -------------------------------- #
all_logits = []
for e in range(100):
    # forward
    h = model(train_g, train_g.ndata['feat'])
    pos_score = pred(train_pos_g, h)
    neg_score = pred(train_neg_g, 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))




# ----------- 5. check results ------------------------ #
from sklearn.metrics import roc_auc_score
with torch.no_grad():
    pos_score = pred(test_pos_g, h)
    neg_score = pred(test_neg_g, h)
    print('AUC', compute_auc(pos_score, neg_score))


# Thumbnail credits: Link Prediction with Neo4j, Mark Needham
# sphinx_gallery_thumbnail_path = '_static/blitz_4_link_predict.png'

#embeddings
h.shape

if 0:
  # Create network for challenge
  df = pd.read_csv("http://interactome.dfci.harvard.edu/S_cerevisiae/download/CCSB-Y2H.txt",sep="\t",header=None)
  df = pd.concat([df, df.rename(columns={0:1, 1:0})])

  G = nx.from_pandas_edgelist(df, source=0, target=1)
  G.remove_edges_from(nx.selfloop_edges(G))
  Gcc = sorted(nx.connected_components(G), key=len, reverse=True)
  G0 = G.subgraph(Gcc[0])

  df = pd.DataFrame(G0.edges())
  df = pd.concat([df, df.rename(columns={0:1, 1:0})]).drop_duplicates()

  nodes = df[0].unique()
  d_conv = dict(zip(nodes, range(len(nodes))))
  df[0] = df[0].map(d_conv)
  df[1] = df[1].map(d_conv)

  G = nx.from_pandas_edgelist(df, source=0, target=1)

  adj = nx.to_scipy_sparse_array(G, nodelist=range(len(G)))


  u,v = np.nonzero(adj)
  eids = np.arange(len(u))
  eids = np.random.permutation(eids)
  test_size = int(len(eids) * 0.1)
  train_size = len(u) - test_size
  test_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]]
  train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]

  # Find all negative edges and split them for training and testing
  adj_neg =  adj.todense() + np.eye(len(G)) - 1
  neg_u, neg_v = np.nonzero(adj_neg)

  neg_eids = np.random.choice(len(neg_u), len(u))
  test_neg_u, test_neg_v = neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]]
  train_neg_u, train_neg_v = neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]]
  # Positive and negative examples to test (50-50)
  pos = pd.DataFrame(np.array([test_pos_u,test_pos_v]).T)
  pos["label"] = 1
  neg = pd.DataFrame(np.array([test_neg_u,test_neg_v]).T)
  neg["label"] = 0
  df = pd.concat([pos,neg])
  df.to_csv("../Data/ppi_network_prediction_test.csv", sep="\t")
            
  # Network with 10% removed
  G.remove_edges_from(zip(test_pos_u, test_pos_v))
  nx.set_node_attributes(G, {v:k for k,v in d_conv.items()}, name="label")
  nx.write_graphml(G, "../Data/ppi_network_prediction.graphml")


x=  pd.merge(pd.DataFrame(G.edges()),df, how="outer", indicator=True)
x["_merge"].value_counts()