package dispatch

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.CoreState

class DispatchModel(NumCores: Int = 2, ThreadsPerCore: Int = 4) {
  val ThreadCountWidth = log2Ceil(ThreadsPerCore) + 1

  var core_start        = Array.fill(NumCores)(false)
  var core_reset        = Array.fill(NumCores)(true)
  var core_block_id     = Array.fill(NumCores)(0)
  var core_thread_count = Array.fill(NumCores)(ThreadsPerCore)
  var done              = false

  // interal states
  var start_execution   = false
  var blocks_done       = 0
  var blocks_dispatched = 0

  def update(
    start:        Boolean,
    thread_count: Int,
    core_done:    Array[Boolean]
  ): Unit = {
    if (!start) {
      return
    }

    if (!start_execution) {
      start_execution = true
      core_reset.map(_ => true)
    }

    val total_blocks = (thread_count + ThreadsPerCore - 1) / ThreadsPerCore
    // println(s"total_blocks in model: $total_blocks, thread_count = $thread_count, blocks_done = $blocks_done")
    if (blocks_done == total_blocks) {
      done = true
    }

    var run_cnt         = 0
    var prev_core_start = core_start.clone()
    for (i <- 0 until NumCores) {
      // println(
      //   s"#i = $i, blocks_dispatched = $blocks_dispatched, total_blocks = $total_blocks, core_reset(i) = ${core_reset(i)}"
      // )
      if (core_reset(i)) {
        core_reset(i) = false

        if (blocks_dispatched < total_blocks) {
          core_start(i) = true
          core_block_id(i) = blocks_dispatched
          core_thread_count(i) = if (blocks_dispatched == total_blocks - 1) {
            thread_count - (blocks_dispatched * ThreadsPerCore)

          // blocks_dispatched += 1
          run_cnt += 1
        }
      }

      // println(s"#i = $i, core_start(i) = ${prev_core_start(i)}, core_done(i) = ${core_done(i)}")
      if (prev_core_start(i) && core_done(i)) {
        core_reset(i) = true
        core_start(i) = false
        blocks_done += 1
      }
    }
    blocks_dispatched += run_cnt
  }
}

class DispatchSpec extends AnyFreeSpec with Matchers {
  "Test dispatch" - {
    "should match model behavior with random inputs" in {
      val NumCores:       Int = 2
      val ThreadsPerCore: Int = 4

      simulate(new Dispatch(NumCores, ThreadsPerCore)) { dut =>
        dut.reset.poke(true.B)
        dut.clock.step()
        dut.reset.poke(false.B)
        dut.clock.step()

        var cnt           = 0
        val rng           = new scala.util.Random(42) // 42 is the seed for reproducibility
        val dispatchModel = new DispatchModel()

        while (cnt < 10000) {
          val core_done    = Array.fill(NumCores)(rng.nextBoolean())
          val thread_count = rng.nextInt(ThreadsPerCore)

          dut.io.start.poke(true.B)
          dut.io.thread_count.poke(thread_count.U)
          for (i <- 0 until NumCores) {
            dut.io.core_done(i).poke(core_done(i).B)
          }

          dut.clock.step()

          // Update model
          dispatchModel.update(
            start = true,
            thread_count = thread_count,
            core_done = core_done
          )

          // Compare model with DUT
          dut.io.done.expect(dispatchModel.done.B)
          for (i <- 0 until NumCores) {
            dut.io.core_start(i).expect(dispatchModel.core_start(i).B)
            dut.io.core_reset(i).expect(dispatchModel.core_reset(i).B)
            dut.io.core_block_id(i).expect(dispatchModel.core_block_id(i).U)
            dut.io.core_thread_count(i).expect(dispatchModel.core_thread_count(i).U)
          }

          cnt += 1
        }
      }
    }
  }
}


package fetcher

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.CoreState
import statecode.FetcherState

class FetcherModel {
  var fetcher_state = FetcherState.IDLE
  var read_valid    = false
  var read_address  = 0
  var instruction   = 0

  def update(
    core_state:     CoreState.Type,
    current_pc:     Int,
    mem_read_ready: Boolean,
    mem_read_data:  Int
  ) = {
    // State machine logic matching the Fetcher module
    fetcher_state match {
      case FetcherState.IDLE =>

          instruction = mem_read_data
          read_valid = false
        }

      case FetcherState.FETCHED =>
        if (core_state == CoreState.DECODE) {
          fetcher_state = FetcherState.IDLE
        }
    }
  }

  def reset(): Unit = {
    fetcher_state = FetcherState.IDLE
    read_valid = false
    read_address = 0
    instruction = 0
  }
}

class FetcherSpec extends AnyFreeSpec with Matchers {
  "Test fetcher" - {
    "should match model behavior with random inputs" in {
      val ProgramMemAddrBits: Int = 8
      val ProgramMemDataBits: Int = 16
      simulate(new Fetcher(ProgramMemAddrBits, ProgramMemDataBits)) { dut =>
        dut.reset.poke(true.B)
        dut.clock.step()
        dut.reset.poke(false.B)
        dut.clock.step()

        def randomCoreState(rng: scala.util.Random): CoreState.Type = {
          val v = rng.nextInt(3)
          v match {
            case 0 => CoreState.FETCH
            case 1 => CoreState.DECODE
            case 2 => CoreState.EXECUTE
          }
        }

        var cnt          = 0
        val rng          = new scala.util.Random(42) // 42 is the seed for reproducibility
        val fetcherModel = new FetcherModel()

        while (cnt < 10000) {
          val core_state     = randomCoreState(rng)
          val current_pc     = rng.nextInt(256)
          val mem_read_ready = rng.nextInt(2) == 1
          val mem_read_data  = rng.nextInt(65536) // 2^16 for 16-bit data

          dut.io.core_state.poke(core_state)
          dut.io.current_pc.poke(current_pc.U)
          dut.io.mem_read_address_sender.ready.poke(mem_read_ready.B)
          dut.io.mem_read_data.poke(mem_read_data.U)

          dut.clock.step()

          // Update model
          fetcherModel.update(
            core_state = core_state,
            current_pc = current_pc,
            mem_read_ready = mem_read_ready,
            mem_read_data = mem_read_data
          )

          // Compare model with DUT
          assert(dut.io.fetcher_state.peekValue().asBigInt == fetcherModel.fetcher_state.litValue)
          dut.io.mem_read_address_sender.valid.expect(fetcherModel.read_valid.B)
          dut.io.mem_read_address_sender.bits.expect(fetcherModel.read_address.U)
          dut.io.instruction.expect(fetcherModel.instruction.U)

          cnt += 1
        }
      }
    }
  }
}

        if (core_state == CoreState.FETCH) {
          fetcher_state = FetcherState.FETCHING
          read_valid = true
          read_address = current_pc
        }


// See README.md for license details.

package gcd

import chisel3._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

/** This is a trivial example of how to run this Specification From within sbt use:
  * {{{
  * testOnly gcd.GCDSpec
  * }}}
  * From a terminal shell use:
  * {{{
  * sbt 'testOnly gcd.GCDSpec'
  * }}}
  * Testing from mill:
  * {{{
  * mill %NAME%.test.testOnly gcd.GCDSpec
  * }}}
  */
class GCDSpec extends AnyFreeSpec with Matchers {
  "Gcd should calculate proper greatest common denominator" in {
    simulate(new DecoupledGcd(16)) { dut =>
      val testValues = for {
        x <- 0 to 10
        y <- 0 to 10
      } yield (x, y)
      val inputSeq   = testValues.map { case (x, y) => (new GcdInputBundle(16)).Lit(_.value1 -> x.U, _.value2 -> y.U) }
      val resultSeq  = testValues.map { case (x, y) =>
        (new GcdOutputBundle(16)).Lit(_.value1 -> x.U, _.value2 -> y.U, _.gcd -> BigInt(x).gcd(BigInt(y)).U)
      }

      dut.reset.poke(true.B)
      dut.clock.step()
      dut.reset.poke(false.B)
      dut.clock.step()

      var sent, received, cycles: Int = 0
      while (sent != 100 && received != 100) {
        assert(cycles <= 1000, "timeout reached")

        if (sent < 100) {
          dut.input.valid.poke(true.B)
          dut.input.bits.value1.poke(testValues(sent)._1.U)
          dut.input.bits.value2.poke(testValues(sent)._2.U)
          if (dut.input.ready.peek().litToBoolean) {
            sent += 1
          }
        }

        if (received < 100) {
          dut.output.ready.poke(true.B)

        dut.clock.step()
        cycles += 1
      }
    }
  }
}

