val path = System.getProperty("user.dir") + "/source/load-ivy.sc"
interp.load.module(ammonite.ops.Path(java.nio.file.FileSystems.getDefault().getPath(path)))

import chisel3._
import chisel3.util._
import chisel3.experimental._
import chisel3.experimental.BundleLiterals._
import chisel3.tester._
import chisel3.tester.RawTester.test

// Chisel Code, but pass in a parameter to set widths of ports
class PassthroughGenerator(width: Int) extends Module { 
  val io = IO(new Bundle {
    val in = Input(UInt(width.W))
    val out = Output(UInt(width.W))
  })
  io.out := io.in
}

test(new PassthroughGenerator(16)) { c =>
    c.io.in.poke(0.U)     // Set our input to value 0
    c.io.out.expect(0.U)  // Assert that the output correctly has 0
    c.io.in.poke(1.U)     // Set our input to value 1
    c.io.out.expect(1.U)  // Assert that the output correctly has 1
    c.io.in.poke(2.U)     // Set our input to value 2
    c.io.out.expect(2.U)  // Assert that the output correctly has 2
}

test(new PassthroughGenerator(16)) { c =>
    c.io.in.poke(0.U)     // Set our input to value 0
    c.clock.step(1)    // advance the clock
    c.io.out.expect(0.U)  // Assert that the output correctly has 0
    c.io.in.poke(1.U)     // Set our input to value 1
    c.clock.step(1)    // advance the clock
    c.io.out.expect(1.U)  // Assert that the output correctly has 1
    c.io.in.poke(2.U)     // Set our input to value 2
    c.clock.step(1)    // advance the clock
    c.io.out.expect(2.U)  // Assert that the output correctly has 2
}

case class QueueModule[T <: Data](ioType: T, entries: Int) extends MultiIOModule {
  val in = IO(Flipped(Decoupled(ioType)))
  val out = IO(Decoupled(ioType))
  out <> Queue(in, entries)
}


test(QueueModule(UInt(9.W), entries = 200)) { c =>
    // Example testsequence showing the use and behavior of Queue
    c.in.initSource()
    c.in.setSourceClock(c.clock)
    c.out.initSink()
    c.out.setSinkClock(c.clock)
    
    val testVector = Seq.tabulate(200){ i => i.U }

    testVector.zip(testVector).foreach { case (in, out) =>
      c.in.enqueueNow(in)
      c.out.expectDequeueNow(out)
    }
}

test(QueueModule(UInt(9.W), entries = 200)) { c =>
    // Example testsequence showing the use and behavior of Queue
    c.in.initSource()
    c.in.setSourceClock(c.clock)
    c.out.initSink()
    c.out.setSinkClock(c.clock)
    
    val testVector = Seq.tabulate(100){ i => i.U }

    c.in.enqueueSeq(testVector)
    c.out.expectDequeueSeq(testVector)
}


test(QueueModule(UInt(9.W), entries = 200)) { c =>
    // Example testsequence showing the use and behavior of Queue
    c.in.initSource()
    c.in.setSourceClock(c.clock)
    c.out.initSink()
    c.out.setSinkClock(c.clock)
    
    val testVector = Seq.tabulate(300){ i => i.U }

    fork {
        c.in.enqueueSeq(testVector)
    }.fork {
        c.out.expectDequeueSeq(testVector)
    }.join()
}

class GcdInputBundle(val w: Int) extends Bundle {
  val value1 = UInt(w.W)
  val value2 = UInt(w.W)
}

class GcdOutputBundle(val w: Int) extends Bundle {
  val value1 = UInt(w.W)
  val value2 = UInt(w.W)
  val gcd    = UInt(w.W)
}

/**
  * Compute GCD using subtraction method.
  * Subtracts the smaller of registers x and y from the larger until register y is zero.
  * value input register x is then the Gcd
  * returns a packet of information with the two input values and their GCD
  */
class DecoupledGcd(width: Int) extends MultiIOModule {

  val input = IO(Flipped(Decoupled(new GcdInputBundle(width))))
  val output = IO(Decoupled(new GcdOutputBundle(width)))

  val xInitial    = Reg(UInt())
  val yInitial    = Reg(UInt())
  val x           = Reg(UInt())
  val y           = Reg(UInt())
  val busy        = RegInit(false.B)
  val resultValid = RegInit(false.B)

  input.ready := ! busy
  output.valid := resultValid
  output.bits := DontCare

  when(busy)  {
    // during computation keep subtracting the smaller from the larger
    when(x > y) {
      x := x - y
    }.otherwise {
      y := y - x
    }
    when(y === 0.U) {
      // when y becomes zero computation is over, signal valid data to output
      output.bits.gcd := x
      output.bits.value1 := xInitial
      output.bits.value2 := yInitial
      output.bits.gcd := x
      output.valid := true.B
      busy := false.B
    }
  }.otherwise {
    when(input.valid) {
      // valid data available and no computation in progress, grab new values and start
      val bundle = input.deq()
      x := bundle.value1
      y := bundle.value2
      xInitial := bundle.value1
      yInitial := bundle.value2
      busy := true.B
      resultValid := false.B
    }
  }
}


test(new DecoupledGcd(16)) { dut =>
  dut.input.initSource().setSourceClock(dut.clock)
  dut.output.initSink().setSinkClock(dut.clock)

  val testValues = for { x <- 1 to 10; y <- 1 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)
  }

  fork {
    dut.input.enqueueSeq(inputSeq)
  }.fork {
    dut.output.expectDequeueSeq(resultSeq)
  }.join()
}