package Core
import chisel3._
import chisel3.core.Input
import chisel3.iotesters.PeekPokeTester
import utilz._

/**
  The daisy multiplier creates two daisy grids, one transposed, and multiplies them.
  */
class daisyMultiplier(dims: Dims, dataWidth: Int) extends Module {

  val io = IO(new Bundle {

    val dataInA     = Input(UInt(dataWidth.W))
    val writeEnableA = Input(Bool())

    val dataInB     = Input(UInt(dataWidth.W))
    val writeEnableB = Input(Bool())

    val dataOut     = Output(UInt(dataWidth.W))
    val dataValid   = Output(Bool())
    val done        = Output(Bool())
  })


  /**
    Your implementation here
    */
  val rowCounter       = RegInit(UInt(8.W), 0.U)
  val colCounter       = RegInit(UInt(8.W), 0.U)

  val rowOutputCounter = RegInit(UInt(8.W), 0.U)

  val calculating      = RegInit(Bool(), false.B)
  val accumulator      = RegInit(UInt(8.W), 0.U)

  val resultReady      = RegInit(Bool(), false.B)


  /**
    Following the same principle behind the the vector matrix multiplication, by
    NOT transposing the dimensions.

    When writing a multiplier for a 3x2 matrix it's implicit that this means a
    3x2 matrix and 2x3, returning a 2x2 matrix. By not transposing the dimensions
    we get the same effect as in VecMat
    */
  val matrixA = Module(new daisyGrid(dims, dataWidth)).io
  val matrixB = Module(new daisyGrid(dims, dataWidth)).io




  matrixA.dataIn := io.dataInA
  matrixA.writeEnable := io.writeEnableA

  matrixB.dataIn := io.dataInB
  matrixB.writeEnable := io.writeEnableB

  ////////////////////////////////////////
  ////////////////////////////////////////
  /// Set up counter statemachine
  io.done := false.B

  when(colCounter === (dims.cols - 1).U){
    colCounter := 0.U

    when(rowCounter === (dims.rows - 1).U){
      rowCounter := 0.U
      calculating := true.B

      when(calculating === true.B){

        when(rowOutputCounter === (dims.rows - 1).U){
          io.done := true.B
        }.otherwise{
          rowOutputCounter := rowOutputCounter + 1.U
        }

      }

    }.otherwise{
      rowCounter := rowCounter + 1.U
    }
  }.otherwise{
    colCounter := colCounter + 1.U
  }



  ////////////////////////////////////////
  ////////////////////////////////////////
  /// set up reading patterns depending on if we are in calculating state or not
  when(calculating === true.B){
    matrixA.rowSelect := rowOutputCounter
  }.otherwise{
    matrixA.rowSelect := rowCounter
  }

  matrixB.rowSelect := rowCounter



  ////////////////////////////////////////
  ////////////////////////////////////////
  /// when we're in calculating mode, check if we have valid output
  resultReady := false.B
  io.dataValid := false.B
  when(calculating === true.B){
    when(colCounter === (dims.cols - 1).U){
      resultReady := true.B
    }
  }


  ////////////////////////////////////////
  ////////////////////////////////////////
  /// when we've got a result ready we need to flush the accumulator
  when(resultReady === true.B){
    // To flush our accumulator we simply disregard previous state
    accumulator := (matrixA.dataOut*matrixB.dataOut)
    io.dataValid := true.B
  }.otherwise{
    accumulator := accumulator + (matrixA.dataOut*matrixB.dataOut)
  }
  io.dataOut := accumulator
}