          if (dut.output.valid.peekValue().asBigInt == 1) {
            dut.output.bits.gcd.expect(BigInt(testValues(received)._1).gcd(testValues(received)._2))
            received += 1
          }


package gpu

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import core.CoreModel
import controller.ControllerModel
import dispatch.DispatchModel
import scala.collection.mutable.ArrayBuffer

class GpuModel(
  DataMemAddrBits:       Int = 8,
  DataMemDataBits:       Int = 8,
  DataMemNumChannels:    Int = 4,
  ProgramMemAddrBits:    Int = 8,
  ProgramMemDataBits:    Int = 16,
  ProgramMemNumChannels: Int = 1,
  NumCores:              Int = 2,
  ThreadsPerBlock: Int = 4) {
  // Submodule models
  private var dispatchModel = new DispatchModel(NumCores, ThreadsPerBlock)

  // Memory controllers
  private val NumLSUs                 = NumCores * ThreadsPerBlock
  private var dataMemoryController    = new ControllerModel(
    DataMemAddrBits,
    DataMemDataBits,
    NumCores * ThreadsPerBlock,
    DataMemNumChannels
  )
  private var programMemoryController = new ControllerModel(
    ProgramMemAddrBits,
    ProgramMemDataBits,
    NumCores,
    ProgramMemNumChannels
  )

  // Core models
  private var coreModels = Array.fill(NumCores)(
    new CoreModel(
      DataMemAddrBits,
      DataMemDataBits,
      ProgramMemAddrBits,
      ProgramMemDataBits,
      ThreadsPerBlock
    )
  )

  // Last cycle datas
  // Device Control Register
  var drc_thread_count = 0

  // Data memory controller outputs
  var data_consumer_read_ready                                       = Array.fill(NumLSUs)(false)
  var data_consumer_read_data                                        = Array.fill(NumLSUs)(0)
  var data_consumer_write_ready                                      = Array.fill(NumLSUs)(false)
  def save_data_memory_controller(model: controller.ControllerModel) = {
    data_consumer_read_ready = model.consumer_read_ready
    data_consumer_read_data = model.consumer_read_data
    data_consumer_write_ready = model.consumer_write_ready
  }

  // Program memory controller without write outputs
  var program_consumer_read_ready                                       = Array.fill(NumCores)(false)
  var program_consumer_read_data                                        = Array.fill(NumCores)(0)
  def save_program_memory_controller(model: controller.ControllerModel) = {
    program_consumer_read_ready = model.consumer_read_ready
    program_consumer_read_data = model.consumer_read_data
  }

  // Dispatcher outputs
  var dispatcher_core_start                          = Array.fill(NumCores)(false)
  var dispatcher_core_block_id                       = Array.fill(NumCores)(0)
  var dispatcher_core_thread_count                   = Array.fill(NumCores)(ThreadsPerBlock)
  def save_dispatcher(model: dispatch.DispatchModel) = {
    dispatcher_core_start = model.core_start
    dispatcher_core_block_id = model.core_block_id
    dispatcher_core_thread_count = model.core_thread_count
  }

  // Cores outputs 8/8
  var cores_done                                 = Array.fill(NumCores)(false)
  var cores_program_mem_read_address_valid       = Array.fill(NumCores)(false)
  var cores_program_mem_read_address_bits        = Array.fill(NumCores)(0)
  var cores_data_mem_read_address_valid          = Array.fill(NumCores)(Array.fill(ThreadsPerBlock)(false))
  var cores_data_mem_read_address_bits           = Array.fill(NumCores)(Array.fill(ThreadsPerBlock)(0))
  var cores_data_mem_write_address_valid         = Array.fill(NumCores)(Array.fill(ThreadsPerBlock)(false))
  var cores_data_mem_write_address_bits          = Array.fill(NumCores)(Array.fill(ThreadsPerBlock)(0))
  var cores_data_mem_write_data_bits             = Array.fill(NumCores)(Array.fill(ThreadsPerBlock)(0))
  def save_core(idx: Int, model: core.CoreModel) = {
    cores_done(idx) = model.done
    cores_program_mem_read_address_valid(idx) = model.program_mem_read_address_valid
    cores_program_mem_read_address_bits(idx) = model.program_mem_read_address_bits
    cores_data_mem_read_address_valid(idx) = model.data_mem_read_address_valid
    cores_data_mem_read_address_bits(idx) = model.data_mem_read_address_bits
    cores_data_mem_write_address_valid(idx) = model.data_mem_write_address_valid
    cores_data_mem_write_address_bits(idx) = model.data_mem_write_address_bits
    cores_data_mem_write_data_bits(idx) = model.data_mem_write_data_bits
  }

  // GPU outputs 8/8
  var gpu_done                   = false
  var gpu_program_mem_read_valid = Array.fill(ProgramMemNumChannels)(false)
  var gpu_program_mem_read_bits  = Array.fill(ProgramMemNumChannels)(0)
  var gpu_data_mem_read_valid    = Array.fill(DataMemNumChannels)(false)
  var gpu_data_mem_read_bits     = Array.fill(DataMemNumChannels)(0)
  var gpu_data_mem_write_valid   = Array.fill(DataMemNumChannels)(false)
  var gpu_data_mem_write_address = Array.fill(DataMemNumChannels)(0)
  var gpu_data_mem_write_data    = Array.fill(DataMemNumChannels)(0)

  def update(
    // Control inputs
    start:                       Boolean,
    device_control_write_enable: Boolean,
    device_control_data:         Int,

    // Memory inputs
    program_mem_read_ready: Array[Boolean],
    program_mem_read_data:  Array[Int],
    data_mem_read_ready:    Array[Boolean],
    data_mem_read_data:     Array[Int],
    data_mem_write_ready:   Array[Boolean]
  ): Unit = {
    // Update dispatcher
    val core_done = coreModels.map(_.done)
    dispatchModel.update(
      start = start,
      thread_count = drc_thread_count,
      core_done = cores_done
    )

    // Update memory controllers
    dataMemoryController.update(
      consumer_read_valid = cores_data_mem_read_address_valid.flatten,
      consumer_read_addr = cores_data_mem_read_address_bits.flatten,
      consumer_write_valid = cores_data_mem_write_address_valid.flatten,
      consumer_write_addr = cores_data_mem_write_address_bits.flatten,
      consumer_write_data = cores_data_mem_write_data_bits.flatten,
      mem_read_ready = data_mem_read_ready,
      mem_read_data = data_mem_read_data,
      mem_write_ready = data_mem_write_ready
    )

    programMemoryController.update(
      consumer_read_valid = cores_program_mem_read_address_valid,
      consumer_read_addr = cores_program_mem_read_address_bits,
      consumer_write_valid = Array.fill(NumCores)(false),
      consumer_write_addr = Array.fill(NumCores)(0),
      consumer_write_data = Array.fill(NumCores)(0),
      mem_read_ready = program_mem_read_ready,
      mem_read_data = program_mem_read_data,
      mem_write_ready = Array.fill(DataMemNumChannels)(false)
    )

    // Update each core
    for (i <- 0 until NumCores) {
      // Update core model
      coreModels(i).update(
        start = dispatcher_core_start(i),
        block_id = dispatcher_core_block_id(i),
        thread_count = dispatcher_core_thread_count(i),
        program_mem_read_ready = program_consumer_read_ready(i),
        program_mem_read_data = program_consumer_read_data(i),
        data_mem_read_data = data_consumer_read_data.slice(i * ThreadsPerBlock, (i + 1) * ThreadsPerBlock),
        data_mem_read_ready = data_consumer_read_ready.slice(i * ThreadsPerBlock, (i + 1) * ThreadsPerBlock),
        data_mem_write_ready = data_consumer_write_ready.slice(i * ThreadsPerBlock, (i + 1) * ThreadsPerBlock)
      )
      save_core(i, coreModels(i))
    }

    // Update outputs
    gpu_done = dispatchModel.done
    gpu_program_mem_read_valid = coreModels.map(_.program_mem_read_address_valid)
    gpu_program_mem_read_bits = coreModels.map(_.program_mem_read_address_bits)
    gpu_data_mem_read_valid = dataMemoryController.mem_read_valid
    gpu_data_mem_read_bits = dataMemoryController.mem_read_address
    gpu_data_mem_write_valid = dataMemoryController.mem_write_valid
    gpu_data_mem_write_address = dataMemoryController.mem_write_address
    gpu_data_mem_write_data = dataMemoryController.mem_write_data

    // update internal states
    if (device_control_write_enable) {
      drc_thread_count = device_control_data
    }
    save_program_memory_controller(programMemoryController)
    save_data_memory_controller(dataMemoryController)
    save_dispatcher(dispatchModel)
  }
}

class GpuSpec extends AnyFreeSpec with Matchers {
  "A cycle based GPU tester" - {
    "should match model behavior with random inputs" in {
      val DataMemAddrBits       = 8
      val DataMemDataBits       = 8
      val DataMemNumChannels    = 4
      val ProgramMemAddrBits    = 8
      val ProgramMemDataBits    = 16
      val ProgramMemNumChannels = 1
      val NumCores              = 2
      val ThreadsPerBlock       = 4

      simulate(
        new Gpu(
          DataMemAddrBits,
          DataMemDataBits,
          DataMemNumChannels,
          ProgramMemAddrBits,
          ProgramMemDataBits,
          ProgramMemNumChannels,
          NumCores,
          ThreadsPerBlock
        )
      ) { dut =>
        // Reset the DUT
        dut.reset.poke(true.B)
        dut.clock.step()
        dut.reset.poke(false.B)
        dut.clock.step()

        var cnt      = 0
        val rng      = new scala.util.Random(42) // 42 is the seed for reproducibility
        val gpuModel = new GpuModel(
          DataMemAddrBits,
          DataMemDataBits,
          DataMemNumChannels,
          ProgramMemAddrBits,
          ProgramMemDataBits,
          ProgramMemNumChannels,
          NumCores,
          ThreadsPerBlock
        )

        while (cnt < 10000) {
          // Generate random inputs
          val device_control_data = rng.nextInt(256)

          val program_mem_read_ready = Array.fill(ProgramMemNumChannels)(rng.nextBoolean())
          val program_mem_read_data  =
            Array.fill(ProgramMemNumChannels)(rng.nextInt(1 << ProgramMemDataBits))

          val data_mem_read_ready  = Array.fill(DataMemNumChannels)(rng.nextBoolean())
          val data_mem_read_data   = Array.fill(DataMemNumChannels)(rng.nextInt(1 << DataMemDataBits))
          val data_mem_write_ready = Array.fill(DataMemNumChannels)(rng.nextBoolean())

          // Apply inputs to DUT
          dut.io.start.poke(true.B)
          dut.io.device_control_write_enable.poke(true.B)
          dut.io.device_control_data.poke(device_control_data.U)

          for (i <- 0 until ProgramMemNumChannels) {
            dut.io.program_mem_read_sender(i).ready.poke(program_mem_read_ready(i).B)
            dut.io.program_mem_read_data(i).poke(program_mem_read_data(i).U)
          }

          for (i <- 0 until DataMemNumChannels) {
            dut.io.data_mem_read_sender(i).ready.poke(data_mem_read_ready(i).B)
            dut.io.data_mem_read_data(i).poke(data_mem_read_data(i).U)
            dut.io.data_mem_write_sender(i).ready.poke(data_mem_write_ready(i).B)
          }

          dut.clock.step()

          // Update model with same inputs
          gpuModel.update(
            start = true,
            device_control_write_enable = true,
            device_control_data = device_control_data,
            program_mem_read_ready = program_mem_read_ready,
            program_mem_read_data = program_mem_read_data,
            data_mem_read_ready = data_mem_read_ready,
            data_mem_read_data = data_mem_read_data,
            data_mem_write_ready = data_mem_write_ready
          )

          // Verify outputs
          dut.io.done.expect(gpuModel.gpu_done.B)

          // Verify program memory interface
          for (i <- 0 until ProgramMemNumChannels) {
            dut.io.program_mem_read_sender(i).valid.expect(gpuModel.gpu_program_mem_read_valid(i).B)
            dut.io.program_mem_read_sender(i).bits.expect(gpuModel.gpu_program_mem_read_bits(i).U)
          }

          // Verify data memory interface
          for (i <- 0 until DataMemNumChannels) {
            dut.io.data_mem_read_sender(i).valid.expect(gpuModel.gpu_data_mem_read_valid(i).B)
            dut.io.data_mem_read_sender(i).bits.expect(gpuModel.gpu_data_mem_read_bits(i).U)
            dut.io.data_mem_write_sender(i).valid.expect(gpuModel.gpu_data_mem_write_valid(i).B)
            dut.io.data_mem_write_sender(i).bits.address.expect(gpuModel.gpu_data_mem_write_address(i).U)
            dut.io.data_mem_write_sender(i).bits.data.expect(gpuModel.gpu_data_mem_write_data(i).U)
          }

          cnt += 1
        }
      }
    }
  }
}

class ProgramMemoryMock(program: Seq[Int]) {
  def step(dut: Gpu) = {
    dut.io.program_mem_read_sender.zip(dut.io.program_mem_read_data).foreach { case (sender, reader) =>
      val readValid = sender.valid.peek().litToBoolean
      if (readValid) {
        val addr = sender.bits.peekValue().asBigInt.toInt
        reader.poke(program(addr).U)
      }
      sender.ready.poke(readValid.B)
    }
  }
}

class DataMemoryMock(data: Seq[Int], ProgramMemAddrBits: Int = 8) {
  var mem = ArrayBuffer.fill(1 << ProgramMemAddrBits)(0)
  for (i <- 0 until data.length) {
    mem(i) = data(i)
  }

