wordle.scala

package wordle

/** Wordle solver, game runner & simulator
  *
  * Optimizes based on a combination of an allowed word list (from the Wordle source code or any
  * other source), word frequency data and the move in the game.
  *
  * @note
  *   [[Wordle.Game]] is mutable to allow for play in an environment without easy STDIN input. Use
  *   [[Wordle.Game.nextMove()]]. All words are in lowercase. Patterns are entered as as strings of
  *   5 numbers: (2 -> green, 1 -> yellow, 0 -> gray).
  * @example
  *   {{{Wordle.Game(arrAllowedWords, arrWordCounts)().run("aloft")}}}
  *
  * @see
  *   https://dbc-029754dd-b0bf.cloud.databricks.com/#notebook/790434
  */
object Wordle {

  final val wordLen = 5
  final val numChoices = 'z' - 'a' + 1
  final val patternCount = Math.pow(3, 5).toInt
  final val aval = 'a'.toInt

  @inline def encode(ch: Char) = ch.toLower - 'a'

  case class Game(var choices: Choices) {
    def nextMove(choice: String, patternString: String): Unit =
      choices = choices.nextChoices(choice, patternString)

    def run(target: String): Int = {
      require(choices.matchInfo.exists(_.word == target))
      var move = 1
      var solved = choices.sorted.headOption.exists(_.word == target)
      while (!solved) {
        val choice = choices.sorted.head
        val word = choice.word
        val pattern = word.zipWithIndex.map { case (ch, i) =>
          if (target(i) == ch) '2' else if (target contains ch) '1' else '0'
        }.mkString
        nextMove(word, pattern)
        move += 1
        solved = choices.sorted.headOption.exists(_.word == target)
      }
      println(s"$target solved ${solved && move <= 6} in $move moves")
      move
    }

    override def toString = productPrefix + s"(${choices.matchInfo.length} words)"
  }

  object Game {
    def apply(allowed: Seq[String], wordCounts: Seq[(String, Long)], silent: Boolean = false)(
      sorter: (WordInfo, Double, Int) => Double = (wi: WordInfo, p: Double, move: Int) => -wi.score
    ): Game = {
      def log(s: String): Unit = if (silent) () else println(s)
      val wordCount = allowed.length
      val matchInfo = Array.ofDim[MatchInfo](wordCount)
      allowed.zipWithIndex.foreach { case (w, i) => matchInfo(i) = MatchInfo(w) }
      Game(Choices(matchInfo, WordCounts(wordCounts), 1, log)(sorter))
    }
  }

  case class Simulation(game: Game, words: Seq[String]) {
    val moves = words.map(word => (word, game.copy().run(word))).toMap
    val avgMoves = moves.values.sum.toDouble / words.length
    override def toString = productPrefix + s"(${words.length} words)"
  }

  case class MatchInfo(word: String) {
    require(word.length == wordLen)

    val encoded = {
      val enc = Array.ofDim[Int](wordLen)
      word.zipWithIndex.foreach { case (ch, j) =>
        enc(j) = encode(ch)
      }
      enc
    }

    val exact = Array.ofDim[Boolean](wordLen, numChoices)
    val any = Array.ofDim[Boolean](numChoices)

    encoded.zipWithIndex.foreach { case (n, i) =>
      exact(i)(n) = true
      any(n) = true
    }

    @inline final def patternValue(o: Array[Int]): Int = {
      var pv = 0
      var mult = 1
      var i = 0
      while (i < wordLen) {
        val och = o(i)
        pv += mult * (if (exact(i)(och)) 2 else if (any(och)) 1 else 0)
        mult *= 3
        i += 1
      }
      pv
    }

    @inline final def wordScore(o: Array[Array[Int]]): Double = {
      @inline def log2(x: Double) = Math.log(x) / Math.log(2.0)

      val patterns = Array.ofDim[Int](patternCount)
      o.foreach(w => patterns(patternValue(w)) += 1)

      var score = 0.0
      probabilitiesFromCounts(patterns).foreach { p: Double =>
        if (p != 0) score += p * log2(1.0 / p)
      }
      score
    }

    @inline private def probabilitiesFromCounts(a: Array[Int]): Array[Double] = {
      val p = Array.ofDim[Double](a.length)
      val total = a.sum.toDouble
      a.indices.foreach(idx => p(idx) = a(idx) / total)
      p
    }
  }

  case class WordInfo(word: String, encoding: Array[Int], score: Double)

  case class Choices(matchInfo: Array[Wordle.MatchInfo], wordCounts: WordCounts, move: Int, log: String => Unit)(
    sorter: (WordInfo, Double, Int) => Double
  ) {
    val wordInfo = {
      val other: Array[Array[Int]] = matchInfo.map(_.encoded)
      matchInfo.map(mi => WordInfo(mi.word, mi.encoded, mi.wordScore(other)))
    }
    val sorted: Array[WordInfo] = wordInfo.sortBy(wi => sorter(wi, wordCounts.wordProbability(wi.word), move))

    log("Next top choices:")
    sorted.take(20).foreach(wi => log(s"${wi.word}: ${wi.score}"))

    def nextChoices(choice: String, patternString: String): Choices = {
      log(s"With $choice resulting in $patternString and starting with ${matchInfo.length} words")
      val minusChoice = matchInfo.filter(_.word != choice)
      val nextInfos = patternString.zipWithIndex.foldLeft(minusChoice) { case (current, (v, i)) =>
        val startCount = current.length
        val ch = choice(i)
        val enc = encode(ch)
        val next = v match {
          case '0' => current.filterNot(_.any(enc))
          case '1' => current.filter(mi => mi.any(enc) && !mi.exact(i)(enc))
          case '2' => current.filter(_.exact(i)(enc))
        }
        log(s"  - dropping ${startCount - next.length} words because of pattern $v for $ch in position ${i + 1}")
        next
      }
      log(s"  - ending with ${nextInfos.length} words")

      val allowed = nextInfos.map(_.word).toSet
      copy(nextInfos, wordCounts.filter(allowed.contains), move + 1)(sorter)
    }

    override def toString = productPrefix + s"(${matchInfo.length} words)"
  }

}

case class WordCounts(words: Map[String, Long]) {
  val total = words.values.sum.toDouble
  @transient val wordProbability = words.mapValues { v => v / total }

  def filter(f: String => Boolean) = WordCounts(words.filterKeys(f).toSeq.toMap)
}

object WordCounts {
  def apply(wordCounts: Seq[(String, Long)]): WordCounts = WordCounts(wordCounts.map(wc => (wc._1, wc._2)).toMap)
}