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 ParameterizedScalaObject(param1: Int, param2: String) {
  println(s"I have parameters: param1 = $param1 and param2 = $param2")
}
val obj1 = new ParameterizedScalaObject(4,     "Hello")
val obj2 = new ParameterizedScalaObject(4 + 2, "World")

class ParameterizedWidthAdder(in0Width: Int, in1Width: Int, sumWidth: Int) extends Module {
  require(in0Width >= 0)
  require(in1Width >= 0)
  require(sumWidth >= 0)
  val io = IO(new Bundle {
    val in0 = Input(UInt(in0Width.W))
    val in1 = Input(UInt(in1Width.W))
    val sum = Output(UInt(sumWidth.W))
  })
  // a +& b includes the carry, a + b does not
  io.sum := io.in0 +& io.in1
}

println(getVerilog(new ParameterizedWidthAdder(1, 4, 6)))

/** Sort4 sorts its 4 inputs to its 4 outputs */
class Sort4(ascending: Boolean) extends Module {
  val io = IO(new Bundle {
    val in0 = Input(UInt(16.W))
    val in1 = Input(UInt(16.W))
    val in2 = Input(UInt(16.W))
    val in3 = Input(UInt(16.W))
    val out0 = Output(UInt(16.W))
    val out1 = Output(UInt(16.W))
    val out2 = Output(UInt(16.W))
    val out3 = Output(UInt(16.W))
  })
    
  // this comparison funtion decides < or > based on the module's parameterization
  def comp(l: UInt, r: UInt): Bool = {
      if (ascending) {
        l < r
      } else {
        l > r
    }
  }

  val row10 = Wire(UInt(16.W))
  val row11 = Wire(UInt(16.W))
  val row12 = Wire(UInt(16.W))
  val row13 = Wire(UInt(16.W))

  when(comp(io.in0, io.in1)) {
    row10 := io.in0            // preserve first two elements
    row11 := io.in1
  }.otherwise {
    row10 := io.in1            // swap first two elements
    row11 := io.in0
  }

  when(comp(io.in2, io.in3)) {
    row12 := io.in2            // preserve last two elements
    row13 := io.in3
  }.otherwise {
    row12 := io.in3            // swap last two elements
    row13 := io.in2
  }

  val row21 = Wire(UInt(16.W))
  val row22 = Wire(UInt(16.W))

  when(comp(row11, row12)) {
    row21 := row11            // preserve middle 2 elements
    row22 := row12
  }.otherwise {
    row21 := row12            // swap middle two elements
    row22 := row11
  }

  val row31 = Wire(UInt(16.W))
  val row32 = Wire(UInt(16.W))
  when(comp(row10, row13)) {
    row31 := row10            // preserve middle 2 elements
    row32 := row13
  }.otherwise {
    row31 := row13            // swap middle two elements
    row32 := row10
  }

  when(comp(row10, row21)) {
    io.out0 := row31            // preserve first two elements
    io.out1 := row21
  }.otherwise {
    io.out0 := row21            // swap first two elements
    io.out1 := row31
  }

  when(comp(row22, row13)) {
    io.out2 := row22            // preserve first two elements
    io.out3 := row32
  }.otherwise {
    io.out2 := row32            // swap first two elements
    io.out3 := row22
  }
}

// verify the inputs are sorted
class Sort4AscendingTester(c: Sort4) extends PeekPokeTester(c) {
  poke(c.io.in0, 3)
  poke(c.io.in1, 6)
  poke(c.io.in2, 9)
  poke(c.io.in3, 12)
  expect(c.io.out0, 3)
  expect(c.io.out1, 6)
  expect(c.io.out2, 9)
  expect(c.io.out3, 12)

  poke(c.io.in0, 13)
  poke(c.io.in1, 4)
  poke(c.io.in2, 6)
  poke(c.io.in3, 1)
  expect(c.io.out0, 1)
  expect(c.io.out1, 4)
  expect(c.io.out2, 6)
  expect(c.io.out3, 13)