  def step(dut: Gpu) = {
    dut.io.data_mem_read_sender.zip(dut.io.data_mem_read_data).foreach { case (sender, reader) =>
      val readValid = sender.valid.peek().litToBoolean
      // Debug print for read interface
      println(

      )
      if (readValid) {
        val addr = sender.bits.peekValue().asBigInt.toInt
        reader.poke(mem(addr).U)
      }
      sender.ready.poke(readValid.B)
    }

    dut.io.data_mem_write_sender.foreach { sender =>
      val writeValid = sender.valid.peek().litToBoolean
      // Debug print for write interface
      println(s"Write: valid=${writeValid}, addr=${sender.bits.address.peekValue().asBigInt.toInt}")
      if (writeValid) {
        val addr = sender.bits.address.peekValue().asBigInt.toInt
        val data = sender.bits.data.peekValue().asBigInt.toInt
        mem(addr) = data
      }
      sender.ready.poke(writeValid.B)
    }
  }

  def getMemory: ArrayBuffer[Int] = mem
}

class GpuBehaviorSpec extends AnyFreeSpec with Matchers {
  "A behavior based GPU tester" - {
    "The final memory state should match GpuVM's data memory" in {
      val DataMemAddrBits       = 8
      val DataMemDataBits       = 8
      val DataMemNumChannels    = 4
      val ProgramMemAddrBits    = 8
      val ProgramMemDataBits    = 16
      val ProgramMemNumChannels = 1
      val NumCores              = 2
      val ThreadsPerBlock       = 4

      // val asmFile = sys.env.get("GPU_ASM")
      // if (asmFile.isEmpty) {
      //   println("Usage: Set the environment variable GPU_ASM to the path of the asm file.")
      //   System.exit(1)
      // }
      // val asm = scala.io.Source.fromFile(asmFile.get).getLines().mkString("\n")
      val asm = MatAddAsm.src

      val lexer  = new Lexer()
      val parser = new AsmParser()
      val vm     = new GpuVM(NumCores, ThreadsPerBlock, 1 << DataMemAddrBits)

      parser.parse(lexer.tokenize(asm))

      vm.init(parser.getDataArrays)
      vm.run(parser.getInstructions)

      val program_mem = new ProgramMemoryMock(MachineCodeEmitter.emit(asm))
      val data_mem    = new DataMemoryMock(parser.getDataArrays.flatten)

      simulate(
        new Gpu(
          DataMemAddrBits,
          DataMemDataBits,
          DataMemNumChannels,
          ProgramMemAddrBits,
          ProgramMemDataBits,
          ProgramMemNumChannels,
          NumCores,
          ThreadsPerBlock
        )
        // ,
        // Seq(WriteVcdAnnotation)
      ) { dut =>
        // Reset the DUT
        dut.reset.poke(true.B)
        dut.clock.step()
        dut.reset.poke(false.B)
        // dut.clock.step()

        val thread = 8
        dut.io.start.poke(true.B)
        dut.io.device_control_write_enable.poke(true.B)
        dut.io.device_control_data.poke(thread.U)
        println(s"DUT io.done = ${dut.io.done.peek().litToBoolean}")
        dut.clock.step()

        dut.io.device_control_write_enable.poke(false.B)

        println(s"DUT io.done = ${dut.io.done.peek().litToBoolean}")
        dut.clock.step()
        println(s"DUT io.done = ${dut.io.done.peek().litToBoolean}")

        var cnt = 0
        while (!dut.io.done.peek().litToBoolean && cnt < 10) {
          println("In while loop")
          program_mem.step(dut)
          data_mem.step(dut)
          dut.clock.step()
          cnt += 1
        }

        data_mem.getMemory.zip(vm.getMemory).zipWithIndex.foreach { case ((a, b), i) =>
          if (a != b) {
            println(s"Mismatch at index $i: $a != $b")
          }
        }
      }
    }
  }
}


package gpu

import scala.collection.mutable.ArrayBuffer

// First, let's define our Token types
sealed trait Token
object Token {
  // Keywords
  case object Thread extends Token
  case object Data   extends Token
//   case object Comment extends Token
  // case object Comma  extends Token
  case object Const  extends Token
  case object Nop    extends Token
  case object Brn    extends Token
  case object Brz    extends Token
  case object Brp    extends Token
  case object Cmp    extends Token
  case object Add    extends Token
  case object Sub    extends Token
  case object Mul    extends Token
  case object Div    extends Token
  case object Ldr    extends Token
  case object Str    extends Token

  // Special variables
  case object BlockIdx  extends Token
  case object BlockDim  extends Token
  case object ThreadIdx extends Token

  // Values
  case class Number(value: Int)    extends Token
  case class Immediate(value: Int) extends Token // Added for #NUMBER format
  case class Register(number: Int) extends Token // Added for Rxx format

  case class LabelDef(name: String) extends Token
  case class LabelUse(name: String) extends Token

  // For any unrecognized tokens
  case class Invalid(value: String) extends Token
}

class Lexer {
  def tokenize(input: String): Vector[Token] = {
    // Split the input into lines and process each line
    val tokens = input
      .replace(",", " ")
      .split("\n")
      .flatMap { line =>
        // Split each line into parts, handling comments
        val parts    = line.split(";", 2)
        val codeLine = parts(0).trim

        // println("*Debug: " + codeLine)

        if (codeLine.isEmpty) Vector.empty
        else {
          // Split the line into words while preserving special characters
          val words = codeLine.split("\\s+").filter(_.nonEmpty)
          words.map(tokenizeWord).toVector
        }
      }
      .toVector

    // Report any invalid tokens
    val invalidTokens = tokens.collect { case Token.Invalid(value) => value }
    if (invalidTokens.nonEmpty) {
      println("Error: Invalid tokens found:")
      invalidTokens.foreach(token => println(s" - $token"))
      Vector.empty
    } else {
      tokens
    }
  }

