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 chisel3.experimental._
import chisel3.internal.firrtl.KnownBinaryPoint

// No inputs or outputs (two versions).
def hello1(): Unit = print("Hello!")
def hello2 = print("Hello again!")

// Math operation: one input and one output.
def times2(x: Int): Int = 2 * x

// Inputs can have default values, and explicitly specifying the return type is optional.
// Note that we recommend specifying the return types to avoid surprises/bugs.
def timesN(x: Int, n: Int = 2) = n * x

// Call the functions listed above.
hello1()
hello2
times2(4)
timesN(4)         // no need to specify n to use the default value
timesN(4, 3)      // argument order is the same as the order where the function was defined
timesN(n=7, x=2)  // arguments may be reordered and assigned to explicitly

// These are normal functions.
def plus1funct(x: Int): Int = x + 1
def times2funct(x: Int): Int = x * 2

// These are functions as vals.
// The first one explicitly specifies the return type.
val plus1val: Int => Int = x => x + 1
val times2val = (x: Int) => x * 2

// Calling both looks the same.
plus1funct(4)
plus1val(4)
plus1funct(x=4)
//plus1val(x=4) // this doesn't work

// create our function
val plus1 = (x: Int) => x + 1
val times2 = (x: Int) => x * 2

// pass it to map, a list function
val myList = List(1, 2, 5, 9)
val myListPlus = myList.map(plus1)
val myListTimes = myList.map(times2)

// create a custom function, which performs an operation on X N times using recursion
def opN(x: Int, n: Int, op: Int => Int): Int = {
  if (n <= 0) { x }
  else { opN(op(x), n-1, op) }
}

opN(7, 3, plus1)
opN(7, 3, times2)

import scala.util.Random

// both x and y call the nextInt function, but x is evaluated immediately and y is a function
val x = Random.nextInt
def y = Random.nextInt

// x was previously evaluated, so it is a constant
println(s"x = $x")
println(s"x = $x")

// y is a function and gets reevaluated at each call, thus these produce different results
println(s"y = $y")
println(s"y = $y")

val myList = List(5, 6, 7, 8)

// add one to every item in the list using an anonymous function
// arguments get passed to the underscore variable
// these all do the same thing
myList.map( (x:Int) => x + 1 )
myList.map(_ + 1)

// a common situation is to use case statements within an anonymous function
val myAnyList = List(1, 2, "3", 4L, myList)
myAnyList.map {
  case (_:Int|_:Long) => "Number"
  case _:String => "String"
  case _ => "error"
}

val exList = List(1, 5, 7, 100)

// write a custom function to add two numbers, then use reduce to find the sum of all values in exList
def add(a: Int, b: Int): Int = ???
val sum = ???

// find the sum of exList using an anonymous function (hint: you've seen this before!)
val anon_sum = ???

// find the moving average of exList from right to left using scan; make the result (ma2) a list of doubles
def avg(a: Int, b: Double): Double = ???
val ma2 = ???

assert(add(88, 88) == 176)
assert(sum == 113)

assert(anon_sum == 113)

assert(avg(100, 100.0) == 100.0)
assert(ma2 == List(8.875, 16.75, 28.5, 50.0, 0.0))

// get some math functions
import scala.math.{abs, round, cos, Pi, pow}

// simple triangular window
val TriangularWindow: (Int, Int) => Seq[Int] = (length, bitwidth) => {
  val raw_coeffs = (0 until length).map( (x:Int) => 1-abs((x.toDouble-(length-1)/2.0)/((length-1)/2.0)) )
  val scaled_coeffs = raw_coeffs.map( (x: Double) => round(x * pow(2, bitwidth)).toInt)
  scaled_coeffs
}

// Hamming window
val HammingWindow: (Int, Int) => Seq[Int] = (length, bitwidth) => {
  val raw_coeffs = (0 until length).map( (x: Int) => 0.54 - 0.46*cos(2*Pi*x/(length-1)))
  val scaled_coeffs = raw_coeffs.map( (x: Double) => round(x * pow(2, bitwidth)).toInt)
  scaled_coeffs
}

// check it out! first argument is the window length, and second argument is the bitwidth
TriangularWindow(10, 16)
HammingWindow(10, 16)

// our FIR has parameterized window length, IO bitwidth, and windowing function
class MyFir(length: Int, bitwidth: Int, window: (Int, Int) => Seq[Int]) extends Module {
  val io = IO(new Bundle {
    val in = Input(UInt(bitwidth.W))
    val out = Output(UInt((bitwidth*2+length-1).W)) // expect bit growth, conservative but lazy
  })

