#=========================================================================
# BehavioralRTLIRVisualizationPass.py
#=========================================================================
"""Provide visualization for behavioral RTLIR AST.

Visualize Behavioral RTLIR using Graphviz packeage. The output graph is in PDF
format.  This file is automatically generated by BehavioralRTLIRImplGen.py.
"""

import os

from pymtl3.passes.BasePass import BasePass
from pymtl3.passes.rtlir.rtype.RTLIRType import BaseRTLIRType

from .BehavioralRTLIR import BehavioralRTLIRNodeVisitor


class BehavioralRTLIRVisualizationPass( BasePass ):
  def __call__( s, model ):
    visitor = BehavioralRTLIRVisualizationVisitor()

    for blk in model.get_update_blocks():
      visitor.init( blk.__name__ )
      visitor.visit( model._pass_behavioral_rtlir_gen.rtlir_upblks[ blk ] )
      visitor.dump()

class BehavioralRTLIRVisualizationVisitor( BehavioralRTLIRNodeVisitor ):
  def __init__( s ):
    s.output = 'unamed'
    s.output_dir = 'rast-viz'
    s.table_header = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0"> '
    s.table_trail = ' </TABLE>>'

  def init( s, name ):
    from graphviz import Digraph
    s.g = Digraph(
      comment = 'BehavioralRTLIR Visualization of ' + name,
      node_attr = { 'shape' : 'plaintext' }
    )
    s.blk_name = name
    s.cur = 0

  def get_str( s, obj ):
    return str(obj).replace('<', '&lt;').replace('>', '&gt;')

  def gen_table_opt( s, node ):
    ret = ''
    if isinstance( node.Type, BaseRTLIRType ):
      ret = ' <TR><TD COLSPAN="2">Type: ' + node.Type.__class__.__name__ + '</TD></TR>'
      for name, obj in vars(node.Type).items():
        obj_str = s.get_str( obj )
        if not isinstance( obj, dict ):
          ret += ' <TR><TD>' + name + '</TD><TD>' + obj_str + '</TD></TR>'
        else:
          ret += ' <TR><TD>' + name + '</TD><TD>{' + obj_str + '}</TD></TR>'
    return ret

  def dump( s ):
    if not os.path.exists( s.output_dir ):
      os.makedirs( s.output_dir )
    s.g.render( s.output_dir + os.sep + s.blk_name )

  def visit_CombUpblk( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">CombUpblk</TD></TR> <TR><TD>name</TD><TD>{name}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(name=s.get_str(node.name))
    s.g.node( str( s.cur ), label = label )
    for i, f in enumerate(node.body):
      s.g.edge( str(local_cur), str(s.cur+1), label = 'body[{idx}]'.format(idx = i) )
      s.visit( f )

  def visit_SeqUpblk( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">SeqUpblk</TD></TR> <TR><TD>name</TD><TD>{name}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(name=s.get_str(node.name))
    s.g.node( str( s.cur ), label = label )
    for i, f in enumerate(node.body):
      s.g.edge( str(local_cur), str(s.cur+1), label = 'body[{idx}]'.format(idx = i) )
      s.visit( f )

  def visit_Assign( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Assign</TD></TR> <TR><TD>blocking</TD><TD>{blocking}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(blocking=s.get_str(node.blocking))
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'target' )
    s.visit( node.target )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'value' )
    s.visit( node.value )

  def visit_If( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">If</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'cond' )
    s.visit( node.cond )
    for i, f in enumerate(node.body):
      s.g.edge( str(local_cur), str(s.cur+1), label = 'body[{idx}]'.format(idx = i) )
      s.visit( f )
    for i, f in enumerate(node.orelse):
      s.g.edge( str(local_cur), str(s.cur+1), label = 'orelse[{idx}]'.format(idx = i) )
      s.visit( f )

  def visit_For( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">For</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'var' )
    s.visit( node.var )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'start' )
    s.visit( node.start )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'end' )
    s.visit( node.end )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'step' )
    s.visit( node.step )
    for i, f in enumerate(node.body):
      s.g.edge( str(local_cur), str(s.cur+1), label = 'body[{idx}]'.format(idx = i) )
      s.visit( f )

  def visit_Number( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Number</TD></TR> <TR><TD>value</TD><TD>{value}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(value=s.get_str(node.value))
    s.g.node( str( s.cur ), label = label )

  def visit_Concat( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Concat</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    for i, f in enumerate(node.values):
      s.g.edge( str(local_cur), str(s.cur+1), label = 'values[{idx}]'.format(idx = i) )
      s.visit( f )

  def visit_ZeroExt( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">ZeroExt</TD></TR> <TR><TD>nbits</TD><TD>{nbits}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(nbits=s.get_str(node.nbits))
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'value' )
    s.visit( node.value )

  def visit_SignExt( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">SignExt</TD></TR> <TR><TD>nbits</TD><TD>{nbits}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(nbits=s.get_str(node.nbits))
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'value' )
    s.visit( node.value )

  def visit_Reduce( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Reduce</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'op' )
    s.visit( node.op )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'value' )
    s.visit( node.value )

  def visit_SizeCast( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">SizeCast</TD></TR> <TR><TD>nbits</TD><TD>{nbits}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(nbits=s.get_str(node.nbits))
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'value' )
    s.visit( node.value )

  def visit_StructInst( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">StructInst</TD></TR> <TR><TD>struct</TD><TD>{struct}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(struct=s.get_str(node.struct))
    s.g.node( str( s.cur ), label = label )
    for i, f in enumerate(node.values):
      s.g.edge( str(local_cur), str(s.cur+1), label = 'values[{idx}]'.format(idx = i) )
      s.visit( f )

  def visit_IfExp( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">IfExp</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'cond' )
    s.visit( node.cond )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'body' )
    s.visit( node.body )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'orelse' )
    s.visit( node.orelse )

  def visit_UnaryOp( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">UnaryOp</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'op' )
    s.visit( node.op )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'operand' )
    s.visit( node.operand )

  def visit_BoolOp( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">BoolOp</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'op' )
    s.visit( node.op )
    for i, f in enumerate(node.values):
      s.g.edge( str(local_cur), str(s.cur+1), label = 'values[{idx}]'.format(idx = i) )
      s.visit( f )

  def visit_BinOp( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">BinOp</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'left' )
    s.visit( node.left )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'op' )
    s.visit( node.op )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'right' )
    s.visit( node.right )

  def visit_Compare( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Compare</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'left' )
    s.visit( node.left )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'op' )
    s.visit( node.op )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'right' )
    s.visit( node.right )

  def visit_Attribute( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Attribute</TD></TR> <TR><TD>attr</TD><TD>{attr}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(attr=s.get_str(node.attr))
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'value' )
    s.visit( node.value )

  def visit_Index( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Index</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'value' )
    s.visit( node.value )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'idx' )
    s.visit( node.idx )

  def visit_Slice( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Slice</TD></TR> <TR><TD>size</TD><TD>{size}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(size=s.get_str(node.size))
    s.g.node( str( s.cur ), label = label )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'value' )
    s.visit( node.value )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'lower' )
    s.visit( node.lower )
    s.g.edge( str(local_cur), str(s.cur+1), label = 'upper' )
    s.visit( node.upper )
    if node.base is not None:
      s.g.edge( str(local_cur), str(s.cur+1), label = 'base' )
      s.visit( node.base )

  def visit_Base( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Base</TD></TR> <TR><TD>base</TD><TD>{base}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(base=s.get_str(node.base))
    s.g.node( str( s.cur ), label = label )

  def visit_LoopVar( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">LoopVar</TD></TR> <TR><TD>name</TD><TD>{name}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(name=s.get_str(node.name))
    s.g.node( str( s.cur ), label = label )

  def visit_FreeVar( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">FreeVar</TD></TR> <TR><TD>name</TD><TD>{name}</TD></TR> <TR><TD>obj</TD><TD>{obj}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(name=s.get_str(node.name), obj=s.get_str(node.obj))
    s.g.node( str( s.cur ), label = label )

  def visit_TmpVar( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">TmpVar</TD></TR> <TR><TD>name</TD><TD>{name}</TD></TR> <TR><TD>upblk_name</TD><TD>{upblk_name}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(name=s.get_str(node.name), upblk_name=s.get_str(node.upblk_name))
    s.g.node( str( s.cur ), label = label )

  def visit_LoopVarDecl( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">LoopVarDecl</TD></TR> <TR><TD>name</TD><TD>{name}</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail).format(name=s.get_str(node.name))
    s.g.node( str( s.cur ), label = label )

  def visit_Invert( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Invert</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Not( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Not</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_UAdd( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">UAdd</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_USub( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">USub</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_And( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">And</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Or( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Or</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Add( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Add</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Sub( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Sub</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Mult( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Mult</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Div( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Div</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Mod( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Mod</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Pow( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Pow</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_ShiftLeft( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">ShiftLeft</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_ShiftRightLogic( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">ShiftRightLogic</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_BitAnd( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">BitAnd</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_BitOr( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">BitOr</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_BitXor( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">BitXor</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Eq( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Eq</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_NotEq( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">NotEq</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Lt( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Lt</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_LtE( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">LtE</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_Gt( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">Gt</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )

  def visit_GtE( s, node ):
    s.cur += 1
    local_cur = s.cur
    table_body = '<TR><TD COLSPAN="2">GtE</TD></TR>'
    table_opt = s.gen_table_opt( node )
    label = (s.table_header + table_body + table_opt + s.table_trail)
    s.g.node( str( s.cur ), label = label )
