"""
#=========================================================================
# 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 __future__ import absolute_import, division, print_function

from graphviz import Digraph

from pymtl3.dsl import CalleePort, NonBlockingCalleeIfc
from pymtl3.dsl.errors import UpblkCyclicError

from .BasePass import BasePass, PassMetadata
from .CLLineTracePass import CLLineTracePass
from .errors import PassOrderError


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.get_all_update_blocks() | top._dag.genblks

    # 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 = list( [ 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 CL line trace pass here.
    cl_trace = CLLineTracePass()
    schedule.insert( 0, cl_trace.process_component( top ) )

    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