  poke(c.io.in0, 13)
  poke(c.io.in1, 6)
  poke(c.io.in2, 4)
  poke(c.io.in3, 1)
  expect(c.io.out0, 1)
  expect(c.io.out1, 4)
  expect(c.io.out2, 6)
  expect(c.io.out3, 13)

}
class Sort4DescendingTester(c: Sort4) extends PeekPokeTester(c) {
  poke(c.io.in0, 3)
  poke(c.io.in1, 6)
  poke(c.io.in2, 9)
  poke(c.io.in3, 12)
  expect(c.io.out0, 12)
  expect(c.io.out1, 9)
  expect(c.io.out2, 6)
  expect(c.io.out3, 3)

  poke(c.io.in0, 13)
  poke(c.io.in1, 4)
  poke(c.io.in2, 6)
  poke(c.io.in3, 1)
  expect(c.io.out0, 13)
  expect(c.io.out1, 6)
  expect(c.io.out2, 4)
  expect(c.io.out3, 1)
    
  poke(c.io.in0, 1)
  poke(c.io.in1, 6)
  poke(c.io.in2, 4)
  poke(c.io.in3, 13)
  expect(c.io.out0, 13)
  expect(c.io.out1, 6)
  expect(c.io.out2, 4)
  expect(c.io.out3, 1)

}

// Here are the testers
val worksAscending = iotesters.Driver(() => new Sort4(true)) { c => new Sort4AscendingTester(c) }
val worksDescending = iotesters.Driver(() => new Sort4(false)) { c => new Sort4DescendingTester(c) }
assert(worksAscending && worksDescending) // Scala Code: if works == false, will throw an error
println("SUCCESS!!") // Scala Code: if we get here, our tests passed!

val map = Map("a" -> 1)
val a = map("a")
println(a)
val b = map("b")
println(b)

val map = Map("a" -> 1)
val a = map.get("a")
println(a)
val b = map.get("b")
println(b)

val some = Some(1)
val none = None
println(some.get)          // Returns 1
// println(none.get)       // Errors!
println(some.getOrElse(2)) // Returns 1
println(none.getOrElse(2)) // Returns 2

class DelayBy1(resetValue: Option[UInt] = None) extends Module {
    val io = IO(new Bundle {
        val in  = Input( UInt(16.W))
        val out = Output(UInt(16.W))
    })
    val reg = if (resetValue.isDefined) { // resetValue = Some(number)
        RegInit(resetValue.get)
    } else { //resetValue = None
        Reg(UInt())
    }
    reg := io.in
    io.out := reg
}

println(getVerilog(new DelayBy1))
println(getVerilog(new DelayBy1(Some(3.U))))

// y is an integer variable defined somewhere else in the code
val y = 7
/// ...
val x = y match {
  case 0 => "zero" // One common syntax, preferred if fits in one line
  case 1 =>        // Another common syntax, preferred if does not fit in one line.
      "one"        // Note the code block continues until the next case
  case 2 => {      // Another syntax, but curly braces are not required
      "two"
  }
  case _ => "many" // _ is a wildcard that matches all values
}
println("y is " + x)

def animalType(biggerThanBreadBox: Boolean, meanAsCanBe: Boolean): String = {
  (biggerThanBreadBox, meanAsCanBe) match {
    case (true, true) => "wolverine"
    case (true, false) => "elephant"
    case (false, true) => "shrew"
    case (false, false) => "puppy"
  }
}
println(animalType(true, true))

val sequence = Seq("a", 1, 0.0)
sequence.foreach { x =>
  x match {
    case s: String => println(s"$x is a String")
    case s: Int    => println(s"$x is an Int")
    case s: Double => println(s"$x is a Double")
    case _ => println(s"$x is an unknown type!")
  }
}

val sequence = Seq("a", 1, 0.0)
sequence.foreach { x =>
  x match {
    case _: Int | _: Double => println(s"$x is a number!")
    case _ => println(s"$x is an unknown type!")
  }
}

val sequence = Seq(Seq("a"), Seq(1), Seq(0.0))
sequence.foreach { x =>
  x match {
    case s: Seq[String] => println(s"$x is a String")
    case s: Seq[Int]    => println(s"$x is an Int")
    case s: Seq[Double] => println(s"$x is a Double")
  }
}

class DelayBy1(resetValue: Option[UInt] = None) extends Module {
  val io = IO(new Bundle {
    val in  = Input( UInt(16.W))
    val out = Output(UInt(16.W))
  })
  val reg = resetValue match {
    case Some(r) => RegInit(r)
    case None    => Reg(UInt())
  }
  reg := io.in
  io.out := reg
}

