package vexriscv.plugin

import vexriscv._
import spinal.core._
import spinal.core.internals.Literal
import spinal.lib._
import vexriscv.demo.GenFull

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer


case class Masked(value : BigInt,care : BigInt){
  assert((value & ~care) == 0)
  var isPrime = true

  def < (that: Masked) = value < that.value || value == that.value && ~care < ~that.care

  def intersects(x: Masked) = ((value ^ x.value) & care & x.care) == 0

  def covers(x: Masked) = ((value ^ x.value) & care | (~x.care) & care) == 0

  def setPrime(value : Boolean) = {
    isPrime = value
    this
  }

  def mergeOneBitDifSmaller(x: Masked) = {
    val bit = value - x.value
    val ret = new Masked(value &~ bit, care & ~bit)
//    ret.isPrime = isPrime || x.isPrime
    isPrime = false
    x.isPrime = false
    ret
  }
  def isSimilarOneBitDifSmaller(x: Masked) = {
    val diff = value - x.value
    care == x.care && value > x.value && (diff & diff - 1) == 0
  }


  def === (hard : Bits) : Bool = (hard & care) === (value & care)

  def toString(bitCount : Int) = (0 until bitCount).map(i => if(care.testBit(i)) (if(value.testBit(i)) "1" else "0") else "-").reverseIterator.reduce(_+_)
}

class DecoderSimplePlugin(catchIllegalInstruction : Boolean = false, forceLegalInstructionComputation : Boolean = false) extends Plugin[VexRiscv] with DecoderService {
  override def add(encoding: Seq[(MaskedLiteral, Seq[(Stageable[_ <: BaseType], Any)])]): Unit = encoding.foreach(e => this.add(e._1,e._2))
  override def add(key: MaskedLiteral, values: Seq[(Stageable[_ <: BaseType], Any)]): Unit = {
    val instructionModel = encodings.getOrElseUpdate(key,ArrayBuffer[(Stageable[_ <: BaseType], BaseType)]())
    values.map{case (a,b) => {
      assert(!instructionModel.contains(a), s"Over specification of $a")
      val value = b match {
        case e: SpinalEnumElement[_] => e()
        case e: BaseType => e
      }
      instructionModel += (a->value)
    }}
  }

  override def addDefault(key: Stageable[_  <: BaseType], value: Any): Unit = {
    assert(!defaults.contains(key))
    defaults(key) = value match{
      case e : SpinalEnumElement[_] => e()
      case e : BaseType => e
    }
  }

  val defaults = mutable.HashMap[Stageable[_ <: BaseType], BaseType]()
  val encodings = mutable.HashMap[MaskedLiteral,ArrayBuffer[(Stageable[_ <: BaseType], BaseType)]]()
  var decodeExceptionPort : Flow[ExceptionCause] = null


  override def setup(pipeline: VexRiscv): Unit = {
    if(catchIllegalInstruction) {
      val exceptionService = pipeline.plugins.filter(_.isInstanceOf[ExceptionService]).head.asInstanceOf[ExceptionService]
      decodeExceptionPort = exceptionService.newExceptionPort(pipeline.decode).setName("decodeExceptionPort")
    }
  }

