// SPDX-License-Identifier: Apache-2.0

package chisel3.util

import chisel3._
import chisel3.experimental.SourceInfo
import chisel3.util.{log2Up, BitPat, Mux1H, PriorityMux}
import chisel3.util.experimental.decode.{decoder, TruthTable}
import chisel3.util.experimental.InlineInstance
import scala.collection.SeqMap
import scala.collection.mutable.{ArrayBuffer, LinkedHashMap, LinkedHashSet}

private[this] object Detail {

  // A passthrough module used to realize the readonly method.
  class ReadOnlyModule[A <: Data](gen: A) extends RawModule with InlineInstance {
    val in = IO(Flipped(gen))
    val out = IO(gen)
    out :<>= in
  }

  // TODO: This method produces sub-par error messages. Replace usages of this
  // with first-class Chisel support once it lands.
  def readOnly[A <: Data](
    gen: A
  ): A = {
    val readOnlyModule = Module(new ReadOnlyModule(chiselTypeOf(gen)))
    readOnlyModule.in :<>= gen
    readOnlyModule.out
  }

}

/** A sparse vector.  This should be created using the [[SparseVec$]] factory. */
class SparseVec[A <: Data] private (size: Int, gen: A, mapping: Seq[(Int, A)]) {

  // An inverse mapping of values to keys.  This is a weak form of compressing
  // the number of actual items in the vector.
  private val dataMap = LinkedHashMap
    .empty[A, LinkedHashSet[Int]]
    .withDefault(_ => LinkedHashSet.empty[Int])

  // Populate the dataMap and determine if a zeroValue exists.  A zeroValue is a
  // value at index zero that the user specified.  The zeroValue is important
  // because it will set the out-of-bounds behavior of the SparseVec.
  // Additionally, check and error if the same index is specified twice.
  private var zeroValue:        Option[(Int, A)] = None
  private val duplicates:       LinkedHashSet[Int] = LinkedHashSet.empty[Int]
  private var nonUniqueMapping: Boolean = false
  mapping.foreach {
    case (k, v) =>
      dataMap(v) = dataMap(v) += k
      if (k == 0)
        zeroValue = Some((k, v))
      if (!duplicates.add(k))
        nonUniqueMapping = true
  }

  if (nonUniqueMapping) {
    // TODO: Improve this error message.
    throw new ChiselException("non-unique mapping in SparseVec")
  }

  // Determine if a default value needs to exist.  This should only exist if
  // there are fewer map items than the size of the Vec.  (As this is a
  // SparseVec, this is expected to be a Some.)
  private val defaultValue: Option[A] = mapping match {
    case full if mapping.size == size => None
    case _                            => Some(WireInit(gen, DontCare))
  }

