package aoc2024

import RelPos
import println
import readInput

/*
--- Day 21: Keypad Conundrum ---
https://adventofcode.com/2024/day/21
*/
fun main() {

    val inlineTestInput = """
029A
980A
179A
456A
379A
"""

    val doorPadMap = hashMapOf(
        '7' to RelPos(-2, -3),
        '8' to RelPos(-1, -3),
        '9' to RelPos(0, -3),
        '4' to RelPos(-2, -2),
        '5' to RelPos(-1, -2),
        '6' to RelPos(0, -2),
        '1' to RelPos(-2, -1),
        '2' to RelPos(-1, -1),
        '3' to RelPos(0, -1),
        '0' to RelPos(-1, 0),
        'A' to RelPos(0, 0),
    )

    val robotPadMap = hashMapOf(
        '^' to RelPos(-1, 0),
        'A' to RelPos(0, 0),
        '<' to RelPos(-2, 1),
        'v' to RelPos(-1, 1),
        '>' to RelPos(0, 1),
    )

    data class PadState(
        var pos: RelPos,
        val padMap: Map<Char, RelPos>,
        val child: PadState? = null,
        val cache: MutableMap<Pair<RelPos, RelPos>, Long> = HashMap()
    ) {
        fun getMovement(c: Char): Long {
            val targetPos = padMap[c]!!
            return if (child == null) {
                1
            } else {
                val movement = cache.getOrPut(pos to targetPos) {
                    val child = child
                    var sum = 0L
                    if (pos.dr == 0 && targetPos.dc == -2) {
                        while (targetPos.dr < pos.dr) {
                            sum += child.getMovement('^')
                            pos = pos.translate(0, -1)
                        }
                        while (targetPos.dr > pos.dr) {
                            sum += child.getMovement('v')
                            pos = pos.translate(0, 1)
                        }
                    }
                    if (pos.dc == -2 && targetPos.dr == 0) {
                        while (targetPos.dc > pos.dc) {
                            sum += child.getMovement('>')
                            pos = pos.translate(1, 0)
                        }
                    }
                    while (targetPos.dc < pos.dc) {
                        sum += child.getMovement('<')
                        pos = pos.translate(-1, 0)
                    }
                    while (targetPos.dr > pos.dr) {
                        sum += child.getMovement('v')
                        pos = pos.translate(0, 1)
                    }
                    while (targetPos.dr < pos.dr) {
                        sum += child.getMovement('^')
                        pos = pos.translate(0, -1)
                    }
                    while (targetPos.dc > pos.dc) {
                        sum += child.getMovement('>')
                        pos = pos.translate(1, 0)
                    }

                    sum += child.getMovement('A')
                    sum
                }

                pos = targetPos
                movement
            }
        }
    }

    fun part1(input: List<String>): Long {
        val home = RelPos(0, 0)
        val myPad = PadState(home, robotPadMap)
        val robotPad2 = PadState(home, robotPadMap, myPad)
        val robotPad1 = PadState(home, robotPadMap, robotPad2)
        val doorPad = PadState(home, doorPadMap, robotPad1)
        var sum = 0L
        for (seq in input) {
            val seqVal = seq.take(3).toLong()
            sum += seq.sumOf { doorPad.getMovement(it) } * seqVal
        }
        return sum
    }

    fun part2(input: List<String>): Long {
        val home = RelPos(0, 0)
        val myPad = PadState(home, robotPadMap)
        var lastPad = myPad
        Array(25) { PadState(home, robotPadMap, lastPad).also { lastPad = it } }
        val doorPad = PadState(home, doorPadMap, lastPad)
        var sum = 0L
        for (seq in input) {
            val seqVal = seq.take(3).toLong()
            sum += seq.sumOf { doorPad.getMovement(it) } * seqVal
        }
        return sum
    }

    // test if implementation meets criteria from the description, like:
    val testInput = inlineTestInput.trim().reader().readLines()
    //val testInput = readInput("aoc2024/Day21_test")
    val testInputPart1Result = part1(testInput)
    println("Part 1 Test: $testInputPart1Result")
    val testInputPart2Result = part2(testInput)
    println("Part 2 Test: $testInputPart2Result")
    check(testInputPart1Result == 126384L)
    //check(testInputPart2Result == 0)

    val input = readInput("aoc2024/Day21")
    part1(input).println()
    part2(input).println()
}
