# This version of the bitcoin experiment imports data preprocessed in Matlab, and uses the GCN baseline
# The point of this script is to do link prediction

# Imports and aliases
import pickle
import torch as t
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.datasets as datasets
import numpy as np
import matplotlib.pyplot as plt
import cProfile
import pandas as pd
import datetime
from scipy.sparse import csr_matrix
import os.path
import embedding_help_functions as ehf
import scipy.io as sio
unsq = t.unsqueeze
sq = t.squeeze

# Settings
alpha_vec = [.75, .76, .77, .78, .79, .80, .81, .82, .83, .84, .85, .86, .87, .88, .89, .90, .91, .92, .93, .94, .95]
no_layers = 1
dataset = "OTC" # OTC or Alpha
no_epochs = 1000
mat_f_name = "saved_content_bitcoin_otc.mat"
no_trials = 1
beta1 = 19
beta2 = 19
cutoff = 95
eval_type = "MAP-MRR" # "MAP-MRR" or "F1"

data_loc = "data/Bitcoin_" + dataset + "/"
S_train, S_val, S_test = 95, 20, 20
lr = 0.01
momentum = 0.9

# Load and return relevant data
A, A_labels, C_train, C_val, C_test, N = ehf.load_data(data_loc, mat_f_name, S_train, S_val, S_test, transformed=False)

# Create features for the nodes
X_train, X_val, X_test = ehf.create_node_features(A, S_train, S_val, S_test, same_block_size=False)

# Extract edges and labels from A_labels, and augment with nonexisting edges
# edges, beta
edges = A_labels._indices()
edges_aug, labels = ehf.augment_edges(edges, N, beta1, beta2, cutoff)

# Divide adjacency matrices and labels into training, validation and testing sets
edges_train, target_train, e_train, edges_val, target_val, e_val, edges_test, target_test, e_test = ehf.split_data(edges_aug, labels, S_train, S_val, S_test, same_block_size = False)

if no_trials > 1:
	ep_acc_loss_vec = []

for tr in range(no_trials):
	for alpha in alpha_vec:
		class_weights = t.tensor([alpha, 1.0-alpha])
		save_res_fname = "results_BASELINE_layers" + str(no_layers) + "_w" + str(round(float(class_weights[0])*100)) + "_" + dataset + "_link_prediction"

		# Create gcn for training
		if no_layers == 2:
			gcn = ehf.EmbeddingKWGCN(C_train[:-1], X_train[:-1], e_train, [6,6,2], nonlin2="selu")
		elif no_layers == 1:
			gcn = ehf.EmbeddingKWGCN(C_train[:-1], X_train[:-1], e_train, [6,2])

		# Train
		optimizer = t.optim.SGD(gcn.parameters(), lr=lr, momentum=momentum)
		criterion = nn.CrossEntropyLoss(weight=class_weights) # Takes arguments (output, target)
		if eval_type == "F1":
			ep_acc_loss = np.zeros((no_epochs,12)) # (precision_train, recall_train, f1_train, loss_train, precision_val, recall_val, f1_val, loss_val, precision_test, recall_test, f1_test, loss_test)
		elif eval_type == "MAP-MRR":
			ep_acc_loss = np.zeros((no_epochs,9)) # (MAP_train, MRR_train, loss_train, MAP_val, MRR_val, loss_val, MAP_test, MRR_test, loss_test)

		for ep in range(no_epochs):
			# Compute loss and take step
			optimizer.zero_grad()
			output_train = gcn()
			loss_train = criterion(output_train, target_train[edges_train[0]!=0])
			loss_train.backward()
			optimizer.step()

			# Things that don't require gradient
			with t.no_grad():
				if ep % 100 == 0:
					# Compute stats for training data; no point in doing more often than this
					guess_train = t.argmax(output_train, dim=1)
					if eval_type == "F1":
						precision_train, recall_train, f1_train = ehf.compute_f1(guess_train, target_train[edges_train[0]!=0])
					elif eval_type == "MAP-MRR":
						MAP_train, MRR_train = ehf.compute_MAP_MRR(output_train, target_train[edges_train[0]!=0], edges_train[:, edges_train[0]!=0])

					# Compute stats for validation data
					output_val = gcn(C_val[:-1], X_val[:-1], e_val)
					guess_val = t.argmax(output_val, dim=1)
					if eval_type == "F1":
						precision_val, recall_val, f1_val = ehf.compute_f1(guess_val, target_val[edges_val[0]!=0])
					elif eval_type == "MAP-MRR":
						MAP_val, MRR_val = ehf.compute_MAP_MRR(output_val, target_val[edges_val[0]!=0], edges_val[:, edges_val[0]!=0])
					loss_val = criterion(output_val, target_val[edges_val[0]!=0])
					
					# Compute stats for test data
					output_test = gcn(C_test[:-1], X_test[:-1], e_test)
					guess_test = t.argmax(output_test, dim=1)
					if eval_type == "F1":
						precision_test, recall_test, f1_test = ehf.compute_f1(guess_test, target_test[edges_test[0]!=0])
					elif eval_type == "MAP-MRR":
						MAP_test, MRR_test = ehf.compute_MAP_MRR(output_test, target_test[edges_test[0]!=0], edges_test[:, edges_test[0]!=0])
					loss_test = criterion(output_test, target_test[edges_test[0]!=0])

					# Print
					if eval_type == "F1":
						ehf.print_f1(precision_train, recall_train, f1_train, loss_train, precision_val, recall_val, f1_val, loss_val, precision_test, recall_test, f1_test, loss_test, alpha, tr, ep)
					elif eval_type == "MAP-MRR":
						print("alpha/Tr/Ep %.2f/%d/%d. Train MAP/MRR %.16f/%.16f. Train loss %.16f." % (alpha, tr, ep, MAP_train, MRR_train, loss_train))
						print("alpha/Tr/Ep %.2f/%d/%d. Val MAP/MRR %.16f/%.16f. Val loss %.16f." % (alpha, tr, ep, MAP_val, MRR_val, loss_val))
						print("alpha/Tr/Ep %.2f/%d/%d. Test MAP/MRR %.16f/%.16f. Test loss %.16f.\n" % (alpha, tr, ep, MAP_test, MRR_test, loss_test))

				# Store values with results
				if eval_type == "F1":							
					ep_acc_loss[ep] = [precision_train, recall_train, f1_train, loss_train, precision_val, recall_val, f1_val, loss_val, precision_test, recall_test, f1_test, loss_test]
				elif eval_type == "MAP-MRR":
					ep_acc_loss[ep] = [MAP_train, MRR_train, loss_train, MAP_val, MRR_val, loss_val, MAP_test, MRR_test, loss_test]

		if eval_type == "F1":
			ehf.print_f1(precision_train, recall_train, f1_train, loss_train, precision_val, recall_val, f1_val, loss_val, precision_test, recall_test, f1_test, loss_test, is_final=True)
		elif eval_type == "MAP-MRR":
			print("FINAL: Train MAP/MRR %.16f/%.16f. Train loss %.16f." % (MAP_train, MRR_train, loss_train))
			print("FINAL: Val MAP/MRR %.16f/%.16f. Val loss %.16f." % (MAP_val, MRR_val, loss_val))
			print("FINAL: Test MAP/MRR %.16f/%.16f. Test loss %.16f.\n" % (MAP_test, MRR_test, loss_test))

		if no_trials == 1:
			pickle.dump(ep_acc_loss, open(save_res_fname, "wb"))
			print("Results saved for single trial")
		else:
			ep_acc_loss_vec.append(ep_acc_loss)

