from IPython.display import Image, display
import os
import subprocess

try:
	import pydot

	has_pydot = True
except ImportError:
	print("pydot not installed")
	has_pydot = False

from cbd.CBDBlocks import *

def draw_to_file(name, cbd, show=False):
	if not has_pydot:
		print("No pydot")
		return

	if cbd is None:
		print("graph_to_dot Error: Empty CBD")
		return

	vattr = ''
	nodes = {}
	graph = pydot.Dot(name, graph_type='digraph')


	for i, block in enumerate(cbd.getBlocks()):

		block_type = block.getBlockType()

		vattr = block.getBlockName()#block_type + str(i)

		fillcolor = get_fill_color(block_type)

		if "Constant" in block_type:
			try:
				vattr += "\\n Value= " + str(str(block.getValue()))
			except KeyError:
				pass
		elif "Gain" in block_type:
			try:
				vattr += "\\n Value= " + str(str(block.gain))
			except KeyError:
				pass
		elif "Test" in block_type:
			try:
				vattr += "\\n Test= " + str(block.expr) + " "+ str(block.switch_value)
			except KeyError:
				pass


		nodes[block.getBlockName()] = pydot.Node(vattr, style="filled", fillcolor=fillcolor)
		graph.add_node(nodes[block.getBlockName()])



	link_colours = {"paired_with": "#505050", "link_in": "#6E2229", "init_val": "#938344",
					"IC": "darkgoldenrod", "reset": "#B9512B"}


	for block in cbd.getBlocks():

		for other in block.linksIN:
			colour = link_colours["link_in"]
			graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour))

		try:
			other = block.IC
			colour = link_colours["IC"]
			graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="IC"))
		except AttributeError:
			pass

		try:
			other = block.reset
			colour = link_colours["reset"]
			graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="reset"))
		except AttributeError:
			pass

		try:
			other = block.init_val
			colour = link_colours["init_val"]
			graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="init_val"))
		except AttributeError:
			pass

		try:
			other = block.TRUE_in
			colour = "#000000"
			graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="TRUE"))
		except AttributeError:
			pass

		try:
			other = block.FALSE_in
			colour = "#000000"
			graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="FALSE"))
		except AttributeError:
			pass


	#
	# 	for other in in_blocks:
	# 		name = other.getBlockName()
	# 		label = ""
	#
	# 		if not name.startswith("IN"):
	# 			label=name
	#
	# 		# if not name.startswith("OUT"):
	# 		# 	label = label + " / " + other.getBlockName()
	#
	# 		write("{a} -> {b} [label=\"{lbl}\"];\n".format(a=other.getBlockName(),
	# 													   b=block.getBlockName(),
	#


	# for e in g.es:
	#
	# 	src, trgt = e.tuple
	# 	src = int(src)
	# 	trgt = int(trgt)
	#
	# 	if src in internal_links.keys():
	# 		internal_links[src]["target"] = trgt
	# 		continue
	# 	elif trgt in internal_links.keys():
	# 		internal_links[trgt]["source"] = src
	# 		continue
	#
	# 	else:
	#
	# 		color = "black"
	# 		if "MatchModel" in mms[src]:
	# 			color = link_colours["match_contains"]
	# 		elif "ApplyModel" in mms[src]:
	# 			color = link_colours["apply_contains"]
	#
	#

	png_filename = name + '.png'

	dot_filename = png_filename.replace(".png", ".dot")
	graph.write(dot_filename, prog='dot')

	command = "dot -Tpng " + dot_filename + " -o " + png_filename
	subprocess.call(command, shell=True)

	command = "rm " + dot_filename
	subprocess.call(command, shell=True)

	if show:
		im = Image(png_filename)
		display(im)


def get_fill_color(node_type):
	nt = node_type

	fillcolor = "#9999FF"

	if nt == 'MatchModel':
		fillcolor = "#E15C34"

	elif "Gain" in nt:
		fillcolor = "#22DDFF"

	elif "Integrator" in nt:
		fillcolor = "lightgoldenrod"

	elif "Test" in nt:
		fillcolor = "coral"

	elif nt == 'ApplyModel':
		fillcolor = "#FED017"

	elif nt == 'Equation':
		fillcolor = "#66FF33"

	elif nt in ['leftExpr', "rightExpr"]:
		fillcolor = "#22DDFF"

	elif nt in ['hasAttr_S', 'MT_pre__hasAttr_S', 'MT_post__hasAttr_S', "hasAttribute_S"]:
		fillcolor = "#E34A0E"

	elif node_type == 'hasAttr_T' or "hasAttribute" in nt:
		fillcolor = "#C75C0D"

	elif nt == 'Attribute':
		fillcolor = "#FFCC00"

	elif "hasArgs" in nt:
		fillcolor = "#FFFF99"

	elif "Concat" in nt:
		fillcolor = "#9999FF"

	elif "Constant" in nt:
		fillcolor = "lightblue"

	elif nt in ['directLink_S', 'directLink_T']:
		fillcolor = "lightyellow"

	elif nt in ['indirectLink_S', 'indirectLink_T']:
		fillcolor = "lightgreen"

	elif node_type == 'backward_link':
		fillcolor = "coral"

	elif node_type == 'trace_link':
		fillcolor = "lightgoldenrod"

	return fillcolor