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.iotesters.{ChiselFlatSpec, Driver, PeekPokeTester}

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
}

class RegisterModuleTester(c: RegisterModule) extends PeekPokeTester(c) {
  for (i <- 0 until 100) {
    poke(c.io.in, i)
    step(1)
    expect(c.io.out, i+1)
  }
}
assert(chisel3.iotesters.Driver(() => new RegisterModule) { c => new RegisterModuleTester(c) })
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)
}

class RegNextModuleTester(c: RegNextModule) extends PeekPokeTester(c) {
  for (i <- 0 until 100) {
    poke(c.io.in, i)
    step(1)
    expect(c.io.out, i+1)
  }
}
assert(chisel3.iotesters.Driver(() => new RegNextModule) { c => new RegNextModuleTester(c) })
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
}
assert(chisel3.iotesters.Driver(() => new FindMax) {
  c => new PeekPokeTester(c) {
    expect(c.io.max, 0)
    poke(c.io.in, 1)
    step(1)
    expect(c.io.max, 1)
    poke(c.io.in, 3)
    step(1)
    expect(c.io.max, 3)
    poke(c.io.in, 2)
    step(1)
    expect(c.io.max, 3)
    poke(c.io.in, 24)
    step(1)
    expect(c.io.max, 24)
  }
})
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)

  ???
}

class MyShiftRegisterTester(c: MyShiftRegister) extends PeekPokeTester(c) {
  var state = c.init
  for (i <- 0 until 10) {
    // poke in LSB of i (i % 2)
    poke(c.io.in, i % 2)
    // update expected state
    state = ((state * 2) + (i % 2)) & 0xf
    step(1)
    expect(c.io.out, state)
  }
}
assert(chisel3.iotesters.Driver(() => new MyShiftRegister()) {
    c => new MyShiftRegisterTester(c)
})
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))

  ???
}

class MyOptionalShiftRegisterTester(c: MyOptionalShiftRegister) extends PeekPokeTester(c) {
  val inSeq = Seq(0, 1, 1, 1, 0, 1, 1, 0, 0, 1)
  var state = c.init
  var i = 0
  poke(c.io.en, 1)
  while (i < 10 * c.n) {
    // poke in repeated inSeq
    val toPoke = inSeq(i % inSeq.length)
    poke(c.io.in, toPoke)
    // update expected state
    state = ((state * 2) + toPoke) & BigInt("1"*c.n, 2)
    step(1)
    expect(c.io.out, state)

    i += 1
  }
}

// test different depths
for (i <- Seq(3, 4, 8, 24, 65)) {
  println(s"Testing n=$i")
  assert(chisel3.iotesters.Driver(() => new MyOptionalShiftRegister(n = i)) {
    c => new MyOptionalShiftRegisterTester(c)
  })
}
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))