import math
import torch
import torch.nn as nn
import torch.nn.init as init
import numpy as np
import pandas as pd
import networkx as nx
import scipy as sp
import seaborn as sns
# from node2vec import Node2Vec
from sklearn.decomposition import PCA
import copy
import time

if torch.cuda.is_available():
  torch.device('cuda')

"""Utils:
Data Loader / Attention / Clones / Embedder"""

def Graph_load_batch(min_num_nodes=20, max_num_nodes=1000, name='ENZYMES', node_attributes=True, graph_labels=True):
    '''
    load many graphs, e.g. enzymes
    :return: a list of graphs
    '''
    print('Loading graph dataset: ' + str(name))
    G = nx.Graph()
    # load data
    # path = '../dataset/' + name + '/'
    path = '/content/gdrive/My Drive/' + name + '/'
    data_adj = np.loadtxt(path + name + '_A.txt', delimiter=',').astype(int)
    if node_attributes:
        data_node_att = np.loadtxt(path + name + '_node_attributes.txt', delimiter=',')
    data_node_label = np.loadtxt(path + name + '_node_labels.txt', delimiter=',').astype(int)
    data_graph_indicator = np.loadtxt(path + name + '_graph_indicator.txt', delimiter=',').astype(int)
    if graph_labels:
        data_graph_labels = np.loadtxt(path + name + '_graph_labels.txt', delimiter=',').astype(int)
    data_tuple = list(map(tuple, data_adj))
    G.add_edges_from(data_tuple)
    for i in range(data_node_label.shape[0]):
        if node_attributes:
            G.add_node(i + 1, feature=data_node_att[i])
        G.add_node(i + 1, label=data_node_label[i])
    G.remove_nodes_from(list(nx.isolates(G)))
    graph_num = data_graph_indicator.max()
    node_list = np.arange(data_graph_indicator.shape[0]) + 1
    graphs = []
    max_nodes = 0
    for i in range(graph_num):
        nodes = node_list[data_graph_indicator == i + 1]
        G_sub = G.subgraph(nodes)
        if graph_labels:
            G_sub.graph['label'] = data_graph_labels[i]

        if G_sub.number_of_nodes() >= min_num_nodes and G_sub.number_of_nodes() <= max_num_nodes:
            graphs.append(G_sub)
            if G_sub.number_of_nodes() > max_nodes:
                max_nodes = G_sub.number_of_nodes()

    print('Loaded')
    return graphs

def attention(query, key, value, d_key):
  scores = torch.matmul(query, key.transpose(-2, -1)) /  math.sqrt(d_key)
  output = torch.matmul(scores, value)
  output = nn.functional.softmax(output)
  return output

def get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

def embedder(graph, dimensions=32, walk_length=8, num_walks=200, workers=4):
  node2vec = Node2Vec(graph, dimensions=dimensions, walk_length=walk_length, num_walks=num_walks, workers=workers)  # Use temp_folder for big graphs
  model = node2vec.fit(window=10, min_count=1, batch_words=4)
  return model.wv.vectors

graphs = Graph_load_batch(min_num_nodes=10, name='ENZYMES')

# G = graphs[1]
# vecs = embedder(G)

# pca = PCA(n_components=2)
# principalComponents = pca.fit_transform(vecs)
# principalDf = pd.DataFrame(data = principalComponents
#              , columns = ['principal component 1', 'principal component 2'])
# principalDf.index = list(G.nodes())

# sns.scatterplot(principalDf['principal component 1'], principalDf['principal component 2'])

"""Sublayers"""

class MultiHeadAttention(nn.Module):
  def __init__(self, heads, d_model, dropout = 0.1):
    super().__init__()
    self.d_model = d_model
    self.d_k = d_model // heads
    self.h = heads
    self.q_linear = nn.Linear(d_model, d_model).cuda()
    self.v_linear = nn.Linear(d_model, d_model).cuda()
    self.k_linear = nn.Linear(d_model, d_model).cuda()
    self.dropout = nn.Dropout(dropout)
    self.out = nn.Linear(d_model, d_model)

  def forward(self, q, k, v):
    # print(q, k, v)
    bs = q.size(0)
    # perform linear operation and split into h heads
    k = self.k_linear(k).view(bs, -1, self.h, self.d_k)
    q = self.q_linear(q).view(bs, -1, self.h, self.d_k)
    v = self.v_linear(v).view(bs, -1, self.h, self.d_k)
    # transpose to get dimensions bs * h * sl * d_model
    k = k.transpose(1,2)
    q = q.transpose(1,2)
    v = v.transpose(1,2)

    scores = attention(q, k, v, self.d_k)
    # concatenate heads and put through final linear layer
    concat = scores.transpose(1,2).contiguous().view(bs, -1, self.d_model)
    output = self.out(concat)

    return output

class FeedForward(nn.Module):
  def __init__(self, d_model, d_ff=2048, dropout = 0.1):
    super().__init__()
    self.linear_1 = nn.Linear(d_model, d_ff).cuda()
    self.dropout = nn.Dropout(dropout)
    self.linear_2 = nn.Linear(d_ff, d_model).cuda()
    
  def forward(self, x):
    x = self.dropout(nn.functional.relu(self.linear_1(x)))
    x = self.linear_2(x)
    return x


class Norm(nn.Module):
  def __init__(self, d_model, eps = 1e-6):
    super().__init__()
    self.size = d_model
    self.alpha = nn.Parameter(torch.ones(self.size))
    self.bias = nn.Parameter(torch.zeros(self.size))
    self.eps = eps
    
  def forward(self, x):
    norm = self.alpha * (x - x.mean(dim=-1, keepdim=True)) / (x.std(dim=-1, keepdim=True) + self.eps) + self.bias
    return norm

