package day14

import java.io.File

fun parseInput(lines: List<String>): Pair<String, Map<String, Char>> {
  val initial = lines.first()
  val insertions =
    lines.drop(2).associate {
      val (pattern, addition) = it.split(" -> ")
      pattern to addition.first()
    }
  return initial to insertions
}

private fun combineAllMaps(vararg maps: Map<Char, Long>): Map<Char, Long> =
  maps.flatMap { it.toList() }
    .groupBy { it.first }
    .mapValues { (_, v) -> v.fold(0L) { acc, (_, v) -> acc + v } }

private fun combine2(map1: Map<Char, Long>, map2: Map<Char, Long>): Map<Char, Long> =
  combineAllMaps(map1, map2)

val memoizeCountDescendantChars = { insertions: Map<String, Char> ->
  // intermediate steps will be stored here
  val cache = mutableMapOf<Pair<String, Int>, Map<Char, Long>>()

  // Recursively build the count, storing all intermediate steps
  // in 'cache' so that they can be reused later on
  fun getAdditionCount(pair: String, steps: Int): Map<Char, Long> =
    when {
      steps == 0 -> emptyMap()
      cache.containsKey(pair to steps) -> cache[pair to steps]!!
      insertions.containsKey(pair) -> {
        val c = insertions[pair]!!
        val result = combineAllMaps(
          mapOf(c to 1L),
          getAdditionCount("${pair[0]}$c", steps - 1),
          getAdditionCount("$c${pair[1]}", steps - 1)
        )
        cache[pair to steps] = result
        result
      }
      else -> emptyMap()
    }

  ::getAdditionCount
}

fun solve(polymer: String, insertions: Map<String, Char>, steps: Int): Long {
  val memoizedGetDescendants = memoizeCountDescendantChars(insertions)

  val finalCount = sequence {
    // Count initial chars
    val initialCharCount = polymer
      .groupingBy { it }.eachCount()
      .mapValues { (_, v) -> v.toLong() }
      .toMutableMap()

    yield(initialCharCount)

    // Count descendant chars generated by each pair or chars ('abc' -> 'ab','bc')
    for (i in 0..(polymer.length - 2)) {
      val pair = polymer.substring(i, i + 2)
      yield(memoizedGetDescendants(pair, steps))
    }
  }.reduce(::combine2)

  val max = finalCount.maxOf { it.value }
  val min = finalCount.minOf { it.value }
  return max - min
}

fun main() {
  File("./input/day14.txt").useLines { lines ->
    val (initialPolymer, insertions) = parseInput(lines.toList())
    println(solve(initialPolymer, insertions, 10))
    println(solve(initialPolymer, insertions, 40))
  }
}