import networkx as nx
import numpy as np
import scipy as sp
import scipy.cluster.vq as vq
import matplotlib.pyplot as plt
import math
import random
import operator
import time

################################################################################
import torch
import torch_geometric
import torch.optim as optim
from torch_geometric.nn import GCNConv
from torch_geometric.utils import from_networkx
from torch_geometric.datasets import Planetoid
from torch_geometric.utils import to_networkx
################################################################################
# Read in data and construct graph

def construct_graph():
  counter = 1
  label_lgenre = dict()
  with open("ent.discogs_lgenre_lgenre.label.name.data") as f:
    for line in f:
      label_lgenre[counter] = line.rstrip()
      counter += 1
  
  counter = 1
  genre_lgenre = dict()
  with open("ent.discogs_lgenre_lgenre.genre.name.data") as f:
    for line in f:
      genre_lgenre[counter] = line.rstrip()
      counter += 1
  
  counter = 1
  label_lstyle = dict()
  with open("ent.discogs_lstyle_lstyle.label.name.data") as f:
    for line in f:
      label_lstyle[counter] = line.rstrip()
      counter += 1
  
  counter = 1
  style_lstyle = dict()
  with open("ent.discogs_lstyle_lstyle.style.name.data") as f:
    for line in f:
      style_lstyle[counter] = line.rstrip()
      counter += 1
  
  G_lgenre = nx.Graph()
  G_lstyle = nx.Graph()
  
  with open("out.discogs_lgenre_lgenre.data") as f:
    for line in f:
      if line[0] == '%':
        continue
      (v1, v2) = line.split()
      G_lgenre.add_edge(label_lgenre[int(v1)], genre_lgenre[int(v2)])
  
  with open("out.discogs_lstyle_lstyle.data") as f:
    for line in f:
      if line[0] == '%':
        continue
      (v1, v2) = line.split()
      G_lstyle.add_edge(label_lstyle[int(v1)], style_lstyle[int(v2)])
  
  G_actual = nx.Graph()
  for v in style_lstyle.values():
    Nv = list()
    for u in G_lstyle.neighbors(v):
      Nv.append(u)
    
    # Construct an E-R graph among all artists with the same style. Thus, our
    # graph will consist of overlapping communities of just artists.
    ratio = 0.001
    num_edges = int(ratio*len(Nv)*(len(Nv)-1)*0.5)
    for i in range(0, num_edges):
      x = random.choice(Nv)
      y = random.choice(Nv)
      G_actual.add_edge(x, y)
  
  vals = label_lgenre.values()
  ground_truth = {}
  for v in list(G_actual.nodes()):
    if v not in vals:
      G_actual.remove_node(v)
    else:
      counts = {}
      for u in G_lgenre.neighbors(v):
        # Ignoring the below since they're extremely pervasive and therefore
        # not that interesting
        if u == 'Electronic' or u == 'Rock':
          continue
        if u not in dict.keys(counts):
          counts[u] = 1
        else:
          counts[u] += 1
      if len(counts.values()) == 0:
        G_actual.remove_node(v)
      else:
        # We'll initially label each artist with the most likely genre among
        # their neighbors in the original artist-genre graph.
        c = np.random.choice([k for k in counts.keys() if counts[k]==max(counts.values())])
        ground_truth[v] = c
  
  #G = G_actual.subgraph(sorted(nx.connected_components(G_actual), key=len, reverse=True)[0])
  #G = G_actual
  return G_actual, ground_truth

################################################################################
# construct graph from pytorch_geometric datasets

def construct_graph_from_pytorch(which):
  if which == "Cora":
    dataset = Planetoid(root='/tmp/Cora', name='Cora')
  elif which == "CiteSeer":
    dataset = Planetoid(root='/tmp/CiteSeer', name='CiteSeer')
  elif which == "PubMed":
    dataset = Planetoid(root='/tmp/PubMed', name='PubMed')
  data = dataset[0]
  G = to_networkx(data, node_attrs=['y'], to_undirected=True)
  gt = {}
  for v in G.nodes():
    gt[v] = G.nodes[v]['y'] 
  return G, gt
  
################################################################################
# method 1 - simple label propagation

def label_prop(G, labels, unlabeled):
  updates = len(unlabeled)
  max_acc = 0.0
  while updates > len(unlabeled)*0.1:
    updates = 0
    for v in sorted(unlabeled, key=lambda k: random.random()):
      counts = {}
      for u in G.neighbors(v):
        if labels[u] == "NA":
          continue
        elif labels[u] not in dict.keys(counts):
          counts[labels[u]] = 1
        else:
          counts[labels[u]] += 1
      if len(counts.values()) > 0:
        c = np.random.choice([k for k in counts.keys() if counts[k]==max(counts.values())])
        if c != labels[v]:
          labels[v] = c
          updates += 1
    test_acc = eval_accuracy(G, gt, labels, unlabeled)
    if test_acc > max_acc:
      max_acc = test_acc
  return max_acc


