package aoc2024

import CharGrid
import RelPos
import println
import readInput

/*
--- Day 16: Reindeer Maze ---
https://adventofcode.com/2024/day/16
*/
fun main() {

    val inlineTestInput = """
###############
#.......#....E#
#.#.###.#.###.#
#.....#.#...#.#
#.###.#####.#.#
#.#.#.......#.#
#.#.#####.###.#
#...........#.#
###.#.#####.#.#
#...#.....#.#.#
#.#.#.###.#.#.#
#.....#...#.#.#
#.###.#.#.#.#.#
#S..#.....#...#
###############
"""
    val inlineTestInput2 = """
#################
#...#...#...#..E#
#.#.#.#.#.#.#.#.#
#.#.#.#...#...#.#
#.#.#.#.###.#.#.#
#...#.#.#.....#.#
#.#.#.#.#.#####.#
#.#...#.#.#.....#
#.#.#####.#.###.#
#.#.#.......#...#
#.#.###.#####.###
#.#.#...#.....#.#
#.#.#.#####.###.#
#.#.#.........#.#
#.#.#.#########.#
#S#.............#
#################
"""

    data class Node(val pos: RelPos, val dir: RelPos, val relLen: Int, val path: LinkedHashSet<Pair<RelPos, RelPos>>)

    val costMap = hashMapOf(
        RelPos(0, -1) to listOf(RelPos(0, -1) to 1, RelPos(1, 0) to 1001, RelPos(-1, 0) to 1001, RelPos(0, 1) to 2001),
        RelPos(1, 0) to listOf(RelPos(1, 0) to 1, RelPos(0, 1) to 1001, RelPos(0, -1) to 1001, RelPos(-1, 0) to 2001),
        RelPos(0, 1) to listOf(RelPos(0, 1) to 1, RelPos(-1, 0) to 1001, RelPos(1, 0) to 1001, RelPos(0, -1) to 2001),
        RelPos(-1, 0) to listOf(RelPos(-1, 0) to 1, RelPos(0, 1) to 1001, RelPos(0, -1) to 1001, RelPos(1, 0) to 2001)
    )

    fun part1(input: List<String>): Int {
        val grid = CharGrid(input, '#')
        val queue = ArrayDeque<Node>()
        val (startc, startr) = grid.findMatches { it == 'S' }[0]
        val (endc, endr) = grid.findMatches { it == 'E' }[0]
        val endpos = RelPos(endc, endr)

        grid[startc, startr] = '.'
        grid[endpos] = '.'
        queue.add(Node(RelPos(startc, startr), RelPos(1, 0), 0, LinkedHashSet()))
        val bestCosts = HashMap<Pair<RelPos, RelPos>, Int>()
        var minCost = Int.MAX_VALUE
        while (queue.isNotEmpty()) {
            val node = queue.removeFirst()
            val pos = node.pos
            val dir = node.dir
            val cost = node.relLen
            val path = node.path
            path.add(pos to dir)
            if (pos == endpos) {
                minCost = minCost.coerceAtMost(cost)
                continue
            }
            if ((bestCosts[pos to dir] ?: Int.MAX_VALUE) <= cost) {
                continue
            }
            bestCosts[pos to dir] = cost
            val nextDirs =
                costMap[dir]!!.filter { grid[pos.translate(it.first)] == '.' && cost + it.second < minCost && !path.contains(pos.translate(it.first) to it.first) }
                    .map { Node(pos.translate(it.first), it.first, cost + it.second, LinkedHashSet(path)) }

            queue.addAll(nextDirs)
        }

        return minCost
    }

    fun part2(input: List<String>): Int {
        val grid = CharGrid(input, '#')
        val queue = ArrayDeque<Node>()
        val (startc, startr) = grid.findMatches { it == 'S' }[0]
        val (endc, endr) = grid.findMatches { it == 'E' }[0]
        val endpos = RelPos(endc, endr)

        grid[startc, startr] = '.'
        grid[endpos] = '.'
        queue.add(Node(RelPos(startc, startr), RelPos(1, 0), 0, LinkedHashSet()))
        val bestCosts = HashMap<Pair<RelPos, RelPos>, Int>()
        var minCost = Int.MAX_VALUE
        val bestNodes = HashSet<RelPos>()
        while (queue.isNotEmpty()) {
            val node = queue.removeFirst()
            val pos = node.pos
            val dir = node.dir
            val cost = node.relLen
            val path = node.path
            path.add(pos to dir)
            if (pos == endpos) {
                if (cost < minCost) {
                    bestNodes.clear()
                    bestNodes.addAll(path.map { it.first })
                }
                if (minCost == cost) {
                    bestNodes.addAll(path.map { it.first })
                }
                minCost = minCost.coerceAtMost(cost)
                continue
            }
            if ((bestCosts[pos to dir] ?: Int.MAX_VALUE) < cost) {
                continue
            }
            bestCosts[pos to dir] = cost
            val nextDirs =
                costMap[dir]!!.filter { grid[pos.translate(it.first)] == '.' && cost + it.second < minCost && !path.contains(pos.translate(it.first) to it.first) }
                    .map { Node(pos.translate(it.first), it.first, cost + it.second, LinkedHashSet(path)) }

            queue.addAll(nextDirs)
        }

        return bestNodes.size
    }

    // test if implementation meets criteria from the description, like:
    val testInput = inlineTestInput.trim().reader().readLines()
    val testInput2 = inlineTestInput2.trim().reader().readLines()
    //val testInput = readInput("aoc2024/Day16_test")
    val testInputPart1Result = part1(testInput)
    println("Part 1 Test: $testInputPart1Result")
    val testInputPart1Result2 = part1(testInput2)
    println("Part 1 Test 2: $testInputPart1Result2")
    val testInputPart2Result = part2(testInput)
    println("Part 2 Test: $testInputPart2Result")
    val testInputPart2Result2 = part2(testInput2)
    println("Part 2 Test2: $testInputPart2Result2")
    check(testInputPart1Result == 7036)
    check(testInputPart1Result2 == 11048)
    check(testInputPart2Result == 45)
    check(testInputPart2Result2 == 64)

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