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.tester._
import chisel3.tester.RawTester.test

class RegisterModule extends Module {
  val io = IO(new Bundle {
    val in  = Input(UInt(12.W))
    val out = Output(UInt(12.W))
  })
  
  val register = Reg(UInt(12.W))
  register := io.in + 1.U
  io.out := register
}

test(new RegisterModule) { c =>
  for (i <- 0 until 100) {
    c.io.in.poke(i.U)
    c.clock.step(1)
    c.io.out.expect((i + 1).U)
  }
}
println("SUCCESS!!")

println(getVerilog(new RegisterModule))

class RegNextModule extends Module {
  val io = IO(new Bundle {
    val in  = Input(UInt(12.W))
    val out = Output(UInt(12.W))
  })
  
  // register bitwidth is inferred from io.out
  io.out := RegNext(io.in + 1.U)
}

test(new RegNextModule) { c =>
  for (i <- 0 until 100) {
    c.io.in.poke(i.U)
    c.clock.step(1)
    c.io.out.expect((i + 1).U)
  }
}
println("SUCCESS!!")

println(getVerilog(new RegNextModule))

class RegInitModule extends Module {
  val io = IO(new Bundle {
    val in  = Input(UInt(12.W))
    val out = Output(UInt(12.W))
  })
  
  val register = RegInit(0.U(12.W))
  register := io.in + 1.U
  io.out := register
}

println(getVerilog(new RegInitModule))

class FindMax extends Module {
  val io = IO(new Bundle {
    val in  = Input(UInt(10.W))
    val max = Output(UInt(10.W))
  })

  val max = RegInit(0.U(10.W))
  when (io.in > max) {
    max := io.in
  }
  io.max := max
}

test(new FindMax) { c =>
    c.io.max.expect(0.U)
    c.io.in.poke(1.U)
    c.clock.step(1)
    c.io.max.expect(1.U)
    c.io.in.poke(3.U)
    c.clock.step(1)
    c.io.max.expect(3.U)
    c.io.in.poke(2.U)
    c.clock.step(1)
    c.io.max.expect(3.U)
    c.io.in.poke(24.U)
    c.clock.step(1)
    c.io.max.expect(24.U)
}
println("SUCCESS!!")

class Comb extends Module {
  val io = IO(new Bundle {
    val in  = Input(SInt(12.W))
    val out = Output(SInt(12.W))
  })

  val delay: SInt = Reg(SInt(12.W))
  delay := io.in
  io.out := io.in - delay
}
println(getVerilog(new Comb))

class MyShiftRegister(val init: Int = 1) extends Module {
  val io = IO(new Bundle {
    val in  = Input(Bool())
    val out = Output(UInt(4.W))
  })

  val state = RegInit(UInt(4.W), init.U)

  ???
}

test(new MyShiftRegister()) { c =>
  var state = c.init
  for (i <- 0 until 10) {
    // poke in LSB of i (i % 2)
    c.io.in.poke(((i % 2) != 0).B)
    // update expected state
    state = ((state * 2) + (i % 2)) & 0xf
    c.clock.step(1)
    c.io.out.expect(state.U)
  }
}
println("SUCCESS!!")

// n is the output width (number of delays - 1)
// init state to init
class MyOptionalShiftRegister(val n: Int, val init: BigInt = 1) extends Module {
  val io = IO(new Bundle {
    val en  = Input(Bool())
    val in  = Input(Bool())
    val out = Output(UInt(n.W))
  })

  val state = RegInit(init.U(n.W))

  ???
}

// test different depths
for (i <- Seq(3, 4, 8, 24, 65)) {
  println(s"Testing n=$i")
  test(new MyOptionalShiftRegister(n = i)) { c =>
    val inSeq = Seq(0, 1, 1, 1, 0, 1, 1, 0, 0, 1)
    var state = c.init
    var i = 0
    c.io.en.poke(true.B)
    while (i < 10 * c.n) {
      // poke in repeated inSeq
      val toPoke = inSeq(i % inSeq.length)
      c.io.in.poke((toPoke != 0).B)
      // update expected state
      state = ((state * 2) + toPoke) & BigInt("1"*c.n, 2)
      c.clock.step(1)
      c.io.out.expect(state.U)

      i += 1
    }
  }
}
println("SUCCESS!!")

// we need to import multi-clock features
import chisel3.experimental.{withClock, withReset, withClockAndReset}

class ClockExamples extends Module {
  val io = IO(new Bundle {
    val in = Input(UInt(10.W))
    val alternateReset    = Input(Bool())
    val alternateClock    = Input(Clock())
    val outImplicit       = Output(UInt())
    val outAlternateReset = Output(UInt())
    val outAlternateClock = Output(UInt())
    val outAlternateBoth  = Output(UInt())
  })

  val imp = RegInit(0.U(10.W))
  imp := io.in
  io.outImplicit := imp

  withReset(io.alternateReset) {
    // everything in this scope with have alternateReset as the reset
    val altRst = RegInit(0.U(10.W))
    altRst := io.in
    io.outAlternateReset := altRst
  }

  withClock(io.alternateClock) {
    val altClk = RegInit(0.U(10.W))
    altClk := io.in
    io.outAlternateClock := altClk
  }

  withClockAndReset(io.alternateClock, io.alternateReset) {
    val alt = RegInit(0.U(10.W))
    alt := io.in
    io.outAlternateBoth := alt
  }
}

println(getVerilog(new ClockExamples))