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}
import scala.collection._

class My4ElementFir(b0: Int, b1: Int, b2: Int, b3: Int) extends Module {
  val io = IO(new Bundle {
    val in = Input(UInt(8.W))
    val out = Output(UInt(8.W))
  })

  val x_n1 = RegNext(io.in, 0.U)
  val x_n2 = RegNext(x_n1, 0.U)
  val x_n3 = RegNext(x_n2, 0.U)
  io.out := io.in * b0.U(8.W) + x_n1 * b1.U(8.W) +
    x_n2 * b2.U(8.W) + x_n3 * b3.U(8.W)
}


/**
  * A naive implementation of an FIR filter with an arbitrary number of taps.
  */
class ScalaFirFilter(taps: Seq[Int]) {
  var pseudoRegisters = List.fill(taps.length)(0)

  def poke(value: Int): Int = {
    pseudoRegisters = value :: pseudoRegisters.take(taps.length - 1)
    var accumulator = 0
    for(i <- taps.indices) {
      accumulator += taps(i) * pseudoRegisters(i)
    }
    accumulator
  }
}

val filter = new ScalaFirFilter(Seq(1, 1, 1, 1))

var out = 0

out = filter.poke(1)
println(s"out = $out")
assert(out == 1)  // 1, 0, 0, 0

out = filter.poke(4)
assert(out == 5)  // 4, 1, 0, 0
println(s"out = $out")

out = filter.poke(3)
assert(out == 8)  // 3, 4, 1, 0
println(s"out = $out")

out = filter.poke(2)
assert(out == 10)  // 2, 3, 4, 1
println(s"out = $out")

out = filter.poke(7)
assert(out == 16)  // 7, 2, 3, 4
println(s"out = $out")

out = filter.poke(0)
assert(out == 12)  // 0, 7, 2, 3
println(s"out = $out")

val goldenModel = new ScalaFirFilter(Seq(1, 1, 1, 1))

Driver(() => new My4ElementFir(1, 1, 1, 1)) {
  c => new PeekPokeTester(c) {
    for(i <- 0 until 100) {
      val input = scala.util.Random.nextInt(8)

      val goldenModelResult = goldenModel.poke(input)

      poke(c.io.in, input)

      expect(c.io.out, goldenModelResult, s"i $i, input $input, gm $goldenModelResult, ${peek(c.io.out)}")

      step(1)
    }
  }
}

class MyManyElementFir(consts: Seq[Int], bitWidth: Int) extends Module {
  val io = IO(new Bundle {
    val in = Input(UInt(bitWidth.W))
    val out = Output(UInt(bitWidth.W))
  })

  val regs = mutable.ArrayBuffer[UInt]()
  for(i <- 0 until consts.length) {
      if(i == 0) regs += io.in
      else       regs += RegNext(regs(i - 1), 0.U)
  }
  
  val muls = mutable.ArrayBuffer[UInt]()
  for(i <- 0 until consts.length) {
      muls += regs(i) * consts(i).U
  }

  val scan = mutable.ArrayBuffer[UInt]()
  for(i <- 0 until consts.length) {
      if(i == 0) scan += muls(i)
      else scan += muls(i) + scan(i - 1)
  }

  io.out := scan.last
}

val goldenModel = new ScalaFirFilter(Seq(1, 1, 1, 1))

Driver(() => new MyManyElementFir(Seq(1, 1, 1, 1), 8)) {
  c => new PeekPokeTester(c) {
    for(i <- 0 until 100) {
      val input = scala.util.Random.nextInt(8)

      val goldenModelResult = goldenModel.poke(input)

      poke(c.io.in, input)

      expect(c.io.out, goldenModelResult, s"i $i, input $input, gm $goldenModelResult, ${peek(c.io.out)}")

      step(1)
    }
  }
}

/** a convenience method to get a random integer
  */
def r(): Int = {
  scala.util.Random.nextInt(1024)
}

/**
  * run a test comparing software and hardware filters
  * run for at least twice as many samples as taps
  */