if no_trials > 1:
	pickle.dump(ep_acc_loss_vec, open(save_res_fname + "_no_trials" + str(no_trials), "wb"))
	print("Results saved for all trials")
[((41, 41, 41, 119), 'embedding_help_functions.load_data', 'ehf.load_data', (), '', True, 'import embedding_help_functions as ehf\n'), ((44, 25, 44, 99), 'embedding_help_functions.create_node_features', 'ehf.create_node_features', (), '', True, 'import embedding_help_functions as ehf\n'), ((49, 20, 49, 69), 'embedding_help_functions.augment_edges', 'ehf.augment_edges', ({(49, 38, 49, 43): 'edges', (49, 45, 49, 46): 'N', (49, 48, 49, 53): 'beta1', (49, 55, 49, 60): 'beta2', (49, 62, 49, 68): 'cutoff'}, {}), '(edges, N, beta1, beta2, cutoff)', True, 'import embedding_help_functions as ehf\n'), ((52, 100, 52, 182), 'embedding_help_functions.split_data', 'ehf.split_data', (), '', True, 'import embedding_help_functions as ehf\n'), ((59, 18, 59, 46), 'torch.tensor', 't.tensor', ({(59, 27, 59, 45): '[alpha, 1.0 - alpha]'}, {}), '([alpha, 1.0 - alpha])', True, 'import torch as t\n'), ((70, 14, 70, 55), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', (), '', True, 'import torch.nn as nn\n'), ((64, 9, 64, 89), 'embedding_help_functions.EmbeddingKWGCN', 'ehf.EmbeddingKWGCN', (), '', True, 'import embedding_help_functions as ehf\n'), ((72, 17, 72, 41), 'numpy.zeros', 'np.zeros', ({(72, 26, 72, 40): '(no_epochs, 12)'}, {}), '((no_epochs, 12))', True, 'import numpy as np\n'), ((127, 3, 127, 177), 'embedding_help_functions.print_f1', 'ehf.print_f1', (), '', True, 'import embedding_help_functions as ehf\n'), ((66, 9, 66, 71), 'embedding_help_functions.EmbeddingKWGCN', 'ehf.EmbeddingKWGCN', ({(66, 28, 66, 40): 'C_train[:-1]', (66, 42, 66, 54): 'X_train[:-1]', (66, 56, 66, 63): 'e_train', (66, 65, 66, 70): '[6, 2]'}, {}), '(C_train[:-1], X_train[:-1], e_train, [6, 2])', True, 'import embedding_help_functions as ehf\n'), ((74, 17, 74, 40), 'numpy.zeros', 'np.zeros', ({(74, 26, 74, 39): '(no_epochs, 9)'}, {}), '((no_epochs, 9))', True, 'import numpy as np\n'), ((85, 8, 85, 19), 'torch.no_grad', 't.no_grad', ({}, {}), '()', True, 'import torch as t\n'), ((88, 19, 88, 48), 'torch.argmax', 't.argmax', (), '', True, 'import torch as t\n'), ((96, 17, 96, 44), 'torch.argmax', 't.argmax', (), '', True, 'import torch as t\n'), ((105, 18, 105, 46), 'torch.argmax', 't.argmax', (), '', True, 'import torch as t\n'), ((90, 48, 90, 108), 'embedding_help_functions.compute_f1', 'ehf.compute_f1', ({(90, 63, 90, 74): 'guess_train', (90, 76, 90, 107): 'target_train[edges_train[0] != 0]'}, {}), '(guess_train, target_train[edges_train[0] != 0])', True, 'import embedding_help_functions as ehf\n'), ((98, 42, 98, 96), 'embedding_help_functions.compute_f1', 'ehf.compute_f1', ({(98, 57, 98, 66): 'guess_val', (98, 68, 98, 95): 'target_val[edges_val[0] != 0]'}, {}), '(guess_val, target_val[edges_val[0] != 0])', True, 'import embedding_help_functions as ehf\n'), ((107, 45, 107, 102), 'embedding_help_functions.compute_f1', 'ehf.compute_f1', ({(107, 60, 107, 70): 'guess_test', (107, 72, 107, 101): 'target_test[edges_test[0] != 0]'}, {}), '(guess_test, target_test[edges_test[0] != 0])', True, 'import embedding_help_functions as ehf\n'), ((114, 6, 114, 180), 'embedding_help_functions.print_f1', 'ehf.print_f1', ({(114, 19, 114, 34): 'precision_train', (114, 36, 114, 48): 'recall_train', (114, 50, 114, 58): 'f1_train', (114, 60, 114, 70): 'loss_train', (114, 72, 114, 85): 'precision_val', (114, 87, 114, 97): 'recall_val', (114, 99, 114, 105): 'f1_val', (114, 107, 114, 115): 'loss_val', (114, 117, 114, 131): 'precision_test', (114, 133, 114, 144): 'recall_test', (114, 146, 114, 153): 'f1_test', (114, 155, 114, 164): 'loss_test', (114, 166, 114, 171): 'alpha', (114, 173, 114, 175): 'tr', (114, 177, 114, 179): 'ep'}, {}), '(precision_train, recall_train, f1_train, loss_train,\n    precision_val, recall_val, f1_val, loss_val, precision_test,\n    recall_test, f1_test, loss_test, alpha, tr, ep)', True, 'import embedding_help_functions as ehf\n'), ((92, 29, 92, 130), 'embedding_help_functions.compute_MAP_MRR', 'ehf.compute_MAP_MRR', ({(92, 49, 92, 61): 'output_train', (92, 63, 92, 94): 'target_train[edges_train[0] != 0]', (92, 96, 92, 129): 'edges_train[:, (edges_train[0] != 0)]'}, {}), '(output_train, target_train[edges_train[0] != 0],\n    edges_train[:, (edges_train[0] != 0)])', True, 'import embedding_help_functions as ehf\n'), ((100, 25, 100, 116), 'embedding_help_functions.compute_MAP_MRR', 'ehf.compute_MAP_MRR', ({(100, 45, 100, 55): 'output_val', (100, 57, 100, 84): 'target_val[edges_val[0] != 0]', (100, 86, 100, 115): 'edges_val[:, (edges_val[0] != 0)]'}, {}), '(output_val, target_val[edges_val[0] != 0], edges_val[:,\n    (edges_val[0] != 0)])', True, 'import embedding_help_functions as ehf\n'), ((109, 27, 109, 123), 'embedding_help_functions.compute_MAP_MRR', 'ehf.compute_MAP_MRR', ({(109, 47, 109, 58): 'output_test', (109, 60, 109, 89): 'target_test[edges_test[0] != 0]', (109, 91, 109, 122): 'edges_test[:, (edges_test[0] != 0)]'}, {}), '(output_test, target_test[edges_test[0] != 0],\n    edges_test[:, (edges_test[0] != 0)])', True, 'import embedding_help_functions as ehf\n')]