#! /usr/bin/env mflowgen-python
#=========================================================================
# mflowgen-info
#=========================================================================
# Reads the configure.yml of an mflowgen step and pretty prints content.
#
#  -h --help     Display this message
#  -v --verbose  Verbose mode
#  -y --config   Path to the step configure.yml to parse
#     --simple   Simpler visualization without connectivity info
#
# Author : Christopher Torng
# Date   : March 5, 2020
#

#-------------------------------------------------------------------------
# Output formatting
#-------------------------------------------------------------------------
# The synthesis node might look like this with connectivity shown,
# parameters, flags, and source directory:
#
#            + 1-freepdk-45nm
#            + adk
#            |
#            |       + 3-rtl
#            |       + design.v
#            |       |
#            |       |              + 2-constraints
#            |       |              + constraints.tcl
#            |       |              |
#            V       V              V
#         +---------------------------------------------+
#         | adk | design.v | constraints.tcl | run.saif |
#         +---------------------------------------------+
#         |                                             |
#         |           4-synopsys-dc-synthesis           |
#         |                                             |
#         +---------------------------------------------+
#         | design.v | design.sdc | design.namemap      |
#         +---------------------------------------------+
#              |           |
#              |           V
#              |           + 5-cadence-innovus-flowsetup
#              |           + design.sdc
#              |           +
#              |           + 6-cadence-innovus-place-route
#              |           + design.sdc
#              V
#              + 5-cadence-innovus-flowsetup
#              + design.v
#              +
#              + 6-cadence-innovus-place-route
#              + design.v
#
#     Parameters
#
#     - clock_period   : 2.0
#     - design_name    : GcdUnit
#     - flatten_effort : 0
#     - nthreads       : 16
#     - topographical  : true
#
#     Flags
#
#     - sandbox : False
#
#     Source Directory
#
#     - /path/to/synopsys-dc-synthesis
#
# With the --simple flag, the node might look like this:
#
#            |       |              |
#            V       V              V
#         +---------------------------------------------+
#         | adk | design.v | constraints.tcl | run.saif |
#         +---------------------------------------------+
#         |                                             |
#         |           4-synopsys-dc-synthesis           |
#         |                                             |
#         +---------------------------------------------+
#         | design.v | design.sdc | design.namemap      |
#         +---------------------------------------------+
#              |           |
#              V           V
#

import argparse
import sys
import yaml

#-------------------------------------------------------------------------
# Command line processing
#-------------------------------------------------------------------------

class ArgumentParserWithCustomError(argparse.ArgumentParser):
  def error( self, msg = "" ):
    if ( msg ): print("\n ERROR: %s" % msg)
    print("")
    file = open( sys.argv[0] )
    for ( lineno, line ) in enumerate( file ):
      if ( line[0] != '#' ): sys.exit(msg != "")
      if ( (lineno == 2) or (lineno >= 4) ): print( line[1:].rstrip("\n") )

def parse_cmdline():
  p = ArgumentParserWithCustomError( add_help = False )
  p.add_argument( "-v", "--verbose", action="store_true" )
  p.add_argument( "-h", "--help",    action="store_true" )
  p.add_argument( "-y", "--config",  required=True       )
  p.add_argument(       "--simple",  action="store_true" )
  opts = p.parse_args()
  if opts.help: p.error()
  return opts

#-------------------------------------------------------------------------
# Main
#-------------------------------------------------------------------------