  override def build(pipeline: VexRiscv): Unit = {
    import pipeline.config._
    import pipeline.decode._

    val stageables = (encodings.flatMap(_._2.map(_._1)) ++ defaults.map(_._1)).toSet.toList

    var offset = 0
    var defaultValue, defaultCare = BigInt(0)
    val offsetOf = mutable.HashMap[Stageable[_ <: BaseType],Int]()

    //Build defaults value and field offset map
    stageables.foreach(e => {
      defaults.get(e) match {
        case Some(value) => {
          value.head.source match {
            case literal: EnumLiteral[_] => literal.fixEncoding(e.dataType.asInstanceOf[SpinalEnumCraft[_]].getEncoding)
            case _ =>
          }
          defaultValue += value.head.source .asInstanceOf[Literal].getValue << offset
          defaultCare += ((BigInt(1) << e.dataType.getBitsWidth) - 1) << offset

        }
        case _ =>
      }
      offsetOf(e) = offset
      offset += e.dataType.getBitsWidth
    })

    //Build spec
    val spec = encodings.map { case (key, values) =>
      var decodedValue = defaultValue
      var decodedCare = defaultCare
      for((e, literal) <- values){
        literal.head.source  match{
          case literal : EnumLiteral[_] => literal.fixEncoding(e.dataType.asInstanceOf[SpinalEnumCraft[_]].getEncoding)
          case _ =>
        }
        val offset = offsetOf(e)
        decodedValue |= literal.head.source.asInstanceOf[Literal].getValue << offset
        decodedCare  |= ((BigInt(1) << e.dataType.getBitsWidth)-1) << offset
      }
      (Masked(key.value,key.careAbout),Masked(decodedValue,decodedCare))
    }



    // logic implementation
    val decodedBits = Bits(stageables.foldLeft(0)(_ + _.dataType.getBitsWidth) bits)
    decodedBits := Symplify(input(INSTRUCTION),spec, decodedBits.getWidth)
    if(catchIllegalInstruction || forceLegalInstructionComputation) insert(LEGAL_INSTRUCTION) := Symplify.logicOf(input(INSTRUCTION), SymplifyBit.getPrimeImplicantsByTrueAndDontCare(spec.unzip._1.toSeq, Nil, 32))


    //Unpack decodedBits and insert fields in the pipeline
    offset = 0
    stageables.foreach(e => {
      insert(e).assignFromBits(decodedBits(offset, e.dataType.getBitsWidth bits))
//            insert(e).assignFromBits(RegNext(decodedBits(offset, e.dataType.getBitsWidth bits)))
      offset += e.dataType.getBitsWidth
    })


    if(catchIllegalInstruction){
      decodeExceptionPort.valid := arbitration.isValid && input(INSTRUCTION_READY) && !input(LEGAL_INSTRUCTION) // ?? HalitIt to alow decoder stage to wait valid data from 2 stages cache cache ??
      decodeExceptionPort.code := 2
      decodeExceptionPort.badAddr.assignDontCare()
    }
  }

  def bench(toplevel : VexRiscv): Unit ={
    toplevel.rework{
      import toplevel.config._
      toplevel.getAllIo.foreach{io =>
        if(io.isInput) io.assignDontCare()
        io.setAsDirectionLess()
      }
      toplevel.decode.input(INSTRUCTION).removeAssignments()
      toplevel.decode.input(INSTRUCTION) := Delay((in Bits(32 bits)).setName("instruction"),2)
      val stageables = encodings.flatMap(_._2.map(_._1)).toSet
      stageables.foreach(e => out(RegNext(RegNext(toplevel.decode.insert(e)).setName(e.getName()))))
      if(catchIllegalInstruction) out(RegNext(RegNext(toplevel.decode.insert(LEGAL_INSTRUCTION)).setName(LEGAL_INSTRUCTION.getName())))
    //  toplevel.getAdditionalNodesRoot.clear()
    }
  }
}

object DecodingBench extends App{
  SpinalVerilog{
    val top = GenFull.cpu()
    top.service(classOf[DecoderSimplePlugin]).bench(top)
    top
  }
}


object Symplify{
  val cache = mutable.HashMap[Bits,mutable.HashMap[Masked,Bool]]()
  def getCache(addr : Bits) = cache.getOrElseUpdate(addr,mutable.HashMap[Masked,Bool]())

  //Generate terms logic for the given input
  def logicOf(input : Bits,terms : Seq[Masked]) = terms.map(t => getCache(input).getOrElseUpdate(t,t === input)).asBits.orR