################################################################################
# method 2 - iterative naive bayes

def init_features(G, labels):
  feat_idx = {}
  counter = 0
  for l in labels.values():
    if l not in feat_idx:
      feat_idx[l] = counter
      counter += 1
  features = {}
  for v in G.nodes():
    features[v] = [0.0]*len(feat_idx)
    update_features(G, v, features, feat_idx, labels)
  return features, feat_idx

def update_features(G, v, features, feat_idx, labels):
  Nv_size = 0.0
  for i in range(0, len(feat_idx)):
    features[v][i] = 0.0
  for u in G.neighbors(v):
    if labels[u] != "NA":
      features[v][feat_idx[labels[u]]] += 1.0
      Nv_size += 1.0
  if Nv_size > 0.0:
    for i in range(0, len(feat_idx)):
      features[v][i] /= Nv_size

def init_bayes(G, features, feat_idx, labels):
  C = [0.0]*len(feat_idx)
  count = 0
  for v in G.nodes():
    if labels[v] != "NA":
      C[feat_idx[labels[v]]] += 1
      count += 1
  for i in range(0, len(feat_idx)):
    C[i] /= count
  
  avgs = dict()
  stds = dict()
  counts = dict()
  for c in feat_idx.keys():
    avgs[c] = [0.0]*len(feat_idx)
    stds[c] = [0.0]*len(feat_idx)
    counts[c] = [0.0]*len(feat_idx)
  
  for v in G.nodes():
    if labels[v] != "NA":
      c = labels[v]
      for i in range(0, len(feat_idx)):
        avgs[c][i] += features[v][i]
        counts[c][i] += 1.0
  for c in feat_idx.keys():
    for i in range(0, len(feat_idx)):
      if counts[c][i] > 0.0:
        avgs[c][i] /= counts[c][i]
      else:
        avgs[c][i] = 0.0
  
  for v in G.nodes():
    if labels[v] != "NA":
      c = labels[v]
      for i in range(0, len(feat_idx)):
        stds[c][i] += (features[v][i] - avgs[c][i])**2
  for c in feat_idx.keys():
    for i in range(0, len(feat_idx)):
      if counts[c][i] > 0.0:
        stds[c][i] /= counts[c][i]
        stds[c][i] = math.sqrt(avgs[c][i])
      else:
        stds[c][i] = 0.0
  
  X = (avgs, stds)
  return C, X

def calc_prob(val, avg, std):
  exponent = math.exp(-((val - avg)**2 / (2 * std**2)))
  return (1 / (math.sqrt(2 * math.pi) * std)) * exponent

def update_label(G, v, features, feat_idx, C, X):
  max_c = "NA"
  max_prob = 0.0
  for c in feat_idx.keys():
    prob = C[feat_idx[c]]
    for i in range(0, len(feat_idx)):
      avg = X[0][c][i]
      std = X[1][c][i]
      if avg > 0.0 and std > 0.0:
        prob *= calc_prob(features[v][i], avg, std)
    if prob > max_prob:
      max_prob = prob
      max_c = c
  return max_c

def naive_bayes(G, labels, unlabeled):
  features, feat_idx = init_features(G, labels)
  C, X = init_bayes(G, features, feat_idx, labels)
  updates = len(unlabeled)
  max_acc = 0.0
  while updates > len(unlabeled)*0.01:
    updates = 0
    for v in unlabeled:
      update_features(G, v, features, feat_idx, labels)
      new_label = update_label(G, v, features, feat_idx, C, X)
      if new_label != labels[v]:
        labels[v] = new_label
        updates += 1
    test_acc = eval_accuracy(G, gt, labels, unlabeled)
    if test_acc > max_acc:
      max_acc = test_acc
  return max_acc

################################################################################
# random walk approach

