from prefuse import *
from prefuse.render import *
from prefuse.action import *
from prefuse.action.assignment import *
from prefuse.action.layout import *
from prefuse.action.layout.graph import *
from prefuse.controls import *
from prefuse.visual import *
from prefuse.util import *
from prefuse.util.display import *
from prefuse.data import *
from javax.swing import *
from java.lang import String, Integer

class InteractiveGraphVisualization:

	def __init__(self, size=(800, 600)):
		self.size = size

		self._init_graph()

		self.vis = Visualization()
		self.vis.add("graph", self.graph)

		self._init_vis()

		self.display = Display(self.vis)
		self.display.setSize(size[0], size[1])
		self.display.addControlListener(DragControl())
		self.display.addControlListener(PanControl())
		self.display.addControlListener(ZoomControl())
		self.display.addControlListener(ZoomToFitControl())
		self.display.addControlListener(FocusControl(1))

		self.frame = JFrame()
		self.frame.getContentPane().add(self.display)
		self.frame.pack()
		self.frame.setVisible(True)

		self.vis.run("layout")

	def _init_graph(self):
		nodes = Table()
		nodes.addColumn("val", String().getClass())

		self.graph = Graph(nodes, False)

	def _init_vis(self):
		r = LabelRenderer("val")
		r.setRoundedCorner(8, 8)
		self.vis.setRendererFactory(DefaultRendererFactory(r))
		edges = ColorAction("graph.edges", VisualItem.STROKECOLOR, ColorLib.gray(200))
		text = ColorAction("graph.nodes", VisualItem.TEXTCOLOR, ColorLib.gray(0))
		fill = ColorAction("graph.nodes", VisualItem.FILLCOLOR, ColorLib.rgb(200,200,255))

		color = ActionList()
		color.add(fill)
		color.add(text)
		color.add(edges)
		self.vis.putAction("color", color)
	
		layout = ActionList(-1)
		layout.add(ForceDirectedLayout("graph", True))
		layout.add(RepaintAction())
		self.vis.putAction("layout", layout)

		repaint = ActionList()
		repaint.add(color)
		self.vis.putAction("repaint", repaint)

	def _reposition(self):
		margin = 50
		duration = 2000

		bounds = self.vis.getBounds("graph")
		GraphicsLib.expand(bounds, margin + 1 / self.display.getScale())
		DisplayLib.fitViewToBounds(self.display, bounds, duration)

	def _repaint(self):
		self.vis.run("repaint")
		self._reposition()

	def add_node(self, text):
		node = self.graph.addNode()
		node.set("val", text)

		self._repaint()

		return node
		
	def add_edge(self, n1, n2):
		edge = self.graph.addEdge(n1, n2)

		self._repaint()

		return edge