
"""
    graph_utils
    ~~~~~~~~~~~

    Function for network manipulation.

    :author: Fahad Khalid
    :license: MIT, see LICENSE for more details.
"""

import math
import json
import igraph


# Provides methods for constructing and manipulating genotype
# networks. The methods in this class are agnostic to molecule
# type.
class NetworkBuilder :
	# Constructor
	def __init__(self, seqBitManip) :
		# Store reference to the sequence/bit manipulator object
		# in use
		self.bitManip = seqBitManip

	# Return 'True' if the source sequence is connected to the target
	# sequence, and 'False' otherwise.
	def areConnected(self, source, target) :
		return self.bitManip.areNeighbors( \
					self.bitManip.seqToBits(source), \
					self.bitManip.seqToBits(target))

	# Create a genotype network using the give list of sequences
	def createGenoNet(self, repertoire, sequences, scores) :
		# Initialize the empty network
		network = igraph.Graph()

		# Set repertoire name as a graph attribute
		network["name"] = repertoire

		# Create vertices
		network.add_vertices(len(sequences))

		# Add scores to all vertices
		# FIXME: Change back from 'escores' to 'scores'
		network.vs["escores"]	= scores
		
		# Add sequences to all vertices
		network.vs["sequences"] = sequences
		
		# Add bit formatted sequences to all vertices
		# Note: For very long sequences, the size of the corresponding
		#		integer value becomes too long for igraph to be able
		#		to deal with it (since igraph is C based). Therefore,
		#		bitseqs cannot be stored in the graph.
		bitseqs = [self.bitManip.seqToBits(seq) for seq in sequences]
	
		# Get a list of pairs of indices which represent paris of
		# sequences that are 1-neighbors, i.e., only one mutation
		# apart.
		edges = [(i,j) \
					for i in range(len(bitseqs) - 1) \
					for j in range(i + 1, len(bitseqs)) \
					if self.bitManip.areNeighbors(bitseqs[i], bitseqs[j])]

		# Connect the two with an edge
		network.add_edges(edges)
		
		return network

	# For the given sequence, returns a list of 1-neighbors that are not part of
	# the given genotype network.
	# Note: The list of external neighbors returned comprises bitseqs, not
	#		strings. The caller must be aware of this.
	def getExternalNeighbors(self, sequence, network) :
		# Get a list of all possible 1-neighbors of the given sequence
		allNeighbors = [ neighbor for neighbor in \
			self.bitManip.generateNeighbors(self.bitManip.seqToBits(sequence)) ]

		# Get a list of 1-neighbors that are part of the given genotype
		# network
		genoNetNeighbors = [self.bitManip.seqToBits(network.vs[neighbor]["sequences"]) \
								for neighbor in self.getNeighbors(sequence, network)]

		# Get a list of sequences that are 1-neighbors, but not part of the
		# genotype network. This is achieved by taking the difference of
		# two sets, i.e., elements in allNeighbors not in genoNetNeighbors
		externNeighbors = list(set(allNeighbors) - set(genoNetNeighbors))

		return externNeighbors

	# Get all unique external neighbors for the given genotype network
	def getAllExtNeighbors(self, network) :
		# List to be populated with external neighbors
		extNeighbors = []

		# Get all sequences
		sequences = network.vs["sequences"]

		# For each sequence, get the list of external neighbors and 
		# append it to extNeighbors
		for sequence in sequences :
			extNeighbors.extend(self.getExternalNeighbors(sequence, network))

		# Remove redundant sequences
		extNeighbors = list(set(extNeighbors))

		return extNeighbors

	# Get the vertex based on the sequence string from the
	# given network
	def getVertex(self, sequence, network) :
		try :
			return network.vs.find(sequences=sequence)
		except ValueError :
			print("Error! ... Non-existent vertex requested: " + str(sequence))

	# Returns a list of sequences that constitute the neighborhood
	# of the given sequence within the given network.
	def getNeighbors(self, sequence, network) :
		return network.neighbors(self.getVertex(sequence, network))

	# Plot the given network using the given layout.
	def plotNetwork(self, network, layout, outPath) :
		vertices = [network.vs[i].index for i in range(0, network.vcount())]
		# For GNs,
		# degrees = network.degree(vertices)
		# For ENs,
		degrees = network.outdegree(vertices)

		visual_style = {}
		# For GNs:
		#visual_style["vertex_size"] = [x*2 for x in degrees]
		# For ENs:
		#visual_style["vertex_size"] = [x*2 for x in degrees]
		visual_style["vertex_size"] = [(x*10 * x*10) for x in network.vs["Evolvability"]]
		#visual_style["edge_width"] = [math.log(x) for x in network.es["weight"]]
		#visual_style["vertex_color"] = [int(x*300) for x in network.vs["escores"]]

		layout = network.layout(layout)

		igraph.plot(network, outPath + network["name"] + "_evonet.svg", bbox=(1500, 1500), **visual_style)

	# Return a list, the size of which represents the number of commnected
	# component in the network. Each element in the list corresponds to the
	# size of one of the connected components.
	def getComponents(self, network) :
		vertexCluster = network.components()

		return vertexCluster.sizes()

	# Returns the gian component.
	def getGiantComponent(self, network) :
		return network.components().giant()

	# Create a network that consists of the given sequence and its
	# one-neighbors.
	def createNeighborhoodNet(self, sequence, network) :
		sequences = [sequence]
		escores = []

		# Get the vertex corresponding to the sequence
		vertex = self.getVertex(sequence, network)

		# Get all adjacent vertices
		neighbors = network.neighbors(vertex)

		# Get escore for the given sequence
		escores.append(network.vs[vertex.index]["escores"])
		
		# Get sequences for all neighbors
		for neighbor in neighbors :
			sequences.append(network.vs[neighbor]["sequences"])
			escores.append(network.vs[neighbor]["escores"])

		# Initialize the empty network
		neighborhoodNet = igraph.Graph()

		# Create the neighborhood network
		
		# Create vertices
		neighborhoodNet.add_vertices(len(sequences))

		# Add attributes
		neighborhoodNet.vs["label"] = [sequences[i]+"\n"+str(escores[i]) \
										for i in range(len(sequences))]

		for i in range(1, len(sequences)) :
			neighborhoodNet.add_edge(0, i)

		# Print the list of sequences for validation purposes
		print(sequences)

		return neighborhoodNet

	# Create a network of the given repertoires where repertoire 'A' is
	# connected to repertoire 'B' if 'A' has one or more external
	# neighbors in 'B'.
	# Arguments: title=network name, genoNets=giant for the repertoire
	# TODO: We have the option of either using the GNs as nodes, or just
	#		the GN names. I need to think more about the pros and cons
	#		of each approach ...
	def createEvoNet(self, title, genoNets) :
		# Get a list of repertoire names
		repertoires = [ gn["name"].rpartition("_dominant")[0] \
						if gn["name"].rpartition("_dominant")[0]  \
						else gn["name"] \
						for gn in genoNets ]
		#repertoires = [gn["name"].replace("_dominant", "") for gn in genoNets]
		# repertoires = [gn["name"].rstrip("_dominant") for gn in genoNets]

		# Get a list of evolvability scores corresponding to the
		# repertoires
		evolvability = [gn["Evolvability"] for gn in genoNets]

		# Initialize the empty network
		network = igraph.Graph(directed=True)

		# Set title as the graph name
		network["name"] = title

		# Create vertices: Each repertoire name is a vertex in the network
		network.add_vertices(len(repertoires))

		# Add attributes

		# TODO: Remove this from the tool ...
		labels = [repertoires[i] + "\n" + str(float("{0:.2f}".format(evolvability[i]))) \
		 			for i in range(len(repertoires))]
		# labels = []
		# labels[:] = repertoires

		# Add labels to all vertices
		network.vs["label"]		= labels	# TODO: Remove this from the tool ...
		# Add repertoire names to all vertices
		network.vs["Repertoires"] = repertoires
		network.vs["GenotypeSet"] = repertoires # Redundant; added for GML
		network.vs["Evolvability"] = evolvability
	
		# Get a list of pairs of indices which represent paris of
		# repertoires that are neighbors, i.e., share at least one external
		# neighbor. The edge determination algorithm takes into account the 
		# antisymmetric relation.
		edges = [(i,j) \
					for i in range(len(repertoires)) \
					for j in range(len(repertoires)) \
					if i != j and self.repertoiresAreConnected(genoNets[i], repertoires[j])]

		# Add the above computed edges to the network
		network.add_edges(edges)

		# Add edge weights
		
		# Create a dict: keys=repertoires, values=genonets
		repToNetDict = {repertoires[i] : genoNets[i] for i in range(len(repertoires))}
		
		# Set weight for each edge, where weight=No. of evo target genotypes from i to j
		weights = [self.getEdgeWeight(edges[i], network, repToNetDict) \
									for i in range(len(edges))]

		# Add edge weights and edge labels as attributes
		network.es["weight"] = weights
		network.es["label"] = weights
		
		return network

	# Returns True if 'gn1' has at least one external neighbor in the GN for
	# 'rep2'
	def repertoiresAreConnected(self, gn1, rep2) :
		# Result flag
		areConnected = False

		# Get the list of evolvability target repertoires for gn1
		evoTargets = json.loads(gn1["Evolvability_targets"])

		# If there is at least one target in the list,
		if evoTargets :
			# If 'rep2' is in the list of targets,
			if rep2 in evoTargets :
				# 'gn1' has external neighbors in rep2
				areConnected = True

		return areConnected

	# Return the weight for the given edge of the given EN
	def getEdgeWeight(self, edge, en, repToNetDict) :
		# Get the repertoire names for source and target vertices
		srcRep = en.vs[edge[0]]["Repertoires"]
		trgRep = en.vs[edge[1]]["Repertoires"]

		# Get the source GN
		srcGN = repToNetDict[srcRep]

		# Get the list of target dicts in source GN
		allTargets = srcGN.vs["Evolvability_targets"]

		# External neighbors of source in target; to be
		# populated in the loop below.
		extNeighbors = []

		# Extract the required dict
		for trgDict in allTargets :
			# If the target repertoire is a key in this dict,
			if trgRep in trgDict :
				# Add external neighbors to the global list
				extNeighbors.extend(trgDict[trgRep])

		# Remove redundant genotypes
		extNeighbors = list(set(extNeighbors))

		return len(extNeighbors)