def construct_matrices(G, labels):
  ids = {}
  counter = 0
  for v in G.nodes():
    if labels[v] != "NA":
      ids[v] = counter
      counter += 1
  
  num_labeled = counter
  
  counter = 0
  for v in G.nodes():
    if labels[v] == "NA":
      ids[v] = counter
      counter += 1
  
  num_unlabeled = counter
  
  rows_P_ul = []
  cols_P_ul = []
  vals_P_ul = []
  for v in G.nodes():
    if labels[v] != "NA":
      continue
    for u in G.neighbors(v):
      if labels[u] != "NA":
        rows_P_ul.append(ids[v])
        cols_P_ul.append(ids[u])
        vals_P_ul.append(1.0 / G.degree(v))
  
  rows_P_uu = []
  cols_P_uu = []
  vals_P_uu = []
  for v in G.nodes():
    if labels[v] != "NA":
      continue
    for u in G.neighbors(v):
      if labels[u] == "NA":
        rows_P_uu.append(ids[v])
        cols_P_uu.append(ids[u])
        vals_P_uu.append(1 / G.degree(v))
  
  feat_idx = {}
  counter = 0
  for l in labels.values():
    if l not in feat_idx:
      feat_idx[l] = counter
      counter += 1
  
  num_feat = len(feat_idx)
  
  rows_Y_l = []
  cols_Y_l = []
  vals_Y_l = []
  for v in G.nodes():
    if labels[v] != "NA":
      rows_Y_l.append(ids[v])
      cols_Y_l.append(feat_idx[labels[v]])
      vals_Y_l.append(1.0)
  
  rows_Y_u = []
  cols_Y_u = []
  vals_Y_u = []
  for v in G.nodes():
    if labels[v] == "NA":
      for i in range(0, num_feat):
        rows_Y_u.append(ids[v])
        cols_Y_u.append(1)
        vals_Y_u.append(1.0 / num_feat)
  
  P_ul = sp.sparse.csr_matrix((vals_P_ul, (rows_P_ul, cols_P_ul)), shape=(num_unlabeled, num_labeled))
  P_uu = sp.sparse.csr_matrix((vals_P_uu, (rows_P_uu, cols_P_uu)), shape=(num_unlabeled, num_unlabeled))
  Y_l = sp.sparse.csr_matrix((vals_Y_l, (rows_Y_l, cols_Y_l)), shape=(num_labeled, num_feat))
  Y_u = sp.sparse.csr_matrix((vals_Y_u, (rows_Y_u, cols_Y_u)), shape=(num_unlabeled, num_feat))
  
  return P_ul, P_uu, Y_l, Y_u, ids, feat_idx

def determine_classes(G, labels, unlabeled, feat_idx, ids, Y_u):
  for v in unlabeled:
    max_c = "NA"
    max_Y = 0.0
    for c in feat_idx.keys():
      if Y_u[ids[v], feat_idx[c]] > max_Y:
        max_Y = Y_u[ids[v], feat_idx[c]]
        max_c = c
    labels[v] = max_c
  return labels

def random_walk(G, labels, unlabeled):
  P_ul, P_uu, Y_l, Y_u, ids, feat_idx = construct_matrices(G, labels)
  max_acc = 0.0
  for t in range(0, 10):
    Y_u_next = P_ul*Y_l + P_uu*Y_u
    Y_u = Y_u_next.copy()
    labels = determine_classes(G, labels, unlabeled, feat_idx, ids, Y_u)
    test_acc = eval_accuracy(G, gt, labels, unlabeled)
    if test_acc > max_acc:
      max_acc = test_acc
  return max_acc

################################################################################
# gcn approach

# our model - 3 layers
class GCN(torch.nn.Module):
  def __init__(self, data):
    super(GCN, self).__init__()
    self.conv1 = torch_geometric.nn.GCNConv(data.num_features, 64)
    self.conv2 = torch_geometric.nn.GCNConv(64, 64)
    self.conv3 = torch_geometric.nn.GCNConv(64, 64)
    self.classifier = torch.nn.Linear(64, data.num_classes)
  
  def forward(self, x, edge_index):
    h = self.conv1(x, edge_index)
    h = h.tanh()
    h = self.conv2(h, edge_index)
    h = h.tanh()
    h = self.conv3(h, edge_index)
    h = h.tanh()
    out = self.classifier(h)
    return out, h

# training function
def train(model, criterion, optimizer, data):
  optimizer.zero_grad()
  out, h = model(data.x, data.edge_index)
  loss = criterion(out[data.train_mask], data.y[data.train_mask])
  loss.backward()
  optimizer.step()
  return loss, h

# test function
@torch.no_grad()
def test(model, data):
  model.eval()
  logits, _ = model(data.x, data.edge_index)
  pred = logits[data.test_mask].max(1)[1]
  acc = pred.eq(data.y[data.test_mask]).sum().item() / data.test_mask.sum().item()
  return acc

# convert our ground truth class strings to a numeric tensor
def gt_to_tensor(G, gt):
  gt_list = []
  gt_map = {}
  class_count = 0
  for v in G.nodes():
    if gt[v] not in gt_map:
      gt_map[gt[v]] = class_count
      class_count += 1
    gt_list.append(gt_map[gt[v]])
  #print(gt_map)
  return torch.tensor(gt_list), class_count

# convert our unlabeled set (test set) to train/test mask tensors
def unlabeled_to_tensor(G, unlabeled):
  unlabeled_map = {}
  for v in unlabeled:
    unlabeled_map[v] = True
  train_mask = []
  test_mask = []
  for v in G.nodes():
    if v in unlabeled_map:
      train_mask.append(False)
      test_mask.append(True)
    else:
      train_mask.append(True)
      test_mask.append(False)
  return torch.tensor(train_mask), torch.tensor(test_mask)


