#!/usr/bin/env python3

import argparse

from graphviz import Digraph

import grimoire.hmm as hmm

## Command line ##

extended_help = """
%(prog)s draws a state diagram for an HMM using Graphviz. Transitions
probabilities can be labeled, but currently there is no support for
emission probabilities because they can be quite complex when
lexicalized. The grimmoire/docs directory contains svg files created
with %(prog)s.
"""

parser = argparse.ArgumentParser(
	description='Draws HMM state diagrams.',
	formatter_class=argparse.RawDescriptionHelpFormatter,
	epilog=extended_help)
parser.add_argument('--hmm', required=True, type=str,
	metavar='<path>', help='path to input HMM file')
parser.add_argument('--svg', required=True, type=str,
	metavar='<path>', help='path to output svg file')
parser.add_argument('--prob', action='store_true',
	help='label transition probabilities')
parser.add_argument('--tb', action='store_true',
	help='top to bottom graph [default LR]')
arg = parser.parse_args()

if __name__ == '__main__':

	# Graph defaults
	graph = Digraph(format='png')
	graph.attr('node', shape='circle')
	graph.attr(label=str(arg.hmm))
	if arg.tb: graph.attr(rankdir='TB')
	else: graph.attr(rankdir='LR')

	# Build graph from HMM
	hmm = hmm.HMM.read(arg.hmm)
	for s1 in hmm.states:
		graph.node(s1.name)
		for s2 in s1.next:
			label = None
			if arg.prob:
				label=str(round(s1.next[s2], 3))
			graph.edge(s1.name, s2, label=label)

	# Save
	with open(arg.svg, 'w') as f:
		f.write(graph._repr_svg_())
