#=========================================================================
# GCD FSM
#=========================================================================

from pymtl       import *
from pclib.ifcs  import InValRdyBundle, OutValRdyBundle
from pclib.rtl   import RegRst
from utils       import *

NUM_STATE       = 5
LOG_NUM_STATE   = 3

class fsm( Model ):

  # Constructor

  def __init__( s ):

    # Interface

    s.in_ = InValRdyBundle [nPE]( 1  )
    s.out = OutValRdyBundle[nPE]( inst_msg() )

    # Register definition
    s.state = RegRst(LOG_NUM_STATE, 0)

    # State definition
    LOAD_TEST_DATA  = 0
    REQ_SIZE        = 1
    RETURN_SIZE     = 2
    SIZE_COPY       = 3
    READ_ELEMENT_1  = 4
    READ_ELEMENT_2  = 5
    COMPARE         = 6
    COMPARE_RESULT  = 7
    DEC_COUNTER     = 8
    CHECK_COUNTER   = 9
    OUTER_LOOP_DEC  = 10
    EXIT_CHECK      = 11
    EXIT_TRANSFER   = 12
    EXIT_DEC        = 13
    END_CHECK       = 14

    # Mux selects
    s.in_mux_sel  = OutPort ( 1 )
    s.out_mux_sel = OutPort ( 1 )

    @s.combinational
    def combinational_module():

      #------------------------------------------------------------------
      # INIT_S0 state
      #------------------------------------------------------------------
      # read the first operand from memory

      if s.state.out == INIT_S0:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.state.in_.value = INIT_S0

        # LOAD to register 0
        s.out[0].msg.ctl.value = LOAD
        s.out[0].msg.des.value = 16 + 0

        if s.out[0].rdy: 
          s.out[0].val.value = 1
          s.state.in_.value  = INIT_S1


      #------------------------------------------------------------------
      # REQ_SIZE
      #------------------------------------------------------------------
      # read the size from memory

      if s.state.out == REQ_SIZE:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.state.in_.value = REQ_SIZE

        s.out[0].msg.ctl.value = MEMREQ
        s.out[0].msg.des.value = 0  # Check what to do with this
        s.out[0].msg.addr.value = s.read_pointer
        s.out[0].msg.type.value = 0  # read

        if s.out[0].rdy: 
          s.out[0].val.value = 1
          s.state.in_.value  = RETURN_SIZE
          s.read_pointer.value = s.read_pointer + 1

      #------------------------------------------------------------------
      # RETURN_SIZE state
      #------------------------------------------------------------------
      # get the size back from memory; also send to PE 1's reg 0

      if s.state.out == RETURN_SIZE:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.state.in_.value = RETURN_SIZE

        s.out[0].msg.ctl.value = MEMRESP
        s.out[0].msg.des.value = 16 + 8  # Store size in PE 0 reg 0, also send to right PE 

        if s.out[0].rdy: 
          s.out[0].val.value = 1
          s.state.in_.value  = RETURN_SIZE

      #------------------------------------------------------------------
      # SIZE_COPY state
      #------------------------------------------------------------------
      # copy the size into PE 0's reg 1 as well

      if s.state.out == SIZE_COPY:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.state.in_.value = SIZE_COPY


        if s.out[0].rdy: 
          s.out[0].val.value = 1
          s.state.in_.value  = READ_ELEMENT_1

      #------------------------------------------------------------------
      # READ_ELEMENT_1 state
      #------------------------------------------------------------------
      # load the first element of the array from the memory

      if s.state.out == READ_ELEMENT_1:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.out[0].msg.ctl.value  = MEMREQ
        s.out[0].msg.addr.value = s.read_pointer
        s.out[0].msg.type.value = 0  # read

        s.state.in_.value = READ_ELEMENT_1

        if s.out[0].rdy: 
          s.out[0].val.value = 1
          s.state.in_.value  = READ_ELEMENT_2
          s.read_pointer.value = s.read_pointer + 1

      #------------------------------------------------------------------
      # WAIT_RESP state
      #------------------------------------------------------------------
      # load the first element of the array from the memory

      if s.state.out == WAIT_RESP:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.out[0].msg.ctl.value  = MEMRESP

        s.state.in_.value = WAIT_RESP

        if s.out[0].rdy: 
          s.out[0].val.value = 1
          s.state.in_.value  = READ_ELEMENT_2

      #------------------------------------------------------------------
      # READ_ELEMENT_2 state
      #------------------------------------------------------------------
      # load the second element of the array from the memory (send request from PE 1

      if s.state.out == READ_ELEMENT_2:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.out[1].msg.ctl.value  = MEMREQ
        s.out[1].msg.addr.value = s.read_pointer
        s.out[1].msg.type.value = 0  # read

        s.state.in_.value = READ_ELEMENT_2

        if s.out[1].rdy: 
          s.out[1].val.value = 1
          s.state.in_.value  = RESP_COMPARE
          s.read_pointer.value = s.read_pointer + 1

      #------------------------------------------------------------------
      # RESP_COMPARE state
      #------------------------------------------------------------------
      # load the first element of the array from the memory, compare in PE0 in parallel

      if s.state.out == RESP_COMPARE:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.out[1].msg.ctl.value  = MEMRESP  # Right PE accepts response
        s.out[1].msg.des.value  = 8 # send to left PE
        s.out[0].msg.ctl.value  = CGT   # Left PE performs compare greater than
        s.out[0].msg.src0.value = 2   # Pick first element from own reg file
        s.out[0].msg.src1.value = 4 + 1   # Pick second element from right PE

        s.state.in_.value = RESP_COMPARE

        if s.out[0].rdy: 
          s.out[0].val.value = 1
          s.state.in_.value  = WRITE_SMALL

      #------------------------------------------------------------------
      # WRITE_SMALL state
      #------------------------------------------------------------------
      # get comparision result, write smaller value to mem

      if s.state.out == WRITE_SMALL:
        
        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        if s.in_[0].val and (s.in_[0].msg == 1): 
          s.out[1].msg.ctl.value  = MEMRESP
          s.out[1].msg.des.value  = 8 # send to left PE

        s.state.in_.value = WRITE_SMALL

        if s.out[0].rdy: 
          s.out[0].val.value = 1
          s.state.in_.value  = DEC_INNER


      #------------------------------------------------------------------
      # CALC_S0 state
      #------------------------------------------------------------------
      # send out CGT instruction

      if s.state.out == CALC_S0:

        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.state.in_.value = CALC_S0

        s.in_[1].rdy.value = 1            # find out what this is
        # termination condition
        if s.in_[1].val and (s.in_[1].msg == 1): 
          s.state.in_.value = DONE

        else:
          # CGT r0 > r1 ?
          s.out[0].msg.ctl.value  = CGT
          s.out[0].msg.src0.value = 0
          s.out[0].msg.src1.value = 1

          if s.out[0].rdy: 
            s.out[0].val.value = 1
            s.state.in_.value  = CALC_S1


      #------------------------------------------------------------------
      # CALC_S1 state
      #------------------------------------------------------------------
      # send out SUB instruction

      if s.state.out == CALC_S1:

        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.state.in_.value  = CALC_S1
        s.in_[1].rdy.value = 1
        # termination condition
        if s.in_[1].val and (s.in_[1].msg == 1): 
          s.state.in_.value = DONE

        else:
          # send out SUB instruction based on the result from CGT
          if s.in_[0].val:
            if s.in_[0].msg == 1:  # r0 > r1: do r0-r1 --> r0
              s.out[0].msg.ctl.value  = SUB
              s.out[0].msg.src0.value = 0
              s.out[0].msg.src1.value = 1
              s.out[0].msg.des.value  = 16 + 4 + 0

            elif s.in_[0].msg == 0: # r0 < r1: do r1-r0 --> r1
              s.out[0].msg.ctl.value  = SUB
              s.out[0].msg.src0.value = 1
              s.out[0].msg.src1.value = 0
              s.out[0].msg.des.value  = 16 + 4 + 1

            # also send out CEZ instruction
            s.out[1].msg.ctl.value  = CEZ
            s.out[1].msg.src0.value = 5

            if s.out[0].rdy and s.out[1].rdy: 
              # only goes to the next state when both SUB and CEZ are successfully issued
              s.in_[0].rdy.value = 1

              s.out[0].val.value = 1
              s.out[1].val.value = 1
              s.state.in_.value  = CALC_S0


      #------------------------------------------------------------------
      # DONE state
      #------------------------------------------------------------------
      # send out STORE instruction

      if s.state.out == DONE:

        for x in range(nPE):
          s.in_[x].rdy.value  = 0
          s.out[x].val.value  = 0
          s.out[x].msg.value  = 0

        s.state.in_.value = DONE

        # STORE value of register 0 to memory
        s.out[0].msg.ctl.value  = STORE
        s.out[0].msg.src0.value = 1

        if s.out[0].rdy: 
          s.state.in_.value  = INIT_S0
          s.out[0].val.value = 1


  def line_trace ( s ):
    
    state2char = {
      0  : "INIT_S0",
      1  : "INIT_S1",
      2  : "CALC_S0",
      3  : "CALC_S1",
      4  : "DONE   ",
    }

    s.state_str = state2char[s.state.out.uint()]

    return "({})".format( s.state_str )