def runOneTest(taps: Seq[Int]) {
  val goldenModel = new ScalaFirFilter(taps)

  Driver(() => new MyManyElementFir(taps, 32)) {
    c => new PeekPokeTester(c) {
      for(i <- 0 until 2 * taps.length) {
        val input = r()

        val goldenModelResult = goldenModel.poke(input)

        poke(c.io.in, input)

        expect(c.io.out, goldenModelResult, s"i $i, input $input, gm $goldenModelResult, ${peek(c.io.out)}")

        step(1)
      }
    }
  }
}

for(tapSize <- 2 until 100 by 10) {
  val taps = Seq.fill(tapSize)(r())  // create a sequence of random coefficients

  runOneTest(taps)
}

runOneTest(Seq.fill(500)(r()))

val taps = Seq.fill(500)(r())

val goldenModel = new ScalaFirFilter(taps)

Driver(() => new MyManyElementFir(taps, 32)) {
  c => new PeekPokeTester(c) {
    for(i <- 0 until 100) {
      val input = r()

      val goldenModelResult = goldenModel.poke(input)

      poke(c.io.in, input)

      expect(c.io.out, goldenModelResult, s"i $i, input $input, gm $goldenModelResult, ${peek(c.io.out)}")

      step(1)
    }
  }
}

class MyManyDynamicElementVecFir(length: Int) extends Module {
  val io = IO(new Bundle {
    val in = Input(UInt(8.W))
    val out = Output(UInt(8.W))
    val consts = Input(Vec(length, UInt(8.W)))
  })

  // Reference solution
  val regs = RegInit(VecInit(Seq.fill(length - 1)(0.U(8.W))))
  for(i <- 0 until length - 1) {
      if(i == 0) regs(i) := io.in
      else       regs(i) := regs(i - 1)
  }
  
  val muls = Wire(Vec(length, UInt(8.W)))
  for(i <- 0 until length) {
      if(i == 0) muls(i) := io.in * io.consts(i)
      else       muls(i) := regs(i - 1) * io.consts(i)
  }

  val scan = Wire(Vec(length, UInt(8.W)))
  for(i <- 0 until length) {
      if(i == 0) scan(i) := muls(i)
      else scan(i) := muls(i) + scan(i - 1)
  }

  io.out := scan(length - 1)
}

val goldenModel = new ScalaFirFilter(Seq(1, 1, 1, 1))

Driver(() => new MyManyDynamicElementVecFir(4)) {
  c => new PeekPokeTester(c) {
    poke(c.io.consts(0), 1)
    poke(c.io.consts(1), 1)
    poke(c.io.consts(2), 1)
    poke(c.io.consts(3), 1)
    for(i <- 0 until 100) {
      val input = scala.util.Random.nextInt(8)

      val goldenModelResult = goldenModel.poke(input)

      poke(c.io.in, input)

      expect(c.io.out, goldenModelResult, s"i $i, input $input, gm $goldenModelResult, ${peek(c.io.out)}")

      step(1)
    }
  }
}

class RegisterFile(readPorts: Int) extends Module {
    require(readPorts >= 0)
    val io = IO(new Bundle {
        val wen   = Input(Bool())
        val waddr = Input(UInt(5.W))
        val wdata = Input(UInt(32.W))
        val raddr = Input(Vec(readPorts, UInt(5.W)))
        val rdata = Output(Vec(readPorts, UInt(32.W)))
    })
    
    // A Register of a vector of UInts
    val reg = RegInit(VecInit(Seq.fill(32)(0.U(32.W))))

}

chisel3.iotesters.Driver(() => new RegisterFile(2) ) { c => new PeekPokeTester(c) {
    def readExpect(addr: Int, value: Int, port: Int = 0): Unit = {
        poke(c.io.raddr(port), addr)
        expect(c.io.rdata(port), value)
    }
    def write(addr: Int, value: Int): Unit = {
        poke(c.io.wen, 1)
        poke(c.io.wdata, value)
        poke(c.io.waddr, addr)
        step(1)
        poke(c.io.wen, 0)
    }
    // everything should be 0 on init
    for (i <- 0 until 32) {
        readExpect(i, 0, port = 0)
        readExpect(i, 0, port = 1)
    }

    // write 5 * addr + 3
    for (i <- 0 until 32) {
        write(i, 5 * i + 3)
    }

    // check that the writes worked
    for (i <- 0 until 32) {
        readExpect(i, if (i == 0) 0 else 5 * i + 3, port = i % 2)
    }

}}