  private def tokenizeWord(word: String): Token = {
    import Token._

    word.toLowerCase match {
      case ".thread" | ".threads"                 => Thread
      case ".data"                                => Data
      // case ";"                                    => Comment
      // case ","                                    => Comma
      case "const"                                => Const
      case "nop"                                  => Nop
      case "brn"                                  => Brn
      case "brz"                                  => Brz
      case "brp"                                  => Brp
      case "cmp"                                  => Cmp
      case "add"                                  => Add
      case "sub"                                  => Sub
      case "mul"                                  => Mul
      case "div"                                  => Div
      case "ldr"                                  => Ldr
      case "str"                                  => Str
      case "ret"                                  => Ret
      case "%blockidx"                            => BlockIdx
      case "%blockdim"                            => BlockDim
      case "%threadidx"                           => ThreadIdx
      // Handle register numbers (R0-R99)
      case reg if reg.toLowerCase.startsWith("r") =>
        try {
          val number = reg.substring(1).toInt
          if (number >= 0 && number <= 99) Register(reg.substring(1).toInt)
          else Invalid(reg)
        } catch {
          case _: NumberFormatException => Invalid(reg)
        }
      // Handle immediate numbers (#NUMBER)
      case imm if imm.startsWith("#")             =>
        try {
          val number = imm.substring(1, imm.length).toInt
          Immediate(number)
        } catch {
          case _: NumberFormatException => Invalid(imm)
        }
      case num if num.matches("-?\\d+")           => Number(num.toInt)
      case label if label.endsWith(":")           => LabelDef(label.dropRight(1))
      case use if use.matches("^[a-z_]*$")        => LabelUse(use)
      case other                                  => Invalid(other)
    }
  }
}

// Test the lexer
object LexerTest {
  def main(args: Array[String]): Unit = {
    val lexer = new Lexer()
    val input = """
        .threads 8
        .data 0 1 2 3 4 5 6 7          ; matrix A (1 x 8)
        .data 0 1 2 3 4 5 6 7          ; matrix B (1 x 8)

        MUL R0, %blockIdx, %blockDim
        ADD R0, R0, %threadIdx         ; i = blockIdx * blockDim + threadIdx

        CONST R1, #0                   ; baseA (matrix A base address)
        CONST R2, #8                   ; baseB (matrix B base address)
        CONST R3, #16                  ; baseC (matrix C base address)

        ADD R4, R1, R0                 ; addr(A[i]) = baseA + i
        LDR R4, R4                     ; load A[i] from global memory

        ADD R5, R2, R0                 ; addr(B[i]) = baseB + i
        LDR R5, R5                     ; load B[i] from global memory

        ADD R6, R4, R15                 ; C[i] = A[i] + B[i]

        ADD R7, R3, R0                 ; addr(C[i]) = baseC + i
        STR R7, R6                     ; store C[i] in global memory

        RET                            ; end of kernel
    """

    val tokens = lexer.tokenize(input)
    tokens.foreach(println)
  }
}

sealed trait RegType
object RegType {
  // Keywords
  case class Imm(value: Int)        extends RegType
  case class Reg(value: Int)        extends RegType
  case class LabelUse(name: String) extends RegType
  // case class BlockIdx()             extends RegType
  // case class BlockDim()             extends RegType
  // case class ThreadIdx()            extends RegType
}

class Instruction(op: Token, args: Vector[RegType]) {
  private var _args:                  Vector[RegType] = args
  def getOp:                          Token           = op
  def getArgs:                        Vector[RegType] = _args
  def setArgs(args: Vector[RegType]): Unit            = {
    _args = args
  }
}

class AsmParser() {
  val BlockIdxOffset  = 13
  val BlockDimOffset  = 14
  val ThreadIdxOffset = 15

  private var debug:            Boolean = false
  def setDebug(debug: Boolean): Unit    = {
    this.debug = debug
  }

  private var tokens:       Vector[Token]       = Vector.empty
  private var threadCount:  Int                 = 0
  private var dataArrays:   Vector[Vector[Int]] = Vector.empty
  private var instructions: Vector[Instruction] = Vector.empty
  private var labels:       Map[String, Int]    = Map.empty
  private var idx:          Int                 = 0
  def peek():               Token               = tokens(idx)
  def consume():            Token               = {
    val tok = peek()
    idx += 1
    tok
  }
  def lookahead():          Token               = tokens(idx + 1)
  def isEof():              Boolean             = idx >= tokens.length

  def parseReg(reg: Token): RegType = {
    reg match {
      case Token.Register(number) => RegType.Reg(number)
      case Token.BlockIdx         => RegType.Reg(BlockIdxOffset)
      case Token.BlockDim         => RegType.Reg(BlockDimOffset)
      case Token.ThreadIdx        => RegType.Reg(ThreadIdxOffset)
      case _                      => throw new IllegalArgumentException(s"Invalid register expression in token: $reg")
    }
  }

  def parseBinayOp(op: Token): Instruction = {
    try {
      val (r1, r2) = (consume(), consume())
      new Instruction(
        op,
        Vector(parseReg(r1), parseReg(r2))
      )
    } catch {
      case _: NoSuchElementException => throw new IllegalArgumentException("Invalid binary operation expression")
    }
  }

  def parseTernaryOp(op: Token): Instruction = {
    try {
      val (r1, r2, r3) = (consume(), consume(), consume())
      new Instruction(
        op,
        Vector(parseReg(r1), parseReg(r2), parseReg(r3))
      )
    } catch {
      case _: NoSuchElementException => throw new IllegalArgumentException("Invalid ternary operation expression")
    }
  }

  def parse(tokens: Vector[Token]): Unit = {
    this.tokens = tokens
    while (!isEof()) {
      val tok = consume()
      if (debug) {
        println(s"*Debug current token: $tok")
      }
      tok match {
        case Token.Thread         => {
          val num = consume()
          assert(num.isInstanceOf[Token.Number], "Invalid thread count")
          threadCount = num.asInstanceOf[Token.Number].value
        }
        case Token.Data           => {
          var array = Vector.empty[Int]
          while (peek().isInstanceOf[Token.Number]) {
            val n = consume()
            array = array :+ n.asInstanceOf[Token.Number].value
          }
          dataArrays = dataArrays :+ array
        }
        case Token.Const          => {
          val reg = consume()
          assert(reg.isInstanceOf[Token.Register], "Invalid Const expr")
          val imm = consume()
          assert(imm.isInstanceOf[Token.Immediate], "Invalid Const expr")

          val (r, i)      = (reg.asInstanceOf[Token.Register].number, imm.asInstanceOf[Token.Immediate].value)
          val instruction = new Instruction(Token.Const, Vector(RegType.Reg(r), RegType.Imm(i)))
          instructions = instructions :+ instruction
        }
        case Token.Nop            => {
          val instruction = new Instruction(Token.Nop, Vector.empty)
          instructions = instructions :+ instruction
        }
        case Token.Brn            => {
          val label       = consume()
          assert(label.isInstanceOf[Token.LabelUse], "Invalid Brn expr")
          val instruction =
            new Instruction(Token.Brn, Vector(RegType.LabelUse(label.asInstanceOf[Token.LabelUse].name)))
          instructions = instructions :+ instruction
        }
        case Token.Brz            => {
          val label       = consume()
          assert(label.isInstanceOf[Token.LabelUse], "Invalid Brz expr")
          val instruction =
            new Instruction(Token.Brz, Vector(RegType.LabelUse(label.asInstanceOf[Token.LabelUse].name)))
          instructions = instructions :+ instruction
        }
        case Token.Brp            => {
          val label       = consume()
          assert(label.isInstanceOf[Token.LabelUse], "Invalid Brp expr")
          val instruction =
            new Instruction(Token.Brp, Vector(RegType.LabelUse(label.asInstanceOf[Token.LabelUse].name)))
          instructions = instructions :+ instruction
        }
        case Token.LabelDef(name) => {
          labels = labels + (name -> instructions.length)
        }
        case Token.Cmp            => {
          instructions = instructions :+ parseBinayOp(Token.Cmp)
        }
        case Token.Add            => {
          instructions = instructions :+ parseTernaryOp(Token.Add)
        }
        case Token.Sub            => {
          instructions = instructions :+ parseTernaryOp(Token.Sub)
        }
        case Token.Mul            => {
          instructions = instructions :+ parseTernaryOp(Token.Mul)
        }
        case Token.Div            => {
          instructions = instructions :+ parseTernaryOp(Token.Div)
        }
        case Token.Ldr            => {
          instructions = instructions :+ parseBinayOp(Token.Ldr)
        }
        case Token.Str            => {
          instructions = instructions :+ parseBinayOp(Token.Str)
        }
        case Token.Ret            => {
          instructions = instructions :+ new Instruction(Token.Ret, Vector.empty)
        }
        case other                => {
          println(s"Unrecognized token: $other")
        }
      }
    }

    // post process:
    // Turn labels into immediate values
    instructions
      .filter(inst =>
        inst.getOp == Token.Brn || inst.getOp == Token.Brn || inst.getOp == Token.Brz || inst.getOp == Token.Brp
      )
      .foreach { inst =>
        val label = inst.getArgs.head.asInstanceOf[RegType.LabelUse].name
        val idx   = labels(label)
        inst.setArgs(Vector(RegType.Imm(idx)))
      }

  }