println(getVerilog(new DelayBy1))
println(getVerilog(new DelayBy1(Some(3.U))))

class HalfFullAdder(val hasCarry: Boolean) extends Module {
  val io = IO(new Bundle {
    val a = Input(UInt(1.W))
    val b = Input(UInt(1.W))
    val carryIn = if (hasCarry) Some(Input(UInt(1.W))) else None
    val s = Output(UInt(1.W))
    val carryOut = Output(UInt(1.W))
  })
  val sum = io.a +& io.b +& io.carryIn.getOrElse(0.U)
  io.s := sum(0)
  io.carryOut := sum(1)
}

class HalfAdderTester(c: HalfFullAdder) extends PeekPokeTester(c) {
  require(!c.hasCarry, "DUT must be half adder")
  // 0 + 0 = 0
  poke(c.io.a, 0)
  poke(c.io.b, 0)
  expect(c.io.s, 0)
  expect(c.io.carryOut, 0)
  // 0 + 1 = 1
  poke(c.io.b, 1)
  expect(c.io.s, 1)
  expect(c.io.carryOut, 0)
  // 1 + 1 = 2
  poke(c.io.a, 1)
  expect(c.io.s, 0)
  expect(c.io.carryOut, 1)
  // 1 + 0 = 1
  poke(c.io.b, 0)
  expect(c.io.s, 1)
  expect(c.io.carryOut, 0)
}

class FullAdderTester(c: HalfFullAdder) extends PeekPokeTester(c) {
  require(c.hasCarry, "DUT must be half adder")
  poke(c.io.carryIn.get, 0)
  // 0 + 0 + 0 = 0
  poke(c.io.a, 0)
  poke(c.io.b, 0)
  expect(c.io.s, 0)
  expect(c.io.carryOut, 0)
  // 0 + 0 + 1 = 1
  poke(c.io.b, 1)
  expect(c.io.s, 1)
  expect(c.io.carryOut, 0)
  // 0 + 1 + 1 = 2
  poke(c.io.a, 1)
  expect(c.io.s, 0)
  expect(c.io.carryOut, 1)
  // 0 + 1 + 0 = 1
  poke(c.io.b, 0)
  expect(c.io.s, 1)
  expect(c.io.carryOut, 0)

  poke(c.io.carryIn.get, 1)
  // 1 + 0 + 0 = 1
  poke(c.io.a, 0)
  poke(c.io.b, 0)
  expect(c.io.s, 1)
  expect(c.io.carryOut, 0)
  // 1 + 0 + 1 = 2
  poke(c.io.b, 1)
  expect(c.io.s, 0)
  expect(c.io.carryOut, 1)
  // 1 + 1 + 1 = 3
  poke(c.io.a, 1)
  expect(c.io.s, 1)
  expect(c.io.carryOut, 1)
  // 1 + 1 + 0 = 2
  poke(c.io.b, 0)
  expect(c.io.s, 0)
  expect(c.io.carryOut, 1)
}

val worksHalf = iotesters.Driver(() => new HalfFullAdder(false)) { c => new HalfAdderTester(c) }
val worksFull = iotesters.Driver(() => new HalfFullAdder(true)) { c => new FullAdderTester(c) }
assert(worksHalf && worksFull) // Scala Code: if works == false, will throw an error
println("SUCCESS!!") // Scala Code: if we get here, our tests passed!

class HalfFullAdder(val hasCarry: Boolean) extends Module {
  val io = IO(new Bundle {
    val a = Input(UInt(1.W))
    val b = Input(UInt(1.W))
    val carryIn = Input(if (hasCarry) UInt(1.W) else UInt(0.W))
    val s = Output(UInt(1.W))
    val carryOut = Output(UInt(1.W))
  })
  val sum = io.a +& io.b +& io.carryIn
  io.s := sum(0)
  io.carryOut := sum(1)
}
println("Half Adder:")
println(getVerilog(new HalfFullAdder(false)))
println("\n\nFull Adder:")
println(getVerilog(new HalfFullAdder(true)))