  // calculate the coefficients using the provided window function, mapping to UInts
  val coeffs = window(length, bitwidth).map(_.U)
  
  // create an array holding the output of the delays
  // note: we avoid using a Vec here since we don't need dynamic indexing
  val delays = Seq.fill(length)(Wire(UInt(bitwidth.W))).scan(io.in)( (prev: UInt, next: UInt) => {
    next := RegNext(prev)
    next
  })
  
  // multiply, putting result in "mults"
  val mults = delays.zip(coeffs).map{ case(delay: UInt, coeff: UInt) => delay * coeff }
  
  // add up multiplier outputs with bit growth
  val result = mults.reduce(_+&_)

  // connect output
  io.out := result
}

// math imports
import scala.math.{pow, sin, Pi}
import breeze.signal.{filter, OptOverhang}
import breeze.signal.support.{CanFilter, FIRKernel1D}
import breeze.linalg.DenseVector

// test parameters
val length = 7
val bitwidth = 12 // must be less than 15, otherwise Int can't represent the data, need BigInt
val window = TriangularWindow

// test our FIR
Driver(() => new MyFir(length, bitwidth, window)) {
  c => new PeekPokeTester(c) {
    
    // test data
    val n = 100 // input length
    val sine_freq = 10
    val samp_freq = 100
    
    // sample data, scale to between 0 and 2^bitwidth
    val max_value = pow(2, bitwidth)-1
    val sine = (0 until n).map(i => (max_value/2 + max_value/2*sin(2*Pi*sine_freq/samp_freq*i)).toInt)
    //println(s"input = ${sine.toArray.deep.mkString(", ")}")
    
    // coefficients
    val coeffs = window(length, bitwidth)
    //println(s"coeffs = ${coeffs.toArray.deep.mkString(", ")}")

    // use breeze filter as golden model; need to reverse coefficients
    val expected = filter(DenseVector(sine.toArray), 
                          FIRKernel1D(DenseVector(coeffs.reverse.toArray), 1.0, ""), 
                          OptOverhang.None)
    //println(s"exp_out = ${expected.toArray.deep.mkString(", ")}")

    // push data through our FIR and check the result
    reset(5)
    for (i <- 0 until n) {
      poke(c.io.in, sine(i))
      if (i >= length-1) { // wait for all registers to be initialized since we didn't zero-pad the data
        expect(c.io.out, expected(i-length+1))
        //println(s"cycle $i, got ${peek(c.io.out)}, expect ${expected(i-length+1)}")
      }
      step(1)
    }
  }
}

class Neuron(inputs: Int, act: FixedPoint => FixedPoint) extends Module {
  val io = IO(new Bundle {
    val in      = Input(Vec(inputs, FixedPoint(16.W, 8.BP)))
    val weights = Input(Vec(inputs, FixedPoint(16.W, 8.BP)))
    val out     = Output(FixedPoint(16.W, 8.BP))
  })
  
  ???
}

val Step: FixedPoint => FixedPoint = ???
val ReLU: FixedPoint => FixedPoint = ???

// test our Neuron 
Driver(() => new Neuron(2, Step)) {
  c => new PeekPokeTester(c) {
    
    val inputs = Seq(Seq(0, 0), Seq(0, 1), Seq(1, 0), Seq(1, 1))
    
    // make these a sequence of two values
    val and_weights = ???
    val or_weights  = ???

    // push data through our Neuron and check the result (AND gate)
    reset(5)
    for (i <- inputs) {
      pokeFixedPoint(c.io.in(0), i(0))
      pokeFixedPoint(c.io.in(1), i(1))
      pokeFixedPoint(c.io.weights(0), and_weights(0))
      pokeFixedPoint(c.io.weights(1), and_weights(1))
      expectFixedPoint(c.io.out, i(0) & i(1), "ERROR")
      step(1)
    }
    
    // push data through our Neuron and check the result (OR gate)
    reset(5)
    for (i <- inputs) {
      pokeFixedPoint(c.io.in(0), i(0))
      pokeFixedPoint(c.io.in(1), i(1))
      pokeFixedPoint(c.io.weights(0), or_weights(0))
      pokeFixedPoint(c.io.weights(1), or_weights(1))
      expectFixedPoint(c.io.out, i(0) | i(1), "ERROR")
      step(1)
    }
  }
}