  // Getter methods
  def getThreadCount:  Int                 = threadCount
  def getDataArrays:   Seq[Seq[Int]] = dataArrays
  def getInstructions: Seq[Instruction] = instructions
}

object MatAddAsm {
  val src = """
        .threads 8
        .data 0 1 2 3 4 5 6 7          ; matrix A (1 x 8)
        .data 0 1 2 3 4 5 6 7          ; matrix B (1 x 8)

        MUL R0, %blockIdx, %blockDim
        ADD R0, R0, %threadIdx         ; i = blockIdx * blockDim + threadIdx

        CONST R1, #0                   ; baseA (matrix A base address)
        CONST R2, #8                   ; baseB (matrix B base address)
        CONST R3, #16                  ; baseC (matrix C base address)

        ADD R4, R1, R0                 ; addr(A[i]) = baseA + i
        LDR R4, R4                     ; load A[i] from global memory

        ADD R5, R2, R0                 ; addr(B[i]) = baseB + i
        LDR R5, R5                     ; load B[i] from global memory

        ADD R6, R4, R5                 ; C[i] = A[i] + B[i]

        ADD R7, R3, R0                 ; addr(C[i]) = baseC + i
        STR R7, R6                     ; store C[i] in global memory

        RET                            ; end of kernel
    """
}

object MatMulAsm {
  val src = """
    .threads 4
    .data 1 2 3 4                  ; matrix A (2 x 2)
    .data 1 2 3 4                  ; matrix B (2 x 2)

    MUL R0, %blockIdx, %blockDim
    ADD R0, R0, %threadIdx         ; i = blockIdx * blockDim + threadIdx

    CONST R1, #1                   ; increment
    CONST R2, #2                   ; N (matrix inner dimension)
    CONST R3, #0                   ; baseA (matrix A base address)
    CONST R4, #4                   ; baseB (matrix B base address)
    CONST R5, #8                   ; baseC (matrix C base address)

    DIV R6, R0, R2                 ; row = i // N
    MUL R7, R6, R2
    SUB R7, R0, R7                 ; col = i % N

    CONST R8, #0                   ; acc = 0
    CONST R9, #0                   ; k = 0

    LOOP:
      MUL R10, R6, R2
      ADD R10, R10, R9
      ADD R10, R10, R3             ; addr(A[i]) = row * N + k + baseA
      LDR R10, R10                 ; load A[i] from global memory

      MUL R11, R9, R2
      ADD R11, R11, R7
      ADD R11, R11, R4             ; addr(B[i]) = k * N + col + baseB
      LDR R11, R11                 ; load B[i] from global memory

      MUL R12, R10, R11
      ADD R8, R8, R12              ; acc = acc + A[i] * B[i]

      ADD R9, R9, R1               ; increment k

      CMP R9, R2
      BRn LOOP                    ; loop while k < N

    ADD R9, R5, R0                 ; addr(C[i]) = baseC + i
    STR R9, R8                     ; store C[i] in global memory

    RET                            ; end of kernel
    """
}

object AsmParserTest1 {
  def main(args: Array[String]): Unit = {
    val lexer  = new Lexer()
    val parser = new AsmParser()
    parser.parse(lexer.tokenize(MatAddAsm.src))

    println("Thread count: " + parser.getThreadCount)
    println("\nData arrays:")
    parser.getDataArrays.zipWithIndex.foreach { case (array, i) =>
      println(s"Array $i: ${array.mkString(", ")}")
    }
  }
}

object AsmParserTest2 {
  def main(args: Array[String]): Unit = {
    val lexer  = new Lexer()
    val tokens = lexer.tokenize(MatMulAsm.src)
    tokens.foreach(println)

    val parser = new AsmParser()
    println("\n################## Parse result ##################\n")

    parser.parse(tokens)

    println("Thread count: " + parser.getThreadCount)
    println("\nData arrays:")
    parser.getDataArrays.zipWithIndex.foreach { case (array, i) =>
      println(s"Array $i: ${array.mkString(", ")}")
    }

    parser.getInstructions.zipWithIndex.foreach { case (inst, i) =>
      println(s"Instruction $i: ${inst.getOp} ${inst.getArgs.mkString(", ")}")
    }
  }
}

class GpuVM(NumCores: Int = 2, ThreadsPerBlock: Int = 4, MemSize: Int = 256) {
  private var debug:            Boolean = false
  def setDebug(debug: Boolean): Unit    = {
    this.debug = debug
  }

  val NumOfThread = NumCores * ThreadsPerBlock

  private var registers: Vector[ArrayBuffer[Int]] = Vector.fill(NumOfThread)(ArrayBuffer.fill(16)(0))
  private var memory:    ArrayBuffer[Int]         = ArrayBuffer.empty

  private var pc:  Int                       = 0
  private var nzp: ArrayBuffer[Seq[Boolean]] = ArrayBuffer.fill(NumOfThread)(Seq.fill(3)(false))

  def getRegisters: Vector[ArrayBuffer[Int]] = registers
  def getMemory:    ArrayBuffer[Int]         = memory

  val BlockIdxOffset  = 13
  val BlockDimOffset  = 14
  val ThreadIdxOffset = 15

  def printRegister(threadIdx: Int): Unit = {
    print(s"Thread $threadIdx registers: [")
    // registers(threadIdx).map(v => f"$v%3d").mkString(", ")
    print(registers(threadIdx).foldLeft("")((acc, v) => acc + f"$v%3d, "))
    println("]")
  }

  def printAllRegisters(): Unit = {
    println("################## Registers ##################")
    for (threadIdx <- 0 until NumOfThread) {
      printRegister(threadIdx)
    }
    println()
  }

  def printMemory(): Unit = {
    println("***************** Memory *****************")
    memory.grouped(8).foreach(group => println(group.map(n => f"$n%4d").mkString(" ")))
    println()
  }

  def init(dataArrays: Seq[Seq[Int]]): Unit = {
    // Initialize blockIdx, blockDim, threadIdx for each thread's register file
    for (blockId <- 0 until NumCores) {
      for (threadId <- 0 until ThreadsPerBlock) {
        val threadIdx = blockId * ThreadsPerBlock + threadId
        registers(threadIdx)(BlockIdxOffset) = blockId
        registers(threadIdx)(BlockDimOffset) = ThreadsPerBlock
        registers(threadIdx)(ThreadIdxOffset) = threadId
      }
    }

    // copy data arrays to memory
    memory = ArrayBuffer.fill(MemSize)(0)
    var offset = 0
    for (array <- dataArrays) {
      for (v <- array) {
        memory(offset) = v
        offset += 1
      }
    }
  }

