package Ex0

import chisel3._

// This import statement makes the scala vector invisible, reducing confusion
import scala.collection.immutable.{ Vector => _ }

class Matrix(val rowsDim: Int, val colsDim: Int) extends Module {

  val io = IO(
    new Bundle {
      val colIdx      = Input(UInt(32.W))
      val rowIdx      = Input(UInt(32.W))
      val dataIn      = Input(UInt(32.W))
      val writeEnable = Input(Bool())

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

  // Creates a vector of zero-initialized registers
  val rows = VecInit(List.fill(rowsDim)(Module(new Vector(colsDim)).io))

  for(ii <- 0 until rowsDim){
    // It doesn't matter what we use to drive idx and dataIn, as long as
    // writeEnable is low, so we always use the input values.
    rows(ii).idx         := io.colIdx
    rows(ii).dataIn      := io.dataIn

    // Enable writing when the current row is selected
    rows(ii).writeEnable := io.writeEnable && (io.rowIdx === ii.U)
  }

  io.dataOut := rows(io.rowIdx).dataOut
}