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()


################################################################################
# Let's look at the Zachary Karate graph and its ground-truth communities

G = nx.read_edgelist("karate.data", comments="%")
comms = {}
with open("karate.gt.data") as f:
  for line in f:
    (key, val) = line.split()
    comms[key] = int(val)

draw_comm_graph(G, 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.

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

draw_comm_graph(G, 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.

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
  
  cur_comms = max_comms
  comms = {}
  for c in range(len(cur_comms)):
    for v in cur_comms[c]:
      comms[v] = c
        
  draw_comm_graph(G, comms)

################################################################################
# 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.

G = nx.ring_of_cliques(8, 4)
G = nx.ring_of_cliques(25, 5)
nx.draw(G)
plt.show()