  def run(instructions: Seq[Instruction]): Unit = {
    while (pc < instructions.length) {
      val inst = instructions(pc)
      if (debug) {
        println(s"*Debug current instruction: ${inst.getOp} ${inst.getArgs.mkString(", ")}")
        printAllRegisters()
        printMemory()
        println("\n")
      }
      for (threadIdx <- 0 until NumOfThread) {
        def nextPc(): Int = {
          inst.getOp match {
            case Token.Nop   => {
              // do nothing
            }
            case Token.Brn   => {
              val labelIdx = inst.getArgs.head.asInstanceOf[RegType.Imm].value
              if (nzp(threadIdx)(0)) {
                return labelIdx
              }
            }
            case Token.Brz   => {
              val labelIdx = inst.getArgs.head.asInstanceOf[RegType.Imm].value
              if (nzp(threadIdx)(1)) {
                return labelIdx
              }
            }
            case Token.Brp   => {
              val labelIdx = inst.getArgs.head.asInstanceOf[RegType.Imm].value
              if (nzp(threadIdx)(2)) {
                return labelIdx
              }
            }
            case Token.Cmp   => {
              val s = inst.getArgs(0).asInstanceOf[RegType.Reg].value
              val t = inst.getArgs(1).asInstanceOf[RegType.Reg].value
              nzp(threadIdx) = Seq(
                registers(threadIdx)(s) < registers(threadIdx)(t),  // neg : idx 0
                registers(threadIdx)(s) == registers(threadIdx)(t), // zero: idx 1
                registers(threadIdx)(s) > registers(threadIdx)(t)   // pos : idx 2
              )
            }
            case Token.Add   => {
              val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
              val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
              val t = inst.getArgs(2).asInstanceOf[RegType.Reg].value
              registers(threadIdx)(d) = registers(threadIdx)(s) + registers(threadIdx)(t)
            }
            case Token.Sub   => {
              val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
              val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
              val t = inst.getArgs(2).asInstanceOf[RegType.Reg].value
              registers(threadIdx)(d) = registers(threadIdx)(s) - registers(threadIdx)(t)
            }
            case Token.Mul   => {
              val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
              val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
              val t = inst.getArgs(2).asInstanceOf[RegType.Reg].value
              registers(threadIdx)(d) = registers(threadIdx)(s) * registers(threadIdx)(t)
            }
            case Token.Div   => {
              val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
              val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
              val t = inst.getArgs(2).asInstanceOf[RegType.Reg].value
              registers(threadIdx)(d) = registers(threadIdx)(s) / registers(threadIdx)(t)
            }
            case Token.Ldr   => {
              val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
              val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
              registers(threadIdx)(d) = memory(registers(threadIdx)(s))
            }
            case Token.Str   => {
              val s = inst.getArgs(0).asInstanceOf[RegType.Reg].value
              val t = inst.getArgs(1).asInstanceOf[RegType.Reg].value
              // println(
              //   s"*Debug store: thread=$threadIdx, s=$s, t=$t, registers(threadIdx)(s)=${registers(threadIdx)(s)}, registers(threadIdx)(t)=${registers(threadIdx)(t)}"
              // )
              memory(registers(threadIdx)(s)) = registers(threadIdx)(t)
            }
            case Token.Const => {
              val (reg, imm) =
                (inst.getArgs.head.asInstanceOf[RegType.Reg].value, inst.getArgs.last.asInstanceOf[RegType.Imm].value)
              registers(threadIdx)(reg) = imm
            }
            case Token.Ret   => {
              return instructions.length
            }
            case _           => {
              throw new Exception(s"Unrecognized instruction: ${inst.getOp}")
            }
          }
          return pc + 1
        }
        val npc = nextPc()
        // for each instruction, only update pc once
        if (threadIdx == NumOfThread - 1) {
          pc = npc
        }
        if (debug) {
          println(s"Thread $threadIdx npc: $npc")
          printRegister(threadIdx)
          println(nzp(threadIdx))
          printMemory()
          println("\n")
        }
      }
    }
  }
}

object GpuVMAddTest {
  def main(args: Array[String]): Unit = {
    val lexer  = new Lexer()
    val tokens = lexer.tokenize(MatAddAsm.src)
    val parser = new AsmParser()
    parser.parse(tokens)

    println("################## MatAddTest ##################")
    println("################## Parse result ##################")
    parser.getDataArrays.foreach(println)

    println("################## VM init ##################")
    val vm = new GpuVM()
    // vm.setDebug(true)
    vm.init(parser.getDataArrays)
    vm.printAllRegisters()
    vm.printMemory()

    println("################## Run result ##################")
    vm.run(parser.getInstructions)

    vm.printAllRegisters()
    vm.printMemory()

    val ref = Seq( 0, 2, 4, 6, 8, 10, 12, 14)
    assert(vm.getMemory.slice(16, 24) == ref)
    println("**** GpuVMAddTest PASSED ****")
  }
}

object GpuVMMulTest {
  def main(args: Array[String]): Unit = {
    val lexer  = new Lexer()
    val tokens = lexer.tokenize(MatMulAsm.src)
    val parser = new AsmParser()
    parser.parse(tokens)

    println("################## MatMulTest ##################")
    println("################## Parse result ##################")
    parser.getDataArrays.foreach(println)

    println("################## VM init ##################")
    val vm = new GpuVM()
    // vm.setDebug(true)
    vm.init(parser.getDataArrays)
    vm.printAllRegisters()
    vm.printMemory()

    println("################## Run result ##################")
    vm.run(parser.getInstructions)

    vm.printAllRegisters()
    vm.printMemory()

    val ref = Seq( 7,   10 ,  15   ,22)
    assert(vm.getMemory.slice(8, 12) == ref)
    println("**** GpuVMMulTest PASSED ****")
  }
}

class MachineCodeEmitter() {
  def asmToMachineCode(asm: String): Seq[Int] = {
    val lexer        = new Lexer()
    val parser       = new AsmParser()
    parser.parse(lexer.tokenize(asm))
    val instructions = parser.getInstructions

    var machineCode = Seq.empty[Int]
    for (inst <- instructions) {
      inst.getOp match {
        case Token.Nop   => {
          machineCode = machineCode :+ 0
        }
        case Token.Brn   => {
          val labelIdx = inst.getArgs.head.asInstanceOf[RegType.Imm].value
          machineCode = machineCode :+ ((1 << 12) | (1 << 11) | labelIdx)
        }
        case Token.Brz   => {
          val labelIdx = inst.getArgs.head.asInstanceOf[RegType.Imm].value
          machineCode = machineCode :+ ((1 << 12) | (1 << 10) | labelIdx)
        }
        case Token.Brp   => {
          val labelIdx = inst.getArgs.head.asInstanceOf[RegType.Imm].value
          machineCode = machineCode :+ ((1 << 12) | (1 << 9) | labelIdx)
        }
        case Token.Cmp   => {
          val s = inst.getArgs(0).asInstanceOf[RegType.Reg].value
          val t = inst.getArgs(1).asInstanceOf[RegType.Reg].value
          machineCode = machineCode :+ ((1 << 13) | (s << 4) | t)
        }
        case Token.Add   => {
          val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
          val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
          val t = inst.getArgs(2).asInstanceOf[RegType.Reg].value
          machineCode = machineCode :+ ((3 << 12) | (d << 8) | (s << 4) | t)
        }
        case Token.Sub   => {
          val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
          val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
          val t = inst.getArgs(2).asInstanceOf[RegType.Reg].value
          machineCode = machineCode :+ ((4 << 12) | (d << 8) | (s << 4) | t)
        }
        case Token.Mul   => {
          val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
          val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
          val t = inst.getArgs(2).asInstanceOf[RegType.Reg].value
          machineCode = machineCode :+ ((5 << 12) | (d << 8) | (s << 4) | t)
        }
        case Token.Div   => {
          val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
          val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
          val t = inst.getArgs(2).asInstanceOf[RegType.Reg].value
          machineCode = machineCode :+ ((6 << 12) | (d << 8) | (s << 4) | t)
        }
        case Token.Ldr   => {
          val d = inst.getArgs(0).asInstanceOf[RegType.Reg].value
          val s = inst.getArgs(1).asInstanceOf[RegType.Reg].value
          machineCode = machineCode :+ ((7 << 12) | (d << 8) | (s << 4))
        }
        case Token.Str   => {
          val s = inst.getArgs(0).asInstanceOf[RegType.Reg].value
          val t = inst.getArgs(1).asInstanceOf[RegType.Reg].value
          machineCode = machineCode :+ ((8 << 12) | (s << 4) | t)
        }
        case Token.Const => {
          val (reg, imm) =
            (inst.getArgs.head.asInstanceOf[RegType.Reg].value, inst.getArgs.last.asInstanceOf[RegType.Imm].value)
          machineCode = machineCode :+ ((9 << 12) | (reg << 8) | imm)
        }
        case Token.Ret   => {
          machineCode = machineCode :+ ((0xf << 12))
        }
        case _           => {
          // throw new Exception(s"Unrecognized instruction: ${inst.getOp}")
        }
      }
    }

    machineCode
  }
}

object MachineCodeEmitterTest {
  def main(args: Array[String]): Unit = {
    def verify(src: Seq[Int], ref: Seq[Int]): Boolean = {
      if (src.length != ref.length) {
        println(s"src.length != ref.length: ${src.length} != ${ref.length}")
        return false
      }
      for (i <- 0 until src.length) {
        if (src(i) != ref(i)) {
          println(s"src($i) != ref($i): ${MachineCodeEmitter.toBinary(src(i))} != ${MachineCodeEmitter.toBinary(ref(i))}")
          return false
        }
      }
      return true
    }

    println("Testing MatAddAsm:")
    val machineCode = new MachineCodeEmitter()
    val addMC = machineCode.asmToMachineCode(MatAddAsm.src)
    addMC.zipWithIndex.foreach { case (ins, idx) => println(f"$idx: ${MachineCodeEmitter.toBinary(ins)}") }

    val addProgram = Seq(
        0b0101000011011110,
        0b0011000000001111,
        0b1001000100000000,
        0b1001001000001000,
        0b1001001100010000,
        0b0011010000010000,
        0b0111010001000000,
        0b0011010100100000,
        0b0111010101010000,
        0b0011011001000101,
        0b0011011100110000,
        0b1000000001110110,
        0b1111000000000000
    )

    assert(verify(addMC, addProgram))
    println("**** MachineCodeEmitter MatAddTest  PASSED ****")

    println("\n\nTesting MatMulAsm:")
    val mulMC = machineCode.asmToMachineCode(MatMulAsm.src)
    mulMC.zipWithIndex.foreach { case (ins, idx) => println(f"$idx: ${MachineCodeEmitter.toBinary(ins)}") }

    val mulProgram = Seq(
        0b0101000011011110,
        0b0011000000001111,
        0b1001000100000001,
        0b1001001000000010,
        0b1001001100000000,
        0b1001010000000100,
        0b1001010100001000,
        0b0110011000000010,
        0b0101011101100010,
        0b0100011100000111,
        0b1001100000000000,
        0b1001100100000000,
        0b0101101001100010,
        0b0011101010101001,
        0b0011101010100011,
        0b0111101010100000,
        0b0101101110010010,
        0b0011101110110111,
        0b0011101110110100,
        0b0111101110110000,
        0b0101110010101011,
        0b0011100010001100,
        0b0011100110010001,
        0b0010000010010010,
        0b0001100000001100,
        0b0011100101010000,
        0b1000000010011000,
        0b1111000000000000
    )
    assert(verify(mulMC, mulProgram))
    println("**** MachineCodeEmitter MatMulTest  PASSED ****")
  }
}

object MachineCodeEmitter {
  def toBinary(ins: Int): String = {
    def toDigit4(i: Int) =
      String.format("%4s", i.toBinaryString).replace(' ', '0')

    val (p1, p2, p3, p4) = (
      toDigit4((ins & 0xf000) >> 12),
      toDigit4((ins & 0x0f00) >> 8),
      toDigit4((ins & 0x00f0) >> 4),
      toDigit4((ins & 0x000f))
    )

    f"$p1 $p2 $p3 $p4"
  }

