## If you are running this on your computer you can install the libraries uncommenting these lines
## Don't mix installations with conda and installations with pip (bad practice)
# conda install -c conda-forge graph-tool
# conda install -c pytorch pytorch
# conda install -c dglteam dgl
# mamba install "scipy<1.13" #required for node2vec

!pip install -q condacolab
import condacolab
condacolab.install() # expect a kernel restart

# Install graph-tool 
!wget https://downloads.skewed.de/skewed-keyring/skewed-keyring_1.0_all_$(lsb_release -s -c).deb
!dpkg -i skewed-keyring_1.0_all_$(lsb_release -s -c).deb
!echo "deb [signed-by=/usr/share/keyrings/skewed-keyring.gpg] https://downloads.skewed.de/apt $(lsb_release -s -c) main" > /etc/apt/sources.list.d/skewed.list
!apt-get update
!apt-get install python3-graph-tool python3-matplotlib python3-cairo

# Install node2vec to create embeddings
!pip install scipy==1.12
!pip install -U gensim
!pip install node2vec

# It should be in this order
!mamba install -c dglteam dgl

# networks
import networkx as nx
import graph_tool.all as gt
from graph_tool import topology, inference, generation, stats, correlations, clustering

# machine learning
from sklearn.metrics import classification_report, roc_auc_score
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegressionCV
from sklearn.metrics import RocCurveDisplay

# embeddings
from sklearn.manifold import SpectralEmbedding
from sklearn.decomposition import PCA, TruncatedSVD, NMF

# Some math
import scipy.sparse as sp
import numpy as np
from numpy.linalg import inv

# data wrangling
import pandas as pd
from itertools import permutations, combinations, product

# viz
import pylab as plt

# Disable SettingWithCopyWarning
pd.options.mode.chained_assignment = None

# path_data = "../../Data/"
#Don't change this on colab
path_data = "./"

network = "ppi"

# Read network (download it from internet) 
## ONLY IN COLAB! (otherwise use the data from the Data folder, no need to download)
!wget https://github.com/jgarciab/NetworkScience/raw/main/Data/{network}_network_prediction.graphml 
!wget https://github.com/jgarciab/NetworkScience/raw/main/Data/{network}_network_prediction_test.csv

# Path network
path_network = f"{path_data}/{network}_network_prediction.graphml"

# Read to networkx, this allows us to use the networkx functions
if network == "ppi":
    G = nx.read_graphml(path_network, node_type=int)
else:
    G = nx.read_graphml(path_network)
    
# Read to graph-tool, this allows us to use the graph-tool functions
g = gt.load_graph(path_network)
num_vertices = g.num_vertices()

# The node names in graph-tool are different from the ones in networkx
# let's create a dictionary to map graph-tool node labels to networkx node labels
if network == "ppi":
    gt_to_networkx_labels = dict((i, int(g.vp['_graphml_vertex_id'][i])) for i in range(g.num_vertices()))
else:
    gt_to_networkx_labels = dict((i, g.vp['_graphml_vertex_id'][i]) for i in range(g.num_vertices()))

def create_predictions(df_edgelist, df_test, columns = None, return_test = False, solver="lbfgs"):
    """
    Train a LogisticRegression given the labels from df_all_pairs and the scores from df_s. 
    The model is evaluated in the test dataset

    Input:
    - df_edgelist: Columns 0,1,edge (edge = presence of a link in the training data (the graph)), variables to use as predictors
    - df_ test: Columns 0,1, link (link = presence of a link in the test data (the actual graph))
    - return_test: if True, return the test dataset with the predictions
    - columns: list of columns to use as predictors
    - solver: solver to use in the LogisticRegression 

    Output:
    - If return_test is True, return the test dataset with the predictions
    - If return_test if False, return the predicted probabilities of all potential links
    """
    if columns is None:
        # Use all columns except the first two as independent variables (the first two are the labels)
        columns = list(df_edgelist.columns)[3:]

    # Create model and fit using cross-validation
    model = LogisticRegressionCV(cv=5, max_iter=1000,
                                  class_weight="balanced", n_jobs=-1, scoring="f1",
                                  solver=solver)
    
    # Remove self-loops (not necessarily what you want to do always, espcially in biological networks)
    df_edgelist = df_edgelist.loc[df_edgelist[0] != df_edgelist[1]]

    # Scale features
    X = df_edgelist[columns]
    scaler = StandardScaler()
    X = scaler.fit_transform(X)
    model.fit(X, df_edgelist["edge"])
    print(f"Intercept: {model.intercept_[0]: 2.2f}")
    for var, col in zip(columns, model.coef_[0]):
        print(f"{var}: {col: 2.2f}")

    # Predict labels for all pairs
    df_edgelist["pred"] = model.predict(X)

    # Merge the predictions with the test dataset to evaluate the prediction
    df_test = pd.merge(df_test, df_edgelist, on=[0,1], how="inner", validate="1:1")

    #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
    RocCurveDisplay.from_estimator(model, scaler.transform(df_test[columns]), df_test["label"]) 
    plt.plot([0,1], [0,1], "--", color="lightgray", zorder=0)
    plt.grid(True)

    if return_test:
        df_test["pred_prob"] = model.predict_proba(scaler.transform(df_test[columns]))[:,1]
        return df_test
    
    # return predicted probabilities
    return model.predict_proba(X)[:,1]

