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

################################################################################
# Draw graph and color each community with distinct color

def draw_comm_graph(G, comms):
  colors = [comms[v] for v in G.nodes()]
  nx.draw_kamada_kawai(G, with_labels=True, node_color=colors)
  plt.show()


################################################################################
# ground truth communities

def ground_truth(filename):
  comms = {}
  with open(filename) as f:
    for line in f:
      (key, val) = line.split()
      comms[key] = int(val)
  return comms

################################################################################
# Modularity Maximization - This is a greedy modularity maximization algorithm
# that essentially works as shown in the next algorithm below. We'll combine
# communities in an agglomerative fashion, with the selected communities to 
# merge being the ones that benefit the overall modularity the most.

def mod_max(G):
  C = nx.community.greedy_modularity_communities(G)
  comms = {}
  counter = 0
  for c in C:
    for v in c:
      comms[v] = counter
    counter += 1
  return comms

################################################################################
# Newman modularity maximization - This is the same basic algorithm being 
# performed above. We are greedily optimizing modularity by checking if 
# combining communities will increase our overall modularity score. We select
# the best merge on each iteration.

def mod_max_newman(G):
  cur_comms = []
  for v in G.nodes():
    cur_comms.append({v})
  
  g = G.copy()
  updates = 1
  while len(cur_comms) > 2:
    max_mod = nx.community.modularity(G, cur_comms)
    max_comms = []
    for i in range(len(cur_comms)):
      for j in range(i+1,len(cur_comms)):
        new_comms = []
        new_comms.append(cur_comms[i].union(cur_comms[j]))
        for c in range(len(cur_comms)):
          if c != i and c != j:
            new_comms.append(cur_comms[c])
        new_mod = nx.community.modularity(G, new_comms)
        if new_mod > max_mod:
          max_mod = new_mod
          max_comms = new_comms
    if len(max_comms) < 1:
      break
    else:
      cur_comms = max_comms
  
  comms = {}
  for c in range(len(cur_comms)):
    for v in cur_comms[c]:
      comms[v] = c
  return comms


################################################################################
# run our analysis
G = nx.read_edgelist("karate.data", comments="%")

Cg = ground_truth("karate.gt.data")
Cm = mod_max(G)
Cn = mod_max_newman(G)

draw_comm_graph(G, Cg)
draw_comm_graph(G, Cm)
draw_comm_graph(G, Cn)

################################################################################
# Resolution limit - Here we can observe at how increasing the scale of our
# ring-of-cliques networks affects how a modularity maximization algorithm
# runs on the network. As our scale increases, the maximum modularity 
# communities end up being multiple cliques instead of just one. Logically, we
# would expect each clique to be a community by itself.
n = 10  # number of cliques
k = 3   # size of each clique
G = nx.ring_of_cliques(n, k)

# Check if we theoretically can resolve our cliques
resolution_limit = math.sqrt(2*G.size())
clique_degree_sum = k*(k-1)
print("Resolution Limit:", resolution_limit)
print("Degree Sum:", clique_degree_sum)
print("Can we resolve?", (resolution_limit < clique_degree_sum))

# And let's look to see if our experiments match theory
Cm = mod_max(G)
draw_comm_graph(G, Cm)