  def emit(asm: String): Seq[Int] = {
    val lexer  = new Lexer()
    val parser = new AsmParser()
    parser.parse(lexer.tokenize(asm))

    val emitter = new MachineCodeEmitter()
    val machineCode = emitter.asmToMachineCode(asm)
    machineCode
  }

  def main(args: Array[String]): Unit = {
    if (args.length < 1 || args(0).isEmpty) {
      println("Usage: Tiny-GPU MachineCodeEmitter <asm_file> [--idx: Print instructions with index]")
      return
    }

    val asmFile = args(0)
    val withIndex = args.length > 1 && args(1) == "--idx"

    // Read the ASM file
    val asm = scala.io.Source.fromFile(asmFile).getLines().mkString("\n")

    // Emit machine code
    val machineCode = emit(asm)

    // Print the machine code in binary format
    if (withIndex) {
      machineCode.zipWithIndex.foreach { case (ins, idx) =>
        println(f"$idx%4d: ${toBinary(ins)}")
      }
    } else {
      machineCode.foreach(ins => println(MachineCodeEmitter.toBinary(ins)))
    }
  }
}


package lsu

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.CoreState
import statecode.LSUState

class LsuModel {
  var lsu_state     = LSUState.IDLE
  var output_data   = 0
  var read_valid    = false
  var read_address  = 0
  var write_valid   = false
  var write_address = 0
  var write_data    = 0

  def update(
    enable:          Boolean,
    core_state:      CoreState.Type,
    read_enable:     Boolean,
    write_enable:    Boolean,
    rs:              Int,
    rt:              Int,
    mem_read_data:   Int,
    mem_read_ready:  Boolean,
    mem_write_ready: Boolean
  ) = {
    if (enable) {
      if (read_enable) {
        lsu_state match {
          case LSUState.IDLE       =>
            if (core_state == CoreState.REQUEST) {
              lsu_state = LSUState.REQUESTING
            }
          case LSUState.REQUESTING =>
            read_valid = true
            read_address = rs
            lsu_state = LSUState.WAITING
          case LSUState.WAITING    =>
            if (mem_read_ready) {
              read_valid = false
              output_data = mem_read_data
              lsu_state = LSUState.DONE
            }
          case LSUState.DONE       =>
            if (core_state == CoreState.UPDATE) {
              lsu_state = LSUState.IDLE
            }
        }
      }

      if (write_enable) {
        lsu_state match {
          case LSUState.IDLE       =>
            if (core_state == CoreState.REQUEST) {
              lsu_state = LSUState.REQUESTING
            }
          case LSUState.REQUESTING =>
            write_valid = true
            write_address = rs
            write_data = rt
            lsu_state = LSUState.WAITING
          case LSUState.WAITING    =>
            if (mem_write_ready) {
              write_valid = false
              lsu_state = LSUState.DONE
            }
          case LSUState.DONE       =>
            if (core_state == CoreState.UPDATE) {
              lsu_state = LSUState.IDLE
            }
        }
      }
    }
  }
}

class LsuSpec extends AnyFreeSpec with Matchers {
  "Test Load/Store Unit" - {
    "should match model behavior with random inputs" in {
      simulate(new MemLoadStoreUnit) { dut =>
        dut.reset.poke(true.B)
        dut.clock.step()
        dut.reset.poke(false.B)
        dut.clock.step()

        def randomCoreState(rng: scala.util.Random): CoreState.Type = {
          val v = rng.nextInt(2)
          v match {
            case 0 => CoreState.REQUEST
            case 1 => CoreState.UPDATE
          }
        }

        var cnt      = 0
        val rng      = new scala.util.Random(42) // 42 is the seed for reproducibility
        val lsuModel = new LsuModel()

        while (cnt < 10000) {
          val enable          = rng.nextBoolean()
          val core_state      = randomCoreState(rng)
          val read_enable     = rng.nextBoolean()

          val rs              = rng.nextInt(256)
          val rt              = rng.nextInt(256)
          val mem_read_data   = rng.nextInt(256)
          val mem_read_ready  = rng.nextBoolean()
          val mem_write_ready = rng.nextBoolean()

          dut.io.enable.poke(enable.B)
          dut.io.core_state.poke(core_state)
          dut.io.mem_rw_enable.read_enable.poke(read_enable.B)
          dut.io.mem_rw_enable.write_enable.poke(write_enable.B)
          dut.io.reg_in.rs.poke(rs.U)
          dut.io.reg_in.rt.poke(rt.U)
          dut.io.mem_read_data.poke(mem_read_data.U)
          dut.io.mem_read_address_sender.ready.poke(mem_read_ready.B)
          dut.io.mem_write_sender.ready.poke(mem_write_ready.B)

          dut.clock.step()

          lsuModel.update(
            enable = enable,
            core_state = core_state,
            read_enable = read_enable,
            write_enable = write_enable,
            rs = rs,
            rt = rt,
            mem_read_data = mem_read_data,
            mem_read_ready = mem_read_ready,
            mem_write_ready = mem_write_ready
          )

          // dut.io.lsu_state.expect(LSUState.IDLE)
          assert(dut.io.lsu_state.peekValue().asBigInt == lsuModel.lsu_state.litValue)
          dut.io.lsu_out.expect(lsuModel.output_data.U)
          dut.io.mem_read_address_sender.valid.expect(lsuModel.read_valid.B)
          if (lsuModel.read_valid) {
            dut.io.mem_read_address_sender.bits.expect(lsuModel.read_address.U)
          }
          dut.io.mem_write_sender.valid.expect(lsuModel.write_valid.B)
          if (lsuModel.write_valid) {
            dut.io.mem_write_sender.bits.address.expect(lsuModel.write_address.U)
            dut.io.mem_write_sender.bits.data.expect(lsuModel.write_data.U)
          }

          cnt += 1
        }
      }
    }
  }
}

val write_enable    = if (read_enable) false else rng.nextBoolean() // Ensure not both enabled


package pc

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.CoreState

class PCModel(DataMemWidth: Int = 8, MemAddrWidth: Int = 8) {
  var nzp     = 0
  var next_pc = 0
  val max_pc  = 1 << MemAddrWidth

  def update(
    enable:                   Boolean,
    core_state:               CoreState.Type,
    decoded_nzp:              Int,
    decoded_immediate:        Int,
    decoded_nzp_write_enable: Boolean,
    decoded_pc_mux:           Boolean,
    alu_out:                  Int,
    current_pc:               Int
  ) = {
    if (enable) {
      // Execute state logic
      if (core_state == CoreState.EXECUTE) {
        if (decoded_pc_mux) {
          if ((nzp & decoded_nzp) != 0) {
            // Branch taken
            next_pc = decoded_immediate
          } else {
            // Branch not taken
            next_pc = (current_pc + 1) % max_pc
          }
        } else {
          // Normal PC increment
          next_pc = (current_pc + 1) % max_pc
        }
      }

      // Update state logic
      if (core_state == CoreState.UPDATE) {
        if (decoded_nzp_write_enable) {
          // Update NZP register from ALU output
          nzp = alu_out & 0x7 // Take only lowest 3 bits
        }
      }
    }
  }

  def reset(): Unit = {
    nzp = 0
    next_pc = 0
  }

  // Getter for testing

}

class PCSpec extends AnyFreeSpec with Matchers {
  "Test Program Counter" in {
    val DataMemWidth: Int = 8
    val MemAddrWidth: Int = 8
    simulate(new ProgramCounter(DataMemWidth, MemAddrWidth)) { dut =>
      dut.reset.poke(true.B)
      dut.clock.step()
      dut.reset.poke(false.B)
      dut.clock.step()

      def randomCoreState(rng: scala.util.Random): CoreState.Type = {
        val v = rng.nextInt(2)
        v match {
          case 0 => CoreState.EXECUTE
          case 1 => CoreState.UPDATE
        }
      }

      dut.io.enable.poke(true.B)

      var cnt     = 0
      val rng     = new scala.util.Random(42) // 42 is the seed for reproducibility
      val pcModel = new PCModel()
      while (cnt < 10000) {
        val core_state               = randomCoreState(rng)
        val decoded_nzp              = rng.nextInt(4)
        val decoded_immediate        = rng.nextInt(256)
        val decoded_nzp_write_enable = rng.nextInt(2) == 1
        val decoded_pc_mux           = rng.nextInt(2) == 1
        val alu_out                  = rng.nextInt(256)
        val current_pc               = rng.nextInt(256)

        dut.io.core_state.poke(core_state)
        dut.io.decoded_nzp.poke(decoded_nzp.U)
        dut.io.decoded_immediate.poke(decoded_immediate.U)
        dut.io.decoded_nzp_write_enable.poke(decoded_nzp_write_enable.B)
        dut.io.decoded_pc_mux.poke(decoded_pc_mux.B)
        dut.io.alu_out.poke(alu_out.U)
        dut.io.current_pc.poke(current_pc.U)

        dut.clock.step()

        pcModel
          .update(
            enable = true,
            core_state = core_state,
            decoded_nzp = decoded_nzp,
            decoded_immediate = decoded_immediate,
            decoded_nzp_write_enable = decoded_nzp_write_enable,
            decoded_pc_mux = decoded_pc_mux,
            alu_out = alu_out,
            current_pc = current_pc
          )

        dut.io.next_pc.expect(
          pcModel.next_pc.U
        )

        cnt += 1
      }
    }
  }
}


package registers

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.RegInputOp
import statecode.CoreState

import scala.collection.mutable.ArrayBuffer

class RegModel {
  var registers = ArrayBuffer.fill(16)(0)
  var rs_in     = 0
  var rt_in     = 0