def main():

  opts = parse_cmdline()

  # Read the configure.yml

  with open( opts.config ) as fd:
    try:
      data = yaml.load( fd, Loader=yaml.FullLoader )
    except AttributeError:
      # PyYAML for python2 does not have FullLoader
      data = yaml.load( fd )

  #-----------------------------------------------------------------------
  # Graph Visual
  #-----------------------------------------------------------------------

  name = data['build_dir']

  # Grab the inputs and outputs

  try:
    inputs  = data['inputs']
  except KeyError:
    inputs = []

  try:
    outputs = data['outputs']
  except KeyError:
    outputs = []

  # Grab the edges, which should be in this format
  #
  #  edges_i:
  #    adk:
  #    - f: adk
  #      step: 1-freepdk-45nm
  #    constraints.tcl:
  #    - f: constraints.tcl
  #      step: 2-constraints
  #    design.v:
  #    - f: design.v
  #      step: 3-rtl
  #  edges_o:
  #    design.sdc:
  #    - f: design.sdc
  #      step: 5-cadence-innovus-flowsetup
  #    - f: design.sdc
  #      step: 6-cadence-innovus-place-route
  #    design.v:
  #    - f: design.v
  #      step: 5-cadence-innovus-flowsetup
  #    - f: design.v
  #      step: 6-cadence-innovus-place-route
  #

  try:
    edges_i  = data['edges_i']
  except KeyError:
    edges_i = {}

  try:
    edges_o  = data['edges_o']
  except KeyError:
    edges_o = {}

  n_inputs  = len( inputs  )
  n_outputs = len( outputs )

  # Start drawing ASCII
  #
  # nchars is the length
  #
  #     inputs_str  --> '| adk | design.v | constraints.tcl | run.saif |'
  #     outputs_str --> '|  |'
  #
  #                      ^-----------------  n_chars  -----------------^

  inputs_str  = '| ' + ' | '.join( inputs  ) + ' |'
  outputs_str = '| ' + ' | '.join( outputs ) + ' |'
  name_str    = '| ' + name                  + ' |'

  n_chars = max( len(inputs_str), len(outputs_str), len(name_str) )

  # Recenter the shorter string so that this:
  #
  #     '| foo | bar |'
  #
  # becomes this based on the given length:
  #
  #     '|     foo     |     bar     |'
  #
  # The input string x _must_ have a non-whitespace char on both sides.

  def recenter( x, length ):
    tokens       = [ t.strip() for t in x.split() ]
    extra_space  = length - len( ''.join(tokens) )
    n_items      = len( tokens )
    space_per    = int( extra_space / ( n_items - 1 ) )
    space_left   = int( extra_space % ( n_items - 1 ) )
    spacer       = space_per * ' '
    spaced_str   = spacer.join( tokens )
    spaced_str_f = spaced_str[:-1] + ' '*space_left + spaced_str[-1:]
    return spaced_str_f

  inputs_str  = recenter( inputs_str,  n_chars )
  outputs_str = recenter( outputs_str, n_chars )
  name_str    = recenter( name_str,    n_chars )

  # arrowify
  #
  # Turn this:
  #
  #     '|     foo     |     bar     |'
  #
  # into this:
  #
  #     '       V             V       '
  #

  def arrowify( x ):
    # No arrows for an empty string "|      |"
    if x[1:-1].strip() == '':
      return ''
    # Create arrows
    tokens     = x.split('|')
    token_lens = [ len(t) for t in tokens ]
    templates  = [ '{:^' + str(t) + '}' if t else '' for t in token_lens ]
    arrows     = [ t.format('V') for t in templates ]
    arrows_str = ' '.join( arrows )
    return arrows_str

  # Generate input and output arrows
  #
  # If the "simple" flag was given, just draw arrows for each port::
  #
  #            |                |                    |
  #            v                v                    v
  #
  # otherwise try to also add connectivity information:
  #
  #            + 7-baz-step
  #            + baz.txt
  #            |
  #            |                + 3-bar-step
  #            |                + bar.txt
  #            |                |
  #            |                |                    + 4-foo-step
  #            |                |                    + foo.txt
  #            |                |                    |
  #            V                V                    V
  #     +-----------------------------------------------+
  #     |                   step                        |
  #     +-----------------------------------------------+
  #              |           |
  #              |           V
  #              |           + 5-foofoo
  #              |           + foofoo.gds
  #              |           +
  #              |           + 6-bazbaz
  #              |           + bazbaz.gds
  #              V
  #              + 7-barbar
  #              + barbar.v
  #
  # The high level approach to building these strings is to use Python
  # string replace() with a max_count argument. The i'th input draws i-1
  # pipe symbols, a single plus sign, and then has text. We want the
  # output strings to replace from the right, but the replace() method
  # only takes a positive max_count. To get around this, for the i'th
  # output from the right, we reverse the string, replace the i with empty
  # space, replace a single plus sign, and then replace the rest with pipe
  # symbols before un-reversing the string and adding text.
  #

  input_arrows  = []
  output_arrows = []

  # Build the rows of arrows for the inputs and outputs
  #
  #            v                v                    v
  #

  arrows_i = arrowify( inputs_str  )
  arrows_o = arrowify( outputs_str )

  # Simple visualization

  if opts.simple:
    input_arrows.append( arrowify( inputs_str ).replace('V','|') )
    input_arrows.append( arrowify( inputs_str ) )
    output_arrows.append( arrowify( outputs_str ).replace('V','|') )
    output_arrows.append( arrowify( outputs_str ) )

  # Better visualization

  else:

    # Get a list of edges sorted in the order of the inputs and outputs

    edges_i_items = \
      sorted( edges_i.items(), key = lambda x: inputs.index(x[0]) )

    edges_o_items = \
      sorted( edges_o.items(), key = lambda x: outputs.index(x[0]) )

    # Build a list of bools for whether edges exist for each input/output

    edges_i_f = [ _[0] for _ in edges_i_items ]
    edges_o_f = [ _[0] for _ in edges_o_items ]

    # Mark which incoming edges exist

    exists_i = [ True if i in edges_i_f else False for i in inputs  ]

    exists_str = arrows_i
    for exists in exists_i:
      if exists : exists_str = exists_str.replace( 'V', 'x', 1 )
      else      : exists_str = exists_str.replace( 'V', ' ', 1 )

    # For the i'th incoming existing edge, replace the first i characters
    # with pipe symbols and the next one with a '+' sign. Then generate
    # the strings for printing. The i=2 strings will look like this:
    #
    #     |       |             + 4-cadence-innovus-place-route
    #     |       |             + constraints.tcl
    #     |       |             |
    #

    nth_input = 0
    for input_, steps in edges_i_items:
      l = exists_str.replace( 'x', '|', nth_input ).replace( 'x', '+', 1 )
      index = l.index( '+' )
      plus_line_str = l[:index] + '+'
      pipe_line_str = l[:index] + '|'
      for s in steps:
        input_arrows.append( plus_line_str + ' ' + s['step'] )
        input_arrows.append( plus_line_str + ' ' + s['f']    )
        input_arrows.append( pipe_line_str                   )
      nth_input += 1

    last_row = exists_str.replace( 'x', 'V' )
    input_arrows.append( last_row )

    # Mark which outgoing edges exist. Reverse the arrows string and the
    # list of existence booleans so we can later call string replace with
    # a positive max_count.

    exists_o = [ True if o in edges_o_f else False for o in outputs ]

    exists_str = arrows_o[::-1]
    for exists in exists_o[::-1]:
      if exists : exists_str = exists_str.replace( 'V', 'x', 1 )
      else      : exists_str = exists_str.replace( 'V', ' ', 1 )

    # Do nearly the same thing as before, but there may be more than one
    # output.
    #
    # For the i'th outgoing existing edge, replace the first i characters
    # with space, the next one with a '+' sign, and the remaining with
    # pipe symbols. Then generate the strings for printing. The i=2
    # strings will look like this:
    #
    #     |           |              |
    #     |           |              v
    #     |           |              + 9-cadence-innovus-place-route
    #     |           |              + design.namemap
    #     |           |              +
    #     |           |              + 10-cadence-innovus-place-route
    #     |           |              + design.namemap
    #     |           |
    #

    first_row = exists_str.replace( 'x', '|' )[::-1]
    output_arrows.append( first_row )

    nth_output = 0
    for output, steps in edges_o_items[::-1]:
      l = exists_str.replace( 'x', ' ', nth_output )
      l = l.replace( 'x', '+', 1 )
      l = l.replace( 'x', '|' )
      l = l[::-1] # undo the reversal so we can write next to the '+'
      index = l.index( '+' )
      plus_line_str  = l[:index] + '+'
      arrow_line_str = l[:index] + 'V'
      for j, s in enumerate( \
                    sorted( steps, key=lambda x: \
                      int(x['step'].split('-')[0]) ) ):
        if j == 0:
          output_arrows.append( arrow_line_str                )
        else:
          output_arrows.append( plus_line_str                 )
        output_arrows.append( plus_line_str + ' ' + s['step'] )
        output_arrows.append( plus_line_str + ' ' + s['f']    )
      nth_output += 1

  # Lines
  #
  #     +---------------+    <-- dash_line_str
  #     |               |    <-- empty_line_str
  #

  dash_line_str       = '+' + '-'*(n_chars-2) + '+'
  empty_line_str      = '|' + ' '*(n_chars-2) + '|'
  center_template_str = '{:^' + str(n_chars) + '}'

  print()
  for _ in input_arrows:
    print( ' '*4 + _ )
  print( ' '*4 + dash_line_str  )
  print( ' '*4 + center_template_str.format( inputs_str ) )
  print( ' '*4 + dash_line_str  )
  print( ' '*4 + empty_line_str )
  print( ' '*4 + center_template_str.format( name_str ) )
  print( ' '*4 + empty_line_str )
  print( ' '*4 + dash_line_str  )
  print( ' '*4 + center_template_str.format( outputs_str ) )
  print( ' '*4 + dash_line_str  )
  for _ in output_arrows:
    print( ' '*4 + _ )
  print()

  #-----------------------------------------------------------------------
  # Parameters and Assertions
  #-----------------------------------------------------------------------

  # print_list

  def print_list( v, indent=0 ):
    for _ in v:
      print( ' '*indent + '- ' + str(_) )

  # print_list_section
  #
  # Prints something like:
  #
  #     Title
  #
  #       - foo
  #       - bar
  #       - baz
  #

  def print_list_section( key ):
    if key in data.keys():
      title = key.capitalize()
      print( title + '\n' )
      print_list( data[key] )
      print()

  # print_dict_section
  #
  # Prints something like:
  #
  #     Title
  #
  #     - foo : 1.0
  #     - bar : 2.0
  #     - baz :
  #       - one
  #       - two
  #

  def print_dict_section( key ):

    if key in data.keys():

      p_width      = max( [ len(p) for p in data[key] ] )
      template_str ='- {:' + str(p_width) + '} : {}'

      title = key.capitalize()
      print( title + '\n' )
      for k, v in data[key].items():
        if type(v) == list:
          print( '- ' + k + ':' )
          print_list( v, indent=2 )
        else:
          print( template_str.format( k, v ) )
      print()

  print_dict_section( 'parameters'     )
  #print_list_section( 'preconditions'  )
  #print_list_section( 'postconditions' )

  #-----------------------------------------------------------------------
  # Flags
  #-----------------------------------------------------------------------

  flags = {}

  if 'sandbox' in data.keys():
    flags = { 'sandbox' : data['sandbox'] }

  if flags:

    p_width      = max( [ len(f) for f in flags ] )
    template_str ='- {:' + str(p_width) + '} : {}'

    print( 'Flags\n' )
    for k, v in flags.items():
      print( template_str.format( k, v ) )
    print()

  #-----------------------------------------------------------------------
  # Source Directory
  #-----------------------------------------------------------------------

  print( 'Source Directory\n' )
  print( '- ' + data['source'] )
  print()


if __name__ == '__main__':
  main()

