// SPDX-License-Identifier: Apache-2.0

package chiselTests.util.random

import chisel3._
import chisel3.stage.ChiselStage
import chisel3.util.{Cat, Counter, Enum}
import chisel3.util.random._
import chisel3.testers.{BasicTester, TesterDriver}
import chiselTests.{ChiselFlatSpec, Utils}

import math.pow

class FooLFSR(val reduction: LFSRReduce, seed: Option[BigInt]) extends PRNG(4, seed) with LFSR {
  def delta(s: Seq[Bool]): Seq[Bool] = s
}

class LFSRMaxPeriod(gen: => UInt) extends BasicTester {

  val rv = gen
  val started = RegNext(true.B, false.B)
  val seed = withReset(!started) { RegInit(rv) }

  val (_, wrap) = Counter(started, math.pow(2.0, rv.getWidth).toInt - 1)

  when(rv === seed && started) {
    chisel3.assert(wrap)
    stop()
  }

  val last = RegNext(rv)
  chisel3.assert(rv =/= last, "LFSR last value (0b%b) was equal to current value (0b%b)", rv, last)

}

/**
  * This test creates two 4 sided dice.
  * Each cycle it adds them together and adds a count to the bin corresponding to that value
  * The asserts check that the bins show the correct distribution.
  */
class LFSRDistribution(gen: => UInt, cycles: Int = 10000) extends BasicTester {

  val rv = gen
  val bins = Reg(Vec(8, UInt(32.W)))

  // Use tap points on each LFSR so values are more independent
  val die0 = Cat(Seq.tabulate(2) { i => rv(i) })
  val die1 = Cat(Seq.tabulate(2) { i => rv(i + 2) })

  val (trial, done) = Counter(true.B, cycles)

  val rollValue = die0 +& die1 // Note +& is critical because sum will need an extra bit.

  bins(rollValue) := bins(rollValue) + 1.U

  when(done) {
    printf(p"bins: $bins\n") // Note using the printable interpolator p"" to print out a Vec

    // test that the distribution feels right.
    assert(bins(1) > bins(0))
    assert(bins(2) > bins(1))
    assert(bins(3) > bins(2))
    assert(bins(4) < bins(3))
    assert(bins(5) < bins(4))
    assert(bins(6) < bins(5))
    assert(bins(7) === 0.U)

    stop()
  }
}

/** This tests that after reset an LFSR is not locked up. This manually sets the seed of the LFSR at run-time to the
  * value that would cause it to lock up. It then asserts reset. The next cycle it checks that the value is NOT the
  * locked up value.
  * @param gen an LFSR to test
  * @param lockUpValue the value that would lock up the LFSR
  */
class LFSRResetTester(gen: => LFSR, lockUpValue: BigInt) extends BasicTester {

  val lfsr = Module(gen)
  lfsr.io.seed.valid := false.B
  lfsr.io.seed.bits := DontCare
  lfsr.io.increment := true.B

  val (count, done) = Counter(true.B, 5)

  lfsr.io.seed.valid := count === 1.U
  lfsr.io.seed.bits := lockUpValue.U(lfsr.width.W).asBools
  lfsr.io.increment := true.B

  when(count === 2.U) {
    assert(lfsr.io.out.asUInt === lockUpValue.U, "LFSR is NOT locked up, but should be!")
  }

  lfsr.reset := count === 3.U

  when(count === 4.U) {
    assert(lfsr.io.out.asUInt =/= lockUpValue.U, "LFSR is locked up, but should not be after reset!")
  }

  when(done) {
    stop()
  }

}

class LFSRSpec extends ChiselFlatSpec with Utils {

  def periodCheck(gen: (Int, Set[Int], LFSRReduce) => PRNG, reduction: LFSRReduce, range: Range): Unit = {
    val testName = s"have a maximal period over a range of widths (${range.head} to ${range.last})" +
      s" using ${reduction.getClass}"
    it should testName in {
      range.foreach { width =>
        LFSR.tapsMaxPeriod(width).foreach { taps =>
          info(s"""width $width okay using taps: ${taps.mkString(", ")}""")
          assertTesterPasses(
            new LFSRMaxPeriod(PRNG(gen(width, taps, reduction))),
            annotations = TesterDriver.verilatorOnly
          )
        }
      }
    }
  }

  behavior.of("LFSR")

  it should "throw an exception if initialized to a seed of zero for XOR configuration" in {
    {
      the[IllegalArgumentException] thrownBy extractCause[IllegalArgumentException] {
        ChiselStage.elaborate(new FooLFSR(XOR, Some(0)))
      }
    }.getMessage should include("Seed cannot be zero")
  }

  it should "throw an exception if initialized to a seed of all ones for XNOR configuration" in {
    {
      the[IllegalArgumentException] thrownBy extractCause[IllegalArgumentException] {
        ChiselStage.elaborate(new FooLFSR(XNOR, Some(15)))
      }
    }.getMessage should include("Seed cannot be all ones")
  }

  it should "reset correctly without a seed for XOR configuration" in {
    assertTesterPasses(new LFSRResetTester(new FooLFSR(XOR, None), 0))
  }

  it should "reset correctly without a seed for XNOR configuration" in {
    assertTesterPasses(new LFSRResetTester(new FooLFSR(XNOR, None), 15))
  }

  behavior.of("MaximalPeriodGaloisLFSR")

  it should "throw an exception if no LFSR taps are known" in {
    {
      the[IllegalArgumentException] thrownBy extractCause[IllegalArgumentException] {
        ChiselStage.elaborate(new MaxPeriodGaloisLFSR(787))
      }
    }.getMessage should include("No max period LFSR taps stored for requested width")
  }

  periodCheck((w: Int, t: Set[Int], r: LFSRReduce) => new GaloisLFSR(w, t, reduction = r), XOR, 2 to 16)
  periodCheck((w: Int, t: Set[Int], r: LFSRReduce) => new GaloisLFSR(w, t, reduction = r), XNOR, 2 to 16)

  ignore should "have a sane distribution for larger widths" in {
    ((17 to 32) ++ Seq(64, 128, 256, 512, 1024, 2048, 4096)).foreach { width =>
      info(s"width $width okay!")
      assertTesterPasses(new LFSRDistribution(LFSR(width), math.pow(2, 22).toInt))
    }
  }

  behavior.of("MaximalPeriodFibonacciLFSR")

  periodCheck((w: Int, t: Set[Int], r: LFSRReduce) => new FibonacciLFSR(w, t, reduction = r), XOR, 2 to 16)
  periodCheck((w: Int, t: Set[Int], r: LFSRReduce) => new FibonacciLFSR(w, t, reduction = r), XNOR, 2 to 16)

  behavior.of("LFSR maximal period taps")

  it should "contain all the expected widths" in {
    ((2 to 786) ++ Seq(1024, 2048, 4096)).foreach(LFSR.tapsMaxPeriod.keys should contain(_))
  }

}