  def rs_out(): Int = {
    registers(rs_in)
  }

  def rt_out(): Int = {
    registers(rt_in)
  }

    decoded_reg_write_enable: Boolean,
    decoded_reg_input_mux:    RegInputOp.Type,
    decoded_immediate:        Int,
    alu_out:                  Int,
    lsu_out:                  Int
  ): Unit = {
    if (enable) {
      registers(13) = block_id
      if (core_state == CoreState.REQUEST) {
        rs_in = rs
        rt_in = rt
      } else if (decoded_reg_write_enable && rd < 13) {
        val data = decoded_reg_input_mux match {
          case RegInputOp.ARITHMETIC => alu_out
          case RegInputOp.MEMORY     => lsu_out
          case RegInputOp.CONSTANT   => decoded_immediate
        }
        registers(rd) = data
      }
    }
  }
}

class RegFileSpec extends AnyFreeSpec with Matchers {
  "Test load/store for RegisterFiles" in {
    val threadsPerBlk = 4
    val threadId      = 0
    val dataBits      = 8
    simulate(new RegisterFiles(threadsPerBlk, threadId, dataBits)) { dut =>
      dut.reset.poke(true.B)
      dut.clock.step()
      dut.reset.poke(false.B)
      dut.clock.step()

      dut.io.enable.poke(true.B)

      // test reg14 and reg15 defualt values
      dut.io.core_state.poke(CoreState.REQUEST)
      dut.io.decoded_reg_address.rs.poke(14.U)
      dut.io.decoded_reg_address.rt.poke(15.U)
      dut.clock.step()
      dut.io.reg_out.rs.expect(threadsPerBlk.U)
      dut.io.reg_out.rt.expect(threadId.U)

      // helper functions
      def randomRegOp(rng: scala.util.Random): RegInputOp.Type = {
        val v = rng.nextInt(3)
        v match {
          case 0 => RegInputOp.ARITHMETIC
          case 1 => RegInputOp.MEMORY
          case 2 => RegInputOp.CONSTANT
        }
      }

      def randomCoreState(rng: scala.util.Random): CoreState.Type = {
        val v = rng.nextInt(2)
        v match {
          case 0 => CoreState.REQUEST
          case 1 => CoreState.UPDATE
        }
      }

      // Initialize register model
      val regModel = new RegModel()

      val times = 100
      val rng   = new scala.util.Random(42) // 42 is the seed for reproducibility
      for (i <- 0 until times) {
        // Generate test sequence
        val block_id   = rng.nextInt(256)
        val rd_in      = rng.nextInt(12)
        val rs_in      = rng.nextInt(12)
        val rt_in      = rng.nextInt(12)
        val op         = randomRegOp(rng)
        val core_state = randomCoreState(rng)
        val alu_out    = rng.nextInt(256)
        val lsu_out    = rng.nextInt(256)
        val immediate  = rng.nextInt(256)

        dut.io.core_state.poke(core_state)
        val is_write = core_state == CoreState.UPDATE
        dut.io.decoded_reg_write_enable.poke(is_write.B)

        dut.io.decoded_reg_address.rd.poke(rd_in.U)
        dut.io.decoded_reg_address.rs.poke(rs_in.U)
        dut.io.decoded_reg_address.rt.poke(rt_in.U)

        dut.io.decoded_reg_input_mux.poke(op)
        dut.io.decoded_immediate.poke(immediate.U)
        dut.io.alu_out.poke(alu_out.U)
        dut.io.lsu_out.poke(lsu_out.U)

        // Update model
        regModel.update(
          enable = true,
          block_id = block_id,
          core_state = core_state,
          rd = rd_in,
          rs = rs_in,
          rt = rt_in,
          decoded_reg_write_enable = is_write,
          decoded_reg_input_mux = op,
          decoded_immediate = immediate,
          alu_out = alu_out,
          lsu_out = lsu_out
        )

        dut.clock.step()
        // println(
        //   s"rd_in: $rd_in, rs_in: $rs_in, rt_in: $rt_in, op: $op, core_state: $core_state, immediate: $immediate, alu_out: $alu_out, lsu_out: $lsu_out"
        // )
        // println(s"regModel after 1 cycle: ${regModel.registers}")

        if (core_state == CoreState.REQUEST) {
          dut.io.reg_out.rs.expect(regModel.rs_out().U)
          dut.io.reg_out.rt.expect(regModel.rt_out().U)
        }
      }
    }
  }
}

  def update(
    enable:                   Boolean,
    block_id:                 Int,
    core_state:               CoreState.Type,


package scheduler

import chisel3._
import chisel3.util._
import chisel3.experimental.BundleLiterals._
import chisel3.simulator.EphemeralSimulator._
import org.scalatest.freespec.AnyFreeSpec
import org.scalatest.matchers.must.Matchers

import statecode.CoreState
import statecode.LSUState
import statecode.FetcherState

class SchedulerModel(ThreadsPerBlock: Int = 4) {
  var current_pc = 0
  var core_state = CoreState.IDLE
  var done       = false

  def update(
    start:            Boolean,
    mem_read_enable:  Boolean,
    mem_write_enable: Boolean,
    decoded_ret:      Boolean,
    fetcher_state:    FetcherState.Type,
    lsu_state:        Array[LSUState.Type],
    next_pc:          Array[Int]
  ): Unit = {
    // State machine logic matching the Scheduler module
    core_state match {
      case CoreState.IDLE =>
        if (start) {
          core_state = CoreState.FETCH
        }

      case CoreState.FETCH =>
        if (fetcher_state == FetcherState.FETCHED) {
          core_state = CoreState.DECODE
        }

      case CoreState.DECODE =>
        core_state = CoreState.REQUEST

      case CoreState.REQUEST =>
        core_state = CoreState.WAIT

      case CoreState.WAIT =>
        // Check if any LSU is waiting or requesting
        val any_lsu_waiting = lsu_state.exists(state => state == LSUState.WAITING || state == LSUState.REQUESTING)
        if (!any_lsu_waiting) {
          core_state = CoreState.EXECUTE
        }

      case CoreState.EXECUTE =>
        core_state = CoreState.UPDATE

      case CoreState.UPDATE =>
        if (decoded_ret) {
          done = true

      case CoreState.DONE =>
      // Stay in DONE state
    }
  }
}

class SchedulerSpec extends AnyFreeSpec with Matchers {
  "Test scheduler" - {
    "should match model behavior with random inputs" in {
      val ThreadsPerBlock: Int = 4
      simulate(new Scheduler(ThreadsPerBlock)) { dut =>
        dut.reset.poke(true.B)
        dut.clock.step()
        dut.reset.poke(false.B)
        dut.clock.step()

        var cnt            = 0
        val rng            = new scala.util.Random(42) // 42 is the seed for reproducibility
        val schedulerModel = new SchedulerModel(ThreadsPerBlock)

        def randomLsuState(rng: scala.util.Random): LSUState.Type = {
          val v = rng.nextInt(4)
          v match {
            case 0 => LSUState.IDLE
            case 1 => LSUState.REQUESTING
            case 2 => LSUState.WAITING
            case 3 => LSUState.DONE
          }
        }

        while (cnt < 10000) {
          // Generate random inputs
          val start              = rng.nextBoolean()
          val mem_read_enable    = rng.nextBoolean()
          val mem_write_enable   = rng.nextBoolean()
          val decoded_ret        = rng.nextBoolean()
          val (fetcher_state, _) = FetcherState.safe(rng.nextInt(3).U)              // Assuming 3 states: IDLE, FETCHING, FETCHED
          val lsu_state          = Array.fill(ThreadsPerBlock)(randomLsuState(rng)) // Assuming 4 states in LSUState
          val next_pc            = Array.fill(ThreadsPerBlock)(rng.nextInt(256))    // 8-bit PC values

          // Poke inputs
          dut.io.start.poke(start.B)
          dut.io.mem_rw_enable.read_enable.poke(mem_read_enable.B)
          dut.io.mem_rw_enable.write_enable.poke(mem_write_enable.B)
          dut.io.decoded_ret.poke(decoded_ret.B)
          dut.io.fetcher_state.poke(fetcher_state)
          for (i <- 0 until ThreadsPerBlock) {
            dut.io.lsu_state(i).poke(lsu_state(i))
            dut.io.next_pc(i).poke(next_pc(i).U)
          }

          dut.clock.step()

          // Update model
          schedulerModel.update(
            start = start,
            mem_read_enable = mem_read_enable,
            mem_write_enable = mem_write_enable,
            decoded_ret = decoded_ret,
            fetcher_state = fetcher_state,
            lsu_state = lsu_state,
            next_pc = next_pc
          )

          // Compare model with DUT
          dut.io.current_pc.expect(schedulerModel.current_pc.U)
          assert(dut.io.core_state.peekValue().asBigInt == schedulerModel.core_state.litValue)
          dut.io.done.expect(schedulerModel.done.B)

          cnt += 1
        }
      }
    }
  }
}

          core_state = CoreState.DONE
        } else {
          // Use the last thread's next_pc (same as in module)
          current_pc = next_pc(ThreadsPerBlock - 1)