def gnn(G, gt, unlabeled):
  # Initialize pytorch data using networkx graph
  data = from_networkx(G)
  data.x = torch.eye(data.num_nodes)
  data.num_features = data.num_nodes
  data.y, data.num_classes = gt_to_tensor(G, gt)
  data.train_mask, data.test_mask = unlabeled_to_tensor(G, unlabeled)
  
  # initialize our GCN model
  model = GCN(data)
  criterion = torch.nn.CrossEntropyLoss()
  optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
  
  # Begin training
  max_acc = 0.0
  for epoch in range(50):
    loss, h = train(model, criterion, optimizer, data)
    test_acc = test(model, data)
    if test_acc > max_acc:
      max_acc = test_acc
    if epoch % 10 == 0:
     print(f"Epoch: {epoch}\tLoss: {loss:.4f}, Test Accuracy: {test_acc:.4f}")
  
  return max_acc


################################################################################
# initialize labels

def init_labels(G, gt, ratio):
  labels = {}
  unlabeled = []
  for v in G.nodes():
    if random.random() < ratio:
      labels[v] = "NA"
      unlabeled.append(v)
    else:
      labels[v] = gt[v]
  return labels, unlabeled


################################################################################
# evaluate accuracy

def eval_accuracy(G, gt, labels, unlabeled):
  true_positives = 0
  for v in unlabeled:
    if gt[v] == labels[v]:
      true_positives += 1
  return true_positives / len(unlabeled)

################################################################################
# run tests using constructed graph and benchmarks
graphs = []
ground_truths = []
graph_names = []

# construct graph and get ground_truth
print("Constructing Graph")
start = time.perf_counter()
G, gt = construct_graph()
graphs.append(G)
ground_truths.append(gt)
graph_names.append("Discogs")
print(f"Done: {time.perf_counter() - start} seconds")

print("Reading benchmarks")
start = time.perf_counter()
for dataset in ["Cora", "CiteSeer", "PubMed"]:
  G, gt = construct_graph_from_pytorch(dataset)
  graphs.append(G)
  ground_truths.append(gt)
  graph_names.append(dataset)
print(f"Done: {time.perf_counter() - start} seconds")

# proportion of data in unlabeled set
# (1 - test) gives us the size of training set
test_size = 0.1
num_test_iter = 5

for i in range(len(graphs)):
  G = graphs[i]
  gt = ground_truths[i]
  graph_name = graph_names[i]
  
  max_acc_lp = 0.0
  max_acc_nb = 0.0
  max_acc_rw = 0.0
  max_acc_gnn = 0.0
  print("Beginning test loop for", graph_name)
  for i in range(0, num_test_iter):
    # run label propagation
    #print("Running Label Prop")
    #start = time.perf_counter()
    labels, unlabeled = init_labels(G, gt, test_size)
    cur_acc_lp = label_prop(G, labels, unlabeled)
    if cur_acc_lp > max_acc_lp:
      max_acc_lp = cur_acc_lp
    #print(f"Done: {time.perf_counter() - start} seconds")
    
    # run naive bayes
    #print("Running Naive Bayes")
    #start = time.perf_counter()
    labels, unlabeled = init_labels(G, gt, test_size)
    cur_acc_nb = naive_bayes(G, labels, unlabeled)
    if cur_acc_nb > max_acc_nb:
      max_acc_nb = cur_acc_nb
    #print(f"Done: {time.perf_counter() - start} seconds")
    
    # run random walk
    #print("Running Random Walk")
    #start = time.perf_counter()
    labels, unlabeled = init_labels(G, gt, test_size)
    cur_acc_rw = random_walk(G, labels, unlabeled)
    if cur_acc_rw > max_acc_rw:
      max_acc_rw = cur_acc_rw
    #print(f"Done: {time.perf_counter() - start} seconds")
    
    # Run basic GCN using pytorch
    #print("Running Graph Convolutional Network")
    #start = time.perf_counter()
    labels, unlabeled = init_labels(G, gt, test_size)
    cur_acc_gnn = gnn(G, gt, unlabeled)
    if cur_acc_gnn > max_acc_gnn:
      max_acc_gnn = cur_acc_gnn
    #print(f"Done: {time.perf_counter() - start} seconds")
  print(f"Done with {graph_name}: {time.perf_counter() - start} seconds")
  print("===================================================")
  print("Graph Name:", graph_name)
  print("Label Prop max accuracy:", max_acc_lp)
  print("Naive Bayes max accuracy:", max_acc_nb)
  print("Random walk max accuracy:", max_acc_rw)
  print("GCN max accuracy:", max_acc_gnn)
  print("===================================================")