# 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")

df_edgelist2 = df_edgelist.copy()

# add a column with the jaccard similarity between node pairs
df_edgelist2['jaccard_score'] = np.array([s_jaccard[i] for i in range(num_vertices)]).reshape(-1)
df_edgelist2['jaccard_score'] = df_edgelist2['jaccard_score'].fillna(0) #there are somehow missing values (self-loops, but not only)

## You can change the "jaccard" argument to create a new metric
#gt.vertex_similarity?

# Calculate pref attachment --> Returns the name of the two nodes, and the product of the degree
pa = nx.preferential_attachment(G, all_edges_nx)
df_pa = pd.DataFrame(list(pa))
df_pa.columns = [0,1,"pa_score"]

df_edgelist2 = pd.merge(df_edgelist, df_pa, how="left", on=[0,1], validate="1:1")

A = nx.to_numpy_array(G)

#Make higher order powers
A2 = (A @ A)
np.fill_diagonal(A2, 0)
A3 = (A2 @ A)
np.fill_diagonal(A3, 0)
A4 = (A3 @ A)
np.fill_diagonal(A4, 0)


# Put it into the same format as edgelist
df_edgelist2 = df_edgelist.copy()
df_edgelist2["a2"] = A2.reshape(-1)
df_edgelist2["a3"] = A3.reshape(-1)
df_edgelist2["a4"] = A4.reshape(-1)

def calculate_dist(v1, v2):
    """
    Calculate the L1 distance between two embeddings
    (you could use a different formula, such as (v1*v2))
    """
    #I'm taking the sum, but you could let a downstream model find that out
    return np.sum(np.abs(v1 - v2))

def create_pair_distance(embeddings, G, fun_distance=calculate_dist):
    """
    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 = [fun_distance(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_numpy_array(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_

#Link prediction can be hard!

#d_conv_inv[776], d_conv_inv[560]

# Imports
import pandas as pd
import numpy as np
import itertools

import torch
import torch.nn as nn
import torch.nn.functional as F
#from torch.utils.data import DataLoader

import dgl
#import dgl.nn as dglnn
#from dgl.data.utils import save_graphs, load_graphs
#import dgl.function as fn
from dgl.nn import SAGEConv

# Set up a seed for replicability
np.random.seed(1546)

# Load data based on network type
network = "ppi"  # Example, set the network type

if network == "ppi":
    # Download PPI network
    df = pd.read_csv("https://github.com/jgarciab/NetworkScience/raw/main/Data/CCSB-Y2H.txt", skiprows=1, sep="\t", header=None)     
else:
    # Load Twitter data
    df = pd.read_csv("https://raw.githubusercontent.com/jgarciab/NetworkScience/main/Data/ic2s2_netsci_3.tsv", sep="\t", usecols=["source", "target"])
    df.columns = [0, 1]

# Create bidirectional edges and remove self-loops
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))

# Keep largest connected component
largest_cc = max(nx.connected_components(G), key=len)
G = G.subgraph(largest_cc)

# Convert network back to DataFrame and relabel nodes
df = pd.DataFrame(G.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)

# Create graph and adjacency matrix
G = nx.from_pandas_edgelist(df, source=0, target=1)
adj = nx.to_scipy_sparse_array(G, nodelist=range(len(G)))

# Split edges into train/test sets
u, v = np.nonzero(adj)
eids = np.random.permutation(len(u))
test_size = int(len(eids) * 0.1)
train_pos_u, test_pos_u = u[eids[test_size:]], u[eids[:test_size]]
train_pos_v, test_pos_v = v[eids[test_size:]], v[eids[:test_size]]

# Find negative edges and split them into train/test sets
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))
train_neg_u, test_neg_u = neg_u[neg_eids[test_size:]], neg_u[neg_eids[:test_size]]
train_neg_v, test_neg_v = neg_v[neg_eids[test_size:]], neg_v[neg_eids[:test_size]]

# Create DGL graph and features
g = dgl.graph((df[0].values, df[1].values))
g.ndata['feat'] = F.one_hot(g.nodes(), num_classes=g.number_of_nodes()).float()

# Remove edges in the test dataset
train_g = dgl.remove_edges(g, eids[:test_size])

# Create DGL graphs for training and testing
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())
all_neg_g = dgl.graph((neg_u, neg_v), num_nodes=g.number_of_nodes())

# Define 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

# Define MLP predictor
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):
        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']

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)

def create_network_train_test(df, net_name="twitter"):
    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_test = pd.concat([pos,neg])
    len(df_test)

    df_test.to_csv(f"{path_data}{net_name}_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, f"{path_data}{net_name}_network_prediction.graphml")

if 0: #don't run it
    np.random.seed(1546)
    # Create network for challenge (Twitter)
    df = pd.read_csv(f"{path_data}/ic2s2_netsci_3.tsv", sep="\t", usecols=["source", "target"])
    df.columns = [0,1]
    create_network_train_test(df, net_name="twitter")

    np.random.seed(1546)
    # Create network for challenge (PPI)
    df = pd.read_csv("http://interactome.dfci.harvard.edu/S_cerevisiae/download/CCSB-Y2H.txt",sep="\t",header=None)
    create_network_train_test(df, net_name="ppi")