  /** Read a value from a [[SparseVec]]. The returned value is read-only.
    *
    * @param addr the address of the value to read from the vec
    * @param lookupType the type of lookup, e.g., binary, one-hot, or when-based
    * @param _sourceInfo implicit source locator information
    * @return a read-only value from the specified address
    * @throws ChiselException if the returned value is written to
    */
  def apply(addr: UInt, lookupType: SparseVec.Type = SparseVec.Binary)(implicit _sourceInfo: SourceInfo): A = {
    val result = lookupType match {
      // Generate a lookup using a decoder.
      case d: SparseVec.Decoder =>
        // The number of bits required to represent all addresses of the vector.
        val addrWidth = log2Up(size)

        // The address where user-provided data starts.  If there is no default
        // value needed, then start user data at address zero.
        val baseEncodedAddress = defaultValue match {
          case Some(_) => 1
          case None    => 0
        }

        // The total number of output addresses that need to be encoded.
        val encodedSize = dataMap.size + baseEncodedAddress + {
          zeroValue match {
            case Some(_) => 1
            case None    => 0
          }
        }

        val legalBitPats = dataMap.zipWithIndex.flatMap {
          case ((_, v), i) =>
            v.map { vv =>
              BitPat(vv.U(addrWidth.W)) -> BitPat(d.encoding(i + baseEncodedAddress, encodedSize))
            }
        } ++ {
          // If out-of-bounds behavior is possible, then add BitPats for all
          // out-of-bounds accesses.
          zeroValue match {
            case None => Seq.empty
            case Some((key, _)) =>
              (size until BigInt(addrWidth).pow(2).toInt).map {
                case i =>
                  BitPat(i.U(addrWidth.W)) -> BitPat(d.encoding(baseEncodedAddress, encodedSize))
              }
          }
        }

        // Generate the truth table.  Use a default value of zero.
        val ttable = TruthTable(
          legalBitPats,
          BitPat.N(encodedSize)
        )

        // Generate the lookup.
        d.lookup(decoder(addr, ttable), VecInit((defaultValue ++ dataMap.keys).toSeq))

      // Generate a lookup using a sequence of when statements.
      case SparseVec.IfElse =>
        val result = Wire(gen)

        // The default value _must_ be specified to bypass initialization checking.
        result := 0.U

        // Generate one when statement for each value.  Combine index checking
        // into an or reduction.
        dataMap.foreach {
          case (k, v) =>
            when(VecInit(v.map(addr === _.U).toSeq).reduce(_ || _)) {
              result := k
            }
        }

        // Return the zero value on out-of-bounds indices if these are possible.
        zeroValue match {
          case Some((_, v)) =>
            when(addr >= size.U) {
              result := v
            }
          case None =>
        }
        result
    }

    // Return a read-only value to avoid user mistakes where they try to write
    // to a SparseVec.
    Detail.readOnly(result)
  }

}

/** Utility for creating Sparse Vectors that behave identical to [[Vec]]s.
  *
  * A [[SparseVec]] is like a [[Vec]] that only has some subset of its constituent elements set to values.
  */
object SparseVec {

  /** The root type of how a [[SparseVec$]] can be accessed. */
  sealed trait Type

  /** A [[SparseVec$]] that is accessed using a [[chisel3.util.experimental.decode.decoder]]. */
  sealed trait Decoder extends Type {
    def encoding(index: Int, width: Int): UInt

    def lookup[A <: Data](
      index:  UInt,
      values: VecLike[A]
    )(
      implicit _sourceInfo: SourceInfo
    ): A
  }

  /** A [[SparseVec$]] accessor that uses a binary-encoded lookup. */
  case object Binary extends Decoder {
    override final def encoding(index: Int, width: Int) = index.U(width.W)

    override final def lookup[A <: Data](
      index:  UInt,
      values: VecLike[A]
    )(
      implicit _sourceInfo: SourceInfo
    ) = values(index)
  }

  /** A [[SparseVec$]] accessor that uses a one-hot-encoded lookup. */
  case object OneHot extends Decoder {
    override final def encoding(index: Int, width: Int) = (1 << index).U((width).W)

    override final def lookup[A <: Data](
      index:  UInt,
      values: VecLike[A]
    )(
      implicit _sourceInfo: SourceInfo
    ) = Mux1H(index, values)
  }

  /** A [[SparseVec$]] accessor that uses Chisel's [[when]] abstraction for lookup. */
  case object IfElse extends Type

  /** Create a new [[SparseVec]] of a specified size, type, and mapping.
    *
    * @param size the virtual size of the vector
    * @param gen the type of each element of the vector
    * @param mapping a mapping of index to value
    * @param _sourceInfo implicit source locator information
    * @return a [[SparseVec]] that can be read from
    * @throws ChiselException if the provided mapping has duplicate keys
    */
  def apply[A <: Data](
    size:    Int,
    gen:     A,
    mapping: Seq[(Int, A)]
  )(
    implicit _sourceInfo: SourceInfo
  ): SparseVec[A] = new SparseVec(size, gen, mapping)

}