  //Decode 'input' b using an mapping[key, decoding] specification
  def apply(input: Bits, mapping: Iterable[(Masked, Masked)],resultWidth : Int) : Bits = {
    val addrWidth = widthOf(input)
    (for(bitId <- 0 until resultWidth) yield{
      val trueTerm = mapping.filter { case (k,t) => (t.care.testBit(bitId) && t.value.testBit(bitId))}.map(_._1)
      val falseTerm = mapping.filter { case (k,t) => (t.care.testBit(bitId) &&  !t.value.testBit(bitId))}.map(_._1)
      val symplifiedTerms = SymplifyBit.getPrimeImplicantsByTrueAndFalse(trueTerm.toSeq, falseTerm.toSeq, addrWidth)
      logicOf(input, symplifiedTerms)
    }).asBits
  }
}

object SymplifyBit{

  //Return a new term with only one bit difference with 'term' and not included in falseTerms. above => 0 to 1 dif, else 1 to 0 diff
  def genImplicitDontCare(falseTerms: Seq[Masked], term: Masked, bits: Int, above: Boolean): Masked = {
    for (i <- 0 until bits; if term.care.testBit(i)) {
      var t: Masked = null
      if(above) {
        if (!term.value.testBit(i))
          t = Masked(term.value.setBit(i), term.care)
      } else {
        if (term.value.testBit(i))
          t = Masked(term.value.clearBit(i), term.care)
      }
      if (t != null && !falseTerms.exists(_.intersects(t))) {
        t.isPrime = false
        return t
      }
    }
    null
  }

  //Return primes implicants for the trueTerms, falseTerms spec. Default value is don't care
  def getPrimeImplicantsByTrueAndFalse(trueTerms: Seq[Masked], falseTerms: Seq[Masked], inputWidth : Int): Seq[Masked] = {
    val primes = mutable.LinkedHashSet[Masked]()
    trueTerms.foreach(_.isPrime = true)
    falseTerms.foreach(_.isPrime = true)
    val trueTermByCareCount = (inputWidth to 0 by -1).map(b => trueTerms.filter(b == _.care.bitCount))
    //table[Vector[HashSet[Masked]]](careCount)(bitSetCount)
    val table = trueTermByCareCount.map(c => (0 to inputWidth).map(b => collection.mutable.Set(c.filter(b == _.value.bitCount): _*)))
    for (i <- 0 to inputWidth) {
      //Expends explicit terms
      for (j <- 0 until inputWidth - i){
        for(term <- table(i)(j)){
          table(i+1)(j) ++= table(i)(j+1).withFilter(_.isSimilarOneBitDifSmaller(term)).map(_.mergeOneBitDifSmaller(term))
        }
      }
      //Expends implicit don't care terms
      for (j <- 0 until inputWidth-i) {
        for (prime <- table(i)(j).withFilter(_.isPrime)) {
          val dc = genImplicitDontCare(falseTerms, prime, inputWidth, true)
          if (dc != null)
            table(i+1)(j) += dc mergeOneBitDifSmaller prime
        }
        for (prime <- table(i)(j+1).withFilter(_.isPrime)) {
          val dc = genImplicitDontCare(falseTerms, prime, inputWidth, false)
          if (dc != null)
            table(i+1)(j) += prime mergeOneBitDifSmaller dc
        }
      }
      for (r <- table(i))
        for (p <- r; if p.isPrime)
          primes += p
    }

    def optimise() {
      val duplicateds = primes.filter(prime => verifyTrueFalse(primes.filterNot(_ == prime), trueTerms, falseTerms))
      if(duplicateds.nonEmpty) {
        primes -= duplicateds.maxBy(_.care.bitCount)
        optimise()
      }
    }

    optimise()

    verifyTrueFalse(primes, trueTerms, falseTerms)
    var duplication = 0
    for(prime <- primes){
      if(verifyTrueFalse(primes.filterNot(_ == prime), trueTerms, falseTerms)){
        duplication += 1
      }
    }
    if(duplication != 0){
      PendingError(s"Duplicated primes : $duplication")
    }
    primes.toSeq
  }

