package Ex0

import scala.math.max

import chisel3._
import chisel3.util.Counter

class MatMul(val rowDimsA: Int, val colDimsA: Int) extends Module {

  val io = IO(
    new Bundle {
      val dataInA     = Input(UInt(32.W))
      val dataInB     = Input(UInt(32.W))

      val dataOut     = Output(UInt(32.W))
      val outputValid = Output(Bool())
    }
  )


  val matrixA     = Module(new Matrix(rowDimsA, colDimsA)).io
  val matrixB     = Module(new Matrix(rowDimsA, colDimsA)).io
  val dotProdCalc = Module(new DotProd(colDimsA)).io

  val calculating = RegInit(false.B)
  val resultCol = Counter(rowDimsA)
  val col = Counter(colDimsA)
  val row = Counter(rowDimsA)

  matrixA.rowIdx      := row.value
  matrixA.colIdx      := col.value
  matrixA.dataIn      := io.dataInA
  matrixA.writeEnable := ~calculating

  matrixB.rowIdx      := Mux(calculating, resultCol.value, row.value)
  matrixB.colIdx      := col.value
  matrixB.dataIn      := io.dataInB
  matrixB.writeEnable := ~calculating

  dotProdCalc.dataInA := matrixA.dataOut
  dotProdCalc.dataInB := matrixB.dataOut

  io.dataOut := dotProdCalc.dataOut
  io.outputValid := dotProdCalc.outputValid & calculating

  when (col.inc()) {
    when (calculating) {
      when (resultCol.inc()) {
        when (row.inc()) {
          calculating := false.B
        }
      }
    }.otherwise {
      when (row.inc()) {
        calculating := true.B
      }
    }
  }
}