object CatDog {
  implicit val numberOfCats: Int = 3
  //implicit val numberOfDogs: Int = 5

  def tooManyCats(nDogs: Int)(implicit nCats: Int): Boolean = nCats > nDogs
    
  val imp = tooManyCats(2)    // Argument passed implicitly!
  val exp = tooManyCats(2)(1) // Argument passed explicitly!
}
CatDog.imp
CatDog.exp

sealed trait Verbosity
implicit case object Silent extends Verbosity
case object Verbose extends Verbosity

class ParameterizedWidthAdder(in0Width: Int, in1Width: Int, sumWidth: Int)(implicit verbosity: Verbosity)
extends Module {
  def log(msg: => String): Unit = verbosity match {
    case Silent =>
    case Verbose => println(msg)
  }
  require(in0Width >= 0)
  log(s"in0Width of $in0Width OK")
  require(in1Width >= 0)
  log(s"in1Width of $in1Width OK")
  require(sumWidth >= 0)
  log(s"sumWidth of $sumWidth OK")
  val io = IO(new Bundle {
    val in0 = Input(UInt(in0Width.W))
    val in1 = Input(UInt(in1Width.W))
    val sum = Output(UInt(sumWidth.W))
  })
  log("Made IO")
  io.sum := io.in0 + io.in1
  log("Assigned output")
}

println(getVerilog(new ParameterizedWidthAdder(1, 4, 5)))
println(getVerilog(new ParameterizedWidthAdder(1, 4, 5)(Verbose)))

class Animal(val name: String, val species: String)
class Human(val name: String)
implicit def human2animal(h: Human): Animal = new Animal(h.name, "Homo sapiens")
val me = new Human("Adam")
println(me.species)

// Mealy machine has
case class BinaryMealyParams(
  // number of states
  nStates: Int,
  // initial state
  s0: Int,
  // function describing state transition
  stateTransition: (Int, Boolean) => Int,
  // function describing output
  output: (Int, Boolean) => Int
) {
  require(nStates >= 0)
  require(s0 < nStates && s0 >= 0)
}

class BinaryMealy(val mp: BinaryMealyParams) extends Module {
  val io = IO(new Bundle {
    val in = Input(Bool())
    val out = Output(UInt())
  })

  val state = RegInit(UInt(), mp.s0.U)

  // output zero if no states
  io.out := 0.U
  for (i <- 0 until mp.nStates) {
    when (state === i.U) {
      when (io.in) {
        state  := mp.stateTransition(i, true).U
        io.out := mp.output(i, true).U
      }.otherwise {
        state  := mp.stateTransition(i, false).U
        io.out := mp.output(i, false).U
      }
    }
  }
}

// example from https://en.wikipedia.org/wiki/Mealy_machine
val nStates = 3
val s0 = 2
def stateTransition(state: Int, in: Boolean): Int = {
  if (in) {
    1
  } else {
    0
  }
}
def output(state: Int, in: Boolean): Int = {
  if (state == 2) {
    return 0
  }
  if ((state == 1 && !in) || (state == 0 && in)) {
    return 1
  } else {
    return 0
  }
}

val testParams = BinaryMealyParams(nStates, s0, stateTransition, output)

class BinaryMealyTester(c: BinaryMealy) extends PeekPokeTester(c) {
  poke(c.io.in, false)
  expect(c.io.out, 0)
  step(1)
  poke(c.io.in, false)
  expect(c.io.out, 0)
  step(1)
  poke(c.io.in, false)
  expect(c.io.out, 0)
  step(1)
  poke(c.io.in, true)
  expect(c.io.out, 1)
  step(1)
  poke(c.io.in, true)
  expect(c.io.out, 0)
  step(1)
  poke(c.io.in, false)
  expect(c.io.out, 1)
  step(1)
  poke(c.io.in, true)
  expect(c.io.out, 1)
  step(1)
  poke(c.io.in, false)
  expect(c.io.out, 1)
  step(1)
  poke(c.io.in, true)
  expect(c.io.out, 1)
}
val works = iotesters.Driver(() => new BinaryMealy(testParams)) { c => new BinaryMealyTester(c) }
assert(works) // Scala Code: if works == false, will throw an error
println("SUCCESS!!") // Scala Code: if we get here, our tests passed!