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 communities

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

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

draw_comm_graph(G, comms)

################################################################################
# Modularity Maximization

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

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

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


G = nx.ring_of_cliques(25, 5)