
from pathlib import Path
import graphviz
from matplotlib import colormaps
from matplotlib.colors import to_hex
import matplotlib.pyplot as plt
from IPython.display import Image, display
import matplotlib.image as mpimg
from io import BytesIO

def create_digraph(engine="dot"):

    engine_configs = {
        "dot": {
            'rankdir': 'TB',
            'ranksep': '1.0',
            'nodesep': '0.4',
            'splines': 'spline',
			'constraint' : 'true',
            'concentrate': 'false'
        },
		"neato": {
			'overlap': 'scale',
			'overlap_scaling': '-4',  
			'esep': '+2.5',           
			'sep': '+1.75',           
			'model': 'shortpath',
			'damping': '0.85',
			'epsilon': '0.00001',
			'maxiter': '1000000',
			'start': '5',
		},
        "fdp": {
            'overlap': 'prism',
            'sep': '+1.5',
            'K': '1.0',
            'splines': 'true',
			'len' : '3.0',
            'maxiter': '5000'
        }
    }

    graph_attr = engine_configs.get(engine, {})

    node_attr = {
        'shape': 'box',
        'style': 'rounded, filled',
        'fillcolor': 'lightblue',
        'fontname': 'Helvetica'
    }
    
    edge_attr = {
        'penwidth': '1.2',
        'color': 'gray40'
    }

    return graphviz.Digraph(
        engine=engine,
        graph_attr=graph_attr,
        node_attr=node_attr,
        edge_attr=edge_attr
    )


def draw_graph( 
	aM, 
	output_dir : Path | None = None,
	title="am_graph", 
	view=True, 
	subgraphs=None,
	engine="dot",
	symbols_only=False ):

	GV = create_digraph(engine=engine)

	cmap = colormaps['Set3']

	if subgraphs :
		graph_colors = [to_hex(cmap(i % 12)) for i in range(len(subgraphs))]

	for node in range( len(aM.states) ) :
		
		if subgraphs :
			
			node_subgraph = -1
			for i, sg in enumerate( subgraphs ) :
				if node in sg :
					node_subgraph = i

			node_color = graph_colors[ node_subgraph ] if node_subgraph >= 0 else 'red'

		else :
			node_color = to_hex(cmap(8))

		node_tex = aM.states[ node ].name
		if not node_tex :
			node_tex = str(node)

		GV.node(
			str(node), 
			label=node_tex, 
			shape='circle',
			style='bold,filled',
        	fillcolor=node_color,
        	color='black',
			width='0.5' )


	for tr in aM.transitions :

		u = tr.origin_state_idx
		v = tr.target_state_idx

		pr_str = str( tr.pq ) if aM.is_q_weighted else str( round( tr.prob, 4 ) )
		
		fontsize= '30' if symbols_only else '12'

		if symbols_only :
			label_text = f"{aM.alphabet[tr.symbol_idx]}"
		else :
			label_text = f"{aM.alphabet[tr.symbol_idx]}({pr_str})"

		html_label = (
            f'<<TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0" CELLPADDING="2">'
            f'<TR><TD BGCOLOR="#FFFFFF">{label_text}</TD></TR>'
            f'</TABLE>>'
        )

		GV.edge(
			str(u), str(v), 
			label=html_label,
			fontsize=fontsize,
			fontname='Times-Italic',
			labelfloat="true",
			fontcolor="#0042ad"
		)

	if view == True :

		GV.graph_attr.update(dpi='300')
		png_bytes = GV.pipe(format="png")
		buf = BytesIO( png_bytes )
		img = mpimg.imread( buf )

		plt.figure(figsize=(8, 6))
		plt.imshow(img)
		plt.axis('off')
		plt.tight_layout()
		plt.show()

	if output_dir is not None :
		
		GV.attr(bgcolor='transparent')
		GV.render(title, directory=output_dir, view=False, format='png', cleanup=True)