"""Layers"""

class EncoderLayer(nn.Module):
  def __init__(self, d_model, heads, dropout = 0.1):
    super().__init__()
    self.norm_1 = Norm(d_model)
    self.norm_2 = Norm(d_model)
    self.attn = MultiHeadAttention(heads, d_model)
    self.ff = FeedForward(d_model)
    self.dropout_1 = nn.Dropout(dropout)
    self.dropout_2 = nn.Dropout(dropout)

  def forward(self, x):
    # x2 = self.norm_1(x)
    x = x + self.dropout_1(self.attn(x,x,x))
    # x2 = self.norm_2(x)
    x = x + self.dropout_2(self.ff(x))
    return x

class DecoderLayer(nn.Module):
  def __init__(self, d_model, heads, dropout=0.1):
    super().__init__()
    self.norm_1 = Norm(d_model)
    self.norm_2 = Norm(d_model)
    self.norm_3 = Norm(d_model)
    self.dropout_1 = nn.Dropout(dropout)
    self.dropout_2 = nn.Dropout(dropout)
    self.dropout_3 = nn.Dropout(dropout)
    self.attn_1 = MultiHeadAttention(heads, d_model)
    self.attn_2 = MultiHeadAttention(heads, d_model)
    self.ff = FeedForward(d_model).cuda()

  def forward(self, x, e_outputs):
    # x2 = self.norm_1(x)
    x = x + self.dropout_1(self.attn_1(x, x, x))
    # x2 = self.norm_2(x)
    # x2 = self.norm_2(x)
    x = x + self.dropout_2(self.attn_2(x, e_outputs, e_outputs))
    # x2 = self.norm_3(x)
    x = x + self.dropout_3(self.ff(x))
    return x

class Encoder(nn.Module):
  def __init__(self, vocab_size, d_model, N, heads):
    super().__init__()
    self.N = N
    self.layers = get_clones(EncoderLayer(d_model, heads), N)
    self.norm = Norm(d_model)
      
  def forward(self, src):
    x = src
    for i in range(N):
      x = self.layers[i](x)
    return self.norm(x)

class Decoder(nn.Module):
  def __init__(self, data_size, d_model, N, heads):
    super().__init__()
    self.N = N
    self.layers = get_clones(DecoderLayer(d_model, heads), N)
    self.norm = Norm(d_model)

  def forward(self, trg, e_outputs):
    x = trg
    for i in range(self.N):
      x = self.layers[i](x, e_outputs)
    return self.norm(x)

"""The Mighty Transformer"""

class Transformer(nn.Module):
  def __init__(self, src_graph, trg_graph, d_model, N, heads):
      super().__init__()
      self.encoder = Encoder(src_graph, d_model, N, heads)
      self.decoder = Decoder(trg_graph, d_model, N, heads)
      self.out = nn.Linear(d_model, trg_graph)

  def forward(self, src, trg):
      e_outputs = self.encoder(src)
      d_output = self.decoder(trg, e_outputs)
      output = self.out(d_output)
      return output


def remove_random_node(graph, max_size=40, min_size=10):
  if len(graph.nodes) >= max_size or len(graph.nodes) < min_size:
    return None
  relabeled_graph = nx.relabel.convert_node_labels_to_integers(graph)
  choice = np.random.choice(list(relabeled_graph.nodes))
  remaining_graph = nx.to_numpy_matrix(relabeled_graph.subgraph(filter(lambda x: x != choice, list(relabeled_graph.nodes))))
  removed_node = nx.to_numpy_matrix(relabeled_graph)[choice]
  graph_length = len(remaining_graph)
  source_graph = np.pad(remaining_graph, [(0, max_size - graph_length), (0, max_size - graph_length)])
  target_graph = np.copy(source_graph)
  removed_node_row = np.asarray(removed_node)[0]
  target_graph[graph_length] = np.pad(removed_node_row, [(0, max_size - len(removed_node_row))])
  return source_graph, target_graph

converted_graphs = list(filter(lambda x: x is not None, [remove_random_node(graph) for graph in graphs]))
source_graphs = torch.Tensor([graph[0] for graph in converted_graphs])
target_graphs = torch.Tensor([graph[1] for graph in converted_graphs])



d_model = 40
heads = 8
N = 6
src_size = len(source_graphs)
trg_size = len(target_graphs)

model = Transformer(src_size, trg_size, d_model, N, heads).cuda()

#print(model)

optim = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)


def train_model(epoch, print_every=100):
  model.train()
  start = time.time()
  temp = start
  total_loss = 0
  for i in range(epoch):

    src = source_graphs.cuda()
    trg = target_graphs.cuda()

    preds = model(src.float(), trg.float())
    optim.zero_grad()
    loss = torch.nn.functional.cross_entropy(preds.view(preds.size(-1), -1), trg.view(trg.size(0), -1))
    loss.backward()
    optim.step()
    total_loss += loss.data[0]
    if (i + 1) % print_every == 0:
      loss_avg = total_loss / print_every
      print("time = %dm, epoch %d, iter = %d, loss = %.3f,\
#       %ds per %d iters" % ((time.time() - start) // 60,\
      epoch + 1, i + 1, loss_avg, time.time() - temp,\
      print_every))
      total_loss = 0
      temp = time.time()

train_model(1, 1)

#preds = model(source_graphs[0].cuda(), target_graphs[0].cuda())
#loss = torch.nn.functional.cross_entropy(preds.view(preds.size(-1), -1), target_graphs.view(target_graphs.size(0), -1))
#