# 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")
[((1080, 1158), 'embedding_help_functions.load_data', 'ehf.load_data', (['data_loc', 'mat_f_name', 'S_train', 'S_val', 'S_test'], {'transformed': '(False)'}), '(data_loc, mat_f_name, S_train, S_val, S_test, transformed=False)\n', (1093, 1158), True, 'import embedding_help_functions as ehf\n'), ((1217, 1291), 'embedding_help_functions.create_node_features', 'ehf.create_node_features', (['A', 'S_train', 'S_val', 'S_test'], {'same_block_size': '(False)'}), '(A, S_train, S_val, S_test, same_block_size=False)\n', (1241, 1291), True, 'import embedding_help_functions as ehf\n'), ((1432, 1481), 'embedding_help_functions.augment_edges', 'ehf.augment_edges', (['edges', 'N', 'beta1', 'beta2', 'cutoff'], {}), '(edges, N, beta1, beta2, cutoff)\n', (1449, 1481), True, 'import embedding_help_functions as ehf\n'), ((1665, 1750), 'embedding_help_functions.split_data', 'ehf.split_data', (['edges_aug', 'labels', 'S_train', 'S_val', 'S_test'], {'same_block_size': '(False)'}), '(edges_aug, labels, S_train, S_val, S_test, same_block_size=False\n    )\n', (1679, 1750), True, 'import embedding_help_functions as ehf\n'), ((1861, 1891), 'torch.tensor', 't.tensor', (['[alpha, 1.0 - alpha]'], {}), '([alpha, 1.0 - alpha])\n', (1869, 1891), True, 'import torch as t\n'), ((2368, 2409), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ([], {'weight': 'class_weights'}), '(weight=class_weights)\n', (2387, 2409), True, 'import torch.nn as nn\n'), ((2097, 2184), 'embedding_help_functions.EmbeddingKWGCN', 'ehf.EmbeddingKWGCN', (['C_train[:-1]', 'X_train[:-1]', 'e_train', '[6, 6, 2]'], {'nonlin2': '"""selu"""'}), "(C_train[:-1], X_train[:-1], e_train, [6, 6, 2], nonlin2=\n    'selu')\n", (2115, 2184), True, 'import embedding_help_functions as ehf\n'), ((2486, 2511), 'numpy.zeros', 'np.zeros', (['(no_epochs, 12)'], {}), '((no_epochs, 12))\n', (2494, 2511), True, 'import numpy as np\n'), ((5624, 5806), 'embedding_help_functions.print_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)'}), '(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, is_final=True)\n', (5636, 5806), True, 'import embedding_help_functions as ehf\n'), ((2210, 2273), 'embedding_help_functions.EmbeddingKWGCN', 'ehf.EmbeddingKWGCN', (['C_train[:-1]', 'X_train[:-1]', 'e_train', '[6, 2]'], {}), '(C_train[:-1], X_train[:-1], e_train, [6, 2])\n', (2228, 2273), True, 'import embedding_help_functions as ehf\n'), ((2709, 2733), 'numpy.zeros', 'np.zeros', (['(no_epochs, 9)'], {}), '((no_epochs, 9))\n', (2717, 2733), True, 'import numpy as np\n'), ((3108, 3119), 'torch.no_grad', 't.no_grad', ([], {}), '()\n', (3117, 3119), True, 'import torch as t\n'), ((3241, 3270), 'torch.argmax', 't.argmax', (['output_train'], {'dim': '(1)'}), '(output_train, dim=1)\n', (3249, 3270), True, 'import torch as t\n'), ((3684, 3711), 'torch.argmax', 't.argmax', (['output_val'], {'dim': '(1)'}), '(output_val, dim=1)\n', (3692, 3711), True, 'import torch as t\n'), ((4170, 4198), 'torch.argmax', 't.argmax', (['output_test'], {'dim': '(1)'}), '(output_test, dim=1)\n', (4178, 4198), True, 'import torch as t\n'), ((3346, 3408), 'embedding_help_functions.compute_f1', 'ehf.compute_f1', (['guess_train', 'target_train[edges_train[0] != 0]'], {}), '(guess_train, target_train[edges_train[0] != 0])\n', (3360, 3408), True, 'import embedding_help_functions as ehf\n'), ((3781, 3837), 'embedding_help_functions.compute_f1', 'ehf.compute_f1', (['guess_val', 'target_val[edges_val[0] != 0]'], {}), '(guess_val, target_val[edges_val[0] != 0])\n', (3795, 3837), True, 'import embedding_help_functions as ehf\n'), ((4271, 4330), 'embedding_help_functions.compute_f1', 'ehf.compute_f1', (['guess_test', 'target_test[edges_test[0] != 0]'], {}), '(guess_test, target_test[edges_test[0] != 0])\n', (4285, 4330), True, 'import embedding_help_functions as ehf\n'), ((4605, 4787), 'embedding_help_functions.print_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'], {}), '(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)\n', (4617, 4787), True, 'import embedding_help_functions as ehf\n'), ((3470, 3581), 'embedding_help_functions.compute_MAP_MRR', 'ehf.compute_MAP_MRR', (['output_train', 'target_train[edges_train[0] != 0]', 'edges_train[:, (edges_train[0] != 0)]'], {}), '(output_train, target_train[edges_train[0] != 0],\n    edges_train[:, (edges_train[0] != 0)])\n', (3489, 3581), True, 'import embedding_help_functions as ehf\n'), ((3895, 3996), 'embedding_help_functions.compute_MAP_MRR', 'ehf.compute_MAP_MRR', (['output_val', 'target_val[edges_val[0] != 0]', 'edges_val[:, (edges_val[0] != 0)]'], {}), '(output_val, target_val[edges_val[0] != 0], edges_val[:,\n    (edges_val[0] != 0)])\n', (3914, 3996), True, 'import embedding_help_functions as ehf\n'), ((4390, 4496), 'embedding_help_functions.compute_MAP_MRR', 'ehf.compute_MAP_MRR', (['output_test', 'target_test[edges_test[0] != 0]', 'edges_test[:, (edges_test[0] != 0)]'], {}), '(output_test, target_test[edges_test[0] != 0],\n    edges_test[:, (edges_test[0] != 0)])\n', (4409, 4496), True, 'import embedding_help_functions as ehf\n')]