"""
Deep Learning on Graphs - ALTEGRAD - Dec 2019
"""

import networkx as nx
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.manifold import SpectralEmbedding
from sklearn.metrics import accuracy_score
from deepwalk import deepwalk
import matplotlib.pyplot as plt


# Loads the karate network
G = nx.read_weighted_edgelist('../data/karate.edgelist', delimiter=' ', nodetype=int, create_using=nx.Graph())
print("Number of nodes:", G.number_of_nodes())
print("Number of edges:", G.number_of_edges())

n = G.number_of_nodes()

# Loads the class labels
class_labels = np.loadtxt('../data/karate_labels.txt', delimiter=',', dtype=np.int32)
idx_to_class_label = dict()
for i in range(class_labels.shape[0]):
    idx_to_class_label[class_labels[i,0]] = class_labels[i,1]

y = list()
for node in G.nodes():
    y.append(idx_to_class_label[node])

y = np.array(y)



############## Task 5

##################
# your code here #


f = plt.figure()
layout = nx.fruchterman_reingold_layout(G)
nx.draw_networkx(G, pos = layout, with_labels = True, node_color = y, cmap = plt.cm.coolwarm, vmin = 0, vmax = 1)
plt.title('Karate network')
plt.show()
f.savefig("karate_classification_graph.pdf")


##################



############## Task 6
n_dim = 128
n_walks = 10 
walk_length = 20 
model = deepwalk(G, n_walks, walk_length, n_dim) # your code here

embeddings = np.zeros((n, n_dim))
for i, node in enumerate(G.nodes()):
    embeddings[i,:] = model.wv[str(node)]

idx = np.random.RandomState(seed=42).permutation(n)
idx_train = idx[:int(0.8*n)]
idx_test = idx[int(0.8*n):]

X_train = embeddings[idx_train,:]
X_test = embeddings[idx_test,:]

y_train = y[idx_train]
y_test = y[idx_test]



############## Task 7

##################
# your code here #

model_lr = LogisticRegression().fit(X_train, y_train)
y_pred_lr = model_lr.predict(X_test)

accuracy_lr = accuracy_score(y_test, y_pred_lr, normalize=False)
print('Accuracy using deepwalk for embeddings: ', accuracy_lr, ' correct out of ', len(y_test))


##################




############## Task 8

##################
# your code here #

spectralEmbedding = SpectralEmbedding(n_components = n_dim, affinity='precomputed').fit(nx.convert_matrix.to_numpy_matrix(G))
embeddings_se = spectralEmbedding.embedding_


X_train = embeddings_se[idx_train,:]
X_test = embeddings_se[idx_test,:]


model_lr_se = LogisticRegression().fit(X_train, y_train)
y_pred_lr_se = model_lr_se.predict(X_test)

accuracy_lr_se = accuracy_score(y_test, y_pred_lr_se, normalize=False)

print('Accuracy using spectral embedding: ', accuracy_lr_se, ' correct out of ', len(y_test))

##################
