"""
#=========================================================================
# OpenLoopCLPass.py
#=========================================================================
# Generate a simple schedule (no Mamba techniques here) based on the
# DAG generated by some previous pass.
#
# Author : Shunning Jiang
# Date   : Apr 20, 2019
"""
from pymtl3.dsl import CalleePort, NonBlockingCalleeIfc
from pymtl3.dsl.errors import UpblkCyclicError

from ..BasePass import BasePass, PassMetadata
from ..errors import PassOrderError
from ..sim.SimpleSchedulePass import SimpleSchedulePass
from ..tracing.CLLineTracePass import CLLineTracePass
from ..tracing.CollectSignalPass import CollectSignalPass
from ..tracing.PrintWavePass import PrintWavePass
from ..tracing.VcdGenerationPass import VcdGenerationPass


class OpenLoopCLPass( BasePass ):
  def __call__( self, top ):
    if not hasattr( top._dag, "all_constraints" ):
      raise PassOrderError( "all_constraints" )

    top._sched = PassMetadata()

    self.schedule_with_top_level_callee( top )

  def schedule_with_top_level_callee( self, top ):

    # Construct the graph with top level callee port

    V = top._dag.final_upblks - top.get_all_update_ff()

    # We collect all top level callee ports/nonblocking callee interfaces

    top_level_callee_ports = top.get_all_object_filter(
      lambda x: isinstance(x, CalleePort) and x.get_host_component() is top )

    top_level_nb_ifcs = top.get_all_object_filter(
      lambda x: isinstance(x, NonBlockingCalleeIfc) and x.get_host_component() is top )

    # We still tell the top level
    method_callee_mapping = {}
    method_guard_mapping  = {}

    # First deal with normal calleeports. We map the actual method to the
    # callee port, and add the port to the vertex set
    for x in top_level_callee_ports:
      if not x.in_non_blocking_interface(): # Normal callee port
        V.add(x)
        assert x.method not in method_callee_mapping
        method_callee_mapping[x.method] = x

    # Then deal with non-blocking callee interfaces. Map the method of the
    # interface to the actual method and set up method-rdy mapping

    for x in top_level_nb_ifcs:
      V.add( x.method )
      method_guard_mapping [x.method] = x.rdy
      assert x.method.method not in method_callee_mapping
      method_callee_mapping[x.method.method] = x.method

    E   = top._dag.all_constraints
    Es  = { v: [] for v in V }
    InD = { v: 0  for v in V }

    for (u, v) in E: # u -> v
      InD[v] += 1
      Es [u].append( v )

    # In addition to existing constraints, we process the constraints that
    # involve top level callee ports. NOTE THAT we assume the user never
    # set the constraint on the actual method object inside the CalleePort
    # In GenDAGPass we already collect those constraints between update
    # blocks and ACTUAL METHODs. We use the ACTUAL METHOD to callee
    # mapping we set up above to avoid missing constraints.

    for (xx, yy) in top._dag.top_level_callee_constraints:

      if xx in method_callee_mapping:
        xx = method_callee_mapping[ xx ]

      if yy in method_callee_mapping:
        yy = method_callee_mapping[ yy ]

      InD[yy] += 1
      Es [xx].append( yy )
      E.add( (xx, yy) )

    # Perform topological sort for a serial schedule.

    schedule = []

    Q = [ v for v in V if not InD[v] ]

    while Q:
      import random
      random.shuffle(Q)

      # Prioritize update blocks instead of method
      # TODO make it O(logn) by balanced BST if needed ...

      u = None
      for i in range(len(Q)):
        if Q[i] not in method_guard_mapping:
          u = Q.pop(i)
          break

      if u is None:
        u = Q.pop()

      if u in method_guard_mapping:
        schedule.append( method_guard_mapping[ u ] )
      schedule.append( u )
      for v in Es[u]:
        InD[v] -= 1
        if not InD[v]:
          Q.append( v )

    # Shunning: we call line trace related pass here.
    CLLineTracePass()( top )
    CollectSignalPass()( top )
    VcdGenerationPass()( top )
    PrintWavePass()( top )
    # Shunning: we reuse ff and posedge schedules from SimpleSchedulePass
    simple = SimpleSchedulePass()
    simple.schedule_ff( top )
    simple.schedule_posedge_flip( top )

    # clear trace before any update block
    schedule.insert( 0, top._tracing.clear_cl_trace )

    # call ff blocks after normal schedule
    schedule.extend( top._sched.schedule_ff )

    # work on tracing
    if hasattr( top._tracing, "vcd_func" ):
      schedule.append( top._tracing.vcd_func )
    if hasattr( top._tracing, "collect_text_sigs" ):
      schedule.append( top._tracing.collect_text_sigs )

    schedule.extend( top._sched.schedule_posedge_flip )

    top._sched.new_schedule_index  = 0
    top._sched.orig_schedule_index = 0

    # Here we are trying to avoid scanning the original schedule that
    # contains methods because we will need isinstance in that case.
    # As a result we created a preprocessed list for execution and use
    # the dictionary to look up the new index of functions.

    # The last element is always line trace
    def print_line_trace():
      print(top.line_trace())

    schedule.append( print_line_trace )

    schedule_no_method = [ x for x in schedule if not isinstance(x, CalleePort) ]
    mapping = { x : i for i, x in enumerate( schedule_no_method ) }

    def wrap_method( top, method,
                     my_idx_new, next_idx_new, schedule_no_method,
                     my_idx_orig, next_idx_orig ):

      def actual_method( *args, **kwargs ):
        i = top._sched.new_schedule_index
        j = top._sched.orig_schedule_index

        if j > my_idx_orig:
          # This means we need to advance the current cycle to the end
          # and then normally execute until we get to the same point.
          # We use original schedule index to handle the case where
          # there are two consecutive methods.
          while i < len(schedule_no_method):
            schedule_no_method[i]()
            i += 1
          i = j = 0
          top.num_cycles_executed += 1

        # We advance from the current point i to the method's position in
        # the schedule without method just to execute those blocks
        while i < my_idx_new:
          schedule_no_method[i]()
          i += 1

        # Execute the method
        ret = method( *args, **kwargs )

        # Execute all update blocks before the next method. Note that if
        # there are several consecutive methods, my_idx_new is equal to next_idx_new
        while i < next_idx_new:
          schedule_no_method[i]()
          i += 1
        j = next_idx_orig

        if i == len(schedule_no_method):
          i = j = 0
          top.num_cycles_executed += 1

        top._sched.new_schedule_index = i
        top._sched.orig_schedule_index = j
        return ret

      return actual_method

    for i, x in enumerate( schedule ):
      if isinstance( x, CalleePort ):
        x.original_method = x.method

        # This is to find the next non-method block's position in the
        # original schedule
        next_func   = i + 1
        while next_func < len(schedule):
          if not isinstance( schedule[next_func], CalleePort ):
            break
          next_func += 1

        # Get the index of the block in the schedule without method
        # This always exists because we append a line trace at the end
        map_next_func = mapping[ schedule[next_func] ]

        # Get the index of the next method in the schedule without method
        next_method = i + 1
        while next_method < len(schedule):
          if isinstance( schedule[next_method], CalleePort ):
            break
          next_method += 1

        # If there is another method after me, I calculate the range of
        # blocks that I need to call and then stop before the user calls
        # the next method.
        if next_method < len(schedule):
          next_func = next_method
          while next_func < len(schedule):
            if not isinstance( schedule[next_func], CalleePort ):
              break
            next_func += 1
          # Get the index in the compacted schedule
          map_next_func_of_next_method = mapping[ schedule[next_func] ]
        else:
          map_next_func_of_next_method = len(schedule_no_method)

        x.method = wrap_method( top, x.method,
                                map_next_func, map_next_func_of_next_method,
                                schedule_no_method,
                                i, next_method )
    top.num_cycles_executed = 0

    return schedule
