#!/usr/bin/python

from numpy import *
import os
import sys
import cPickle
import girvan_newman as gn
import networkx as nx
import pylab

cutoff=.5
prefix='gn_apo_cutoff_%s'%cutoff
mi = genfromtxt('/u1/jwereszc/get3_anton/analysis/mutinf/blocks_6x450/apo/n6o4/residues-nsims6-structs4500-bin30_bootstrap_sigavg01_mutinf_res_0diag.txt',skip_header=1,usecols=range(1,709))
mi=mi+1e-8
edges = genfromtxt('/u1/jwereszc/get3_anton/analysis/gn/apo/contacts/contacts_5')
file='edgelist_%s.dat'%cutoff
f=open(file,'w')
for i in range(708):
    for j in range(i+1,708):
        if edges[i,j]>0 :
            f.write("%6i  %6i  %7.4f\n" % (i,j,-.5*log(1-exp(-2*mi[i,j]/3))))
f.close()
save_figs = True
# Load graph from edge list file
G = nx.read_weighted_edgelist(file, nodetype=int)

# Perform Girvan Newman algorithm
itr = int(len(G.edges()) - 9)    # Number of iterations
max_com_graph, l_com_graph, modularity = gn.Girvan_Newman_algorithm(G, itr)

# Show the modularity
pylab.figure()
pylab.plot(modularity)
print 'Maximum modularity = %7.4f' %max(modularity)
pylab.xlabel('Iteration')
pylab.ylabel('Modularity Q-value')
pylab.title('Modularity')
if save_figs:
	pylab.savefig('%s_modularity.pdf' % prefix, format='pdf')

# Show the Original Graph
pylab.figure()
ax = pylab.gca()
positions = nx.layout.spring_layout(G)
nx.draw(G, ax=ax, pos=positions)
pylab.title('Network')
if save_figs:
	pylab.savefig('%s_network.pdf'%prefix, format='pdf')
 
# Show the community structure graph (all nodes, colored by community)
pylab.figure()
communities = gn.find_communities(max_com_graph)
gn.draw_communities(G, communities, pos=positions)
pylab.title('Network with residues colored by Community')
if save_figs:
	pylab.savefig('%s_community_structure.pdf'%prefix, format='pdf')

# Show coarsegrained community graph 
# Note: size of nodes proportional to the number of residues
#       width of edges proportional to the number of shortest paths 
#       between communities
pylab.figure()
cg_graph = gn.coarsegrain_communities(G, communities)

# Make positions of nodes match up to view in VMD
coords = genfromtxt('/u1/jwereszc/get3_anton/analysis/gn/apo/gn/alpha_carbons.pdb')
vmdcolors = [(0,0,1,1),(1,0,0,1),(.35,0.35,.35,1),(1,.5,0,1),(1,1,0,1),(0.5,0.5,0.2,1),(.45,.0,.9,1),(0,1,0,1),(1,1,1,1),(1,.6,.6,1),(.25,.75,.75,1),(0.65,0,.65,1),(0.5,0.9,0.4,1),(0.9,0.4,0.7,1),(0.5,0.3,0,1),(0.5,0.5,0.75,1)]
pos={}
colors=[]
for i in range(len(communities)):
    pos.update({i:(mean(coords[communities[i],6]),mean(coords[communities[i],7]))})
    colors.append(vmdcolors[i])

gn.draw_coarsegrain(cg_graph,pos=pos,colors=colors,node_factor=10.0, edge_factor=50.)
pylab.title('Community Network Coarsegrain')
if save_figs:
    pylab.savefig('%s_community_structure_cg.pdf'%prefix, format='pdf')

pylab.ion()
pylab.show()

raw_input("Press ENTER to exit")

# Save the communities classification to a pickle file
gn.pickle_data(G, communities)

vmdtemp=open('temp.tcl','w')
vmdtemp.write("mol new /u1/jwereszc/get3_anton/analysis/gn/apo/gn/oriented.pdb waitfor all\n")

for i in range(len(communities)):
    val=i
    if (val==6):
        val=26
	for j in range(len(communities[i])):
		vmdtemp.write("[atomselect top \"residue %i\"] set beta %i\n"%(communities[i][j],val))
vmdtemp.write("[atomselect top all] writepdb %s.pdb\nexit\n" % prefix)
vmdtemp.close()
os.system("vmd -dispdev text -e temp.tcl")
# Enter debugger
# import pdb
# pdb.set_trace()