# probs.py
# Andrew Davison, ad@coe.psu.ac.th, May 2025

'''
https://graphviz.org/doc/info/shapes.html
'''

import graphviz

TEMP_FNM = 'temp_graph'

sizes = [('S',0.1), ('M',0.25), ('L',0.3), ('XL', 0.2), 
         ('XXL',0.15)]
colors = [('W',0.5), ('R',0.3), ('G',0.2)]

# dot, TD
tree = graphviz.Digraph(format='png')
tree.attr(label='Probability Tree')

tree.node('Root')
totCount = 0
for size, sprob in sizes:
  tree.node(size, shape='circle')
  tree.edge('Root', size, label=str(sprob))
  for col, cprob in colors:
    pNode = "p"+str(totCount)
    tNode = "t"+str(totCount)
    totCount += 1

    tree.node(pNode, label=col, shape='circle', color='blue')
    tree.edge(size, pNode, label=str(cprob))

    # total prob as blank shape linked by invisible edge
    tree.node(tNode, label=f"{(sprob*cprob):.3f}", 
                     labelloc='t', shape='none')
    tree.edge(pNode, tNode, len='0.5', weight='3', 
                                       style='invis')
                        # shorter and more vertical

tree.render(filename=TEMP_FNM, view=True, cleanup=True)