  //Verify that the 'terms' doesn't violate the trueTerms ++ falseTerms spec
  def verifyTrueFalse(terms : Iterable[Masked], trueTerms : Seq[Masked], falseTerms : Seq[Masked]): Boolean ={
    return (trueTerms.forall(trueTerm => terms.exists(_ covers trueTerm))) && (falseTerms.forall(falseTerm => !terms.exists(_ covers falseTerm)))
  }

  def checkTrue(terms : Iterable[Masked], trueTerms : Seq[Masked]): Boolean ={
    return trueTerms.forall(trueTerm => terms.exists(_ covers trueTerm))
  }


  def getPrimeImplicantsByTrue(trueTerms: Seq[Masked], inputWidth : Int) : Seq[Masked] = getPrimeImplicantsByTrueAndDontCare(trueTerms, Nil, inputWidth)

  // Return primes implicants for the trueTerms, default value is False.
  // You can insert don't care values by adding non-prime implicants in the trueTerms
  // Will simplify the trueTerms from the most constrained ones to the least constrained ones
  def getPrimeImplicantsByTrueAndDontCare(trueTerms: Seq[Masked],dontCareTerms: Seq[Masked], inputWidth : Int): Seq[Masked] = {
    val primes = mutable.LinkedHashSet[Masked]()
    trueTerms.foreach(_.isPrime = true)
    dontCareTerms.foreach(_.isPrime = false)
    val termsByCareCount = (inputWidth to 0 by -1).map(b => (trueTerms ++ dontCareTerms).filter(b == _.care.bitCount))
    //table[Vector[HashSet[Masked]]](careCount)(bitSetCount)
    val table = termsByCareCount.map(c => (0 to inputWidth).map(b => collection.mutable.Set(c.filter(m => b == m.value.bitCount): _*)))
    for (i <- 0 to inputWidth) {
      for (j <- 0 until inputWidth - i){
        for(term <- table(i)(j)){
          table(i+1)(j) ++= table(i)(j+1).withFilter(_.isSimilarOneBitDifSmaller(term)).map(_.mergeOneBitDifSmaller(term))
        }
      }
      for (r <- table(i))
        for (p <- r; if p.isPrime)
          primes += p
    }


    def optimise() {
      val duplicateds = primes.filter(prime => checkTrue(primes.filterNot(_ == prime), trueTerms))
      if(duplicateds.nonEmpty) {
        primes -= duplicateds.maxBy(_.care.bitCount)
        optimise()
      }
    }

    optimise()


    var duplication = 0
    for(prime <- primes){
      if(checkTrue(primes.filterNot(_ == prime), trueTerms)){
        duplication += 1
      }
    }
    if(duplication != 0){
      PendingError(s"Duplicated primes : $duplication")
    }
    primes.toSeq
  }

  def main(args: Array[String]) {
    {
//      val default = Masked(0, 0xF)
//      val primeImplicants = List(4, 8, 10, 11, 12, 15).map(v => Masked(v, 0xF))
//      val dcImplicants = List(9, 14).map(v => Masked(v, 0xF).setPrime(false))
//      val reducedPrimeImplicants = getPrimeImplicantsByTrueAndDontCare(primeImplicants, dcImplicants, 4)
//      println("UUT")
//      println(reducedPrimeImplicants.map(_.toString(4)).mkString("\n"))
//      println("REF")
//      println("-100\n10--\n1--0\n1-1-")
    }

    {
      val primeImplicants = List(0).map(v => Masked(v, 0xF))
      val dcImplicants = (1 to 15).map(v => Masked(v, 0xF))
      val reducedPrimeImplicants = getPrimeImplicantsByTrueAndDontCare(primeImplicants, dcImplicants, 4)
      println("UUT")
      println(reducedPrimeImplicants.map(_.toString(4)).mkString("\n"))
    }
    {
      val trueTerms = List(0, 15).map(v => Masked(v, 0xF))
      val falseTerms = List(3).map(v => Masked(v, 0xF))
      val primes =  getPrimeImplicantsByTrueAndFalse(trueTerms, falseTerms, 4)
      println(primes.map(_.toString(4)).mkString("\n"))
    }
  }
}