package aoc2023

import CharGrid
import RelPos
import println
import readInput
import java.util.*
import kotlin.math.min

/*
--- Day 17: Clumsy Crucible ---
https://adventofcode.com/2023/day/17
*/
fun main() {

    val inlineTestInput = """
2413432311323
3215453535623
3255245654254
3446585845452
4546657867536
1438598798454
4457876987766
3637877979653
4654967986887
4564679986453
1224686865563
2546548887735
4322674655533
"""

    var best = Int.MAX_VALUE
    var grid = CharGrid(arrayOf(charArrayOf(' ')))
    var bestHeat = IntArray(0)

    data class LavaContext(val c: Int, val r: Int, val dir: RelPos, val dirCount: Int, val heatloss: Int, val parent: LavaContext? = null) {
        override fun equals(other: Any?): Boolean {
            if (this === other) return true
            if (javaClass != other?.javaClass) return false

            other as LavaContext

            if (c != other.c) return false
            if (r != other.r) return false
            if (dir != other.dir) return false
            if (dirCount != other.dirCount) return false
            if (heatloss != other.heatloss) return false

            return true
        }

        override fun hashCode(): Int {
            var result = c
            result = 31 * result + r
            result = 31 * result + dir.hashCode()
            result = 31 * result + dirCount
            result = 31 * result + heatloss
            return result
        }
    }

    val contexts = PriorityQueue(compareBy<LavaContext> { it.heatloss - it.c - it.r })
    val visitCache = HashSet<LavaContext>()

    fun astar(minSame: Int, maxSame: Int) {
        while (contexts.isNotEmpty()) {
            do {
                val ctx = contexts.remove()!!
                if (ctx.heatloss >= best) break
                val newHeatLoss = ctx.heatloss + grid[ctx.c, ctx.r].digitToInt()
                if (newHeatLoss > best) break
                if (ctx.c == grid.width - 1 && ctx.r == grid.height - 1) {
                    if (ctx.dirCount >= minSame) {
                        best = best.coerceAtMost(newHeatLoss)
                        val path = grid.copyOf()
                        var bt: LavaContext? = ctx
                        while (bt != null) {
                            path[bt.c, bt.r] = '*'
                            bt = bt.parent
                        }
                    }
                    break
                }
                val p = ctx.r * grid.width + ctx.c
                if (newHeatLoss - 10 > bestHeat[p]) break
                bestHeat[p] = min(bestHeat[p], newHeatLoss)
                val newDirs = (grid.matchRelative(ctx.c, ctx.r, CharGrid.PLUS_POS) { it != '#' }
                    .filter { ctx.dir.dc != -it.dc || ctx.dir.dr != -it.dr })
                newDirs.forEach {
                    val take = ((it == ctx.dir) && (ctx.dirCount + 1 < maxSame)) ||
                        ((it != ctx.dir) && (ctx.dirCount + 1 > minSame))
                    if (take) {
                        val newDirCount = if (it == ctx.dir) ctx.dirCount + 1 else 0
                        val newCtx = LavaContext(ctx.c + it.dc, ctx.r + it.dr, it, newDirCount, newHeatLoss, ctx)
                        if (!visitCache.contains(newCtx)) {
                            visitCache.add(newCtx)
                            contexts.add(newCtx)
                        }
                    }
                }
            } while (false)
        }
    }

    fun part1(input: List<String>): Int {
        grid = CharGrid(input, '#')
        best = (grid.width + grid.height) * 9
        bestHeat = IntArray(grid.width * grid.height) { best }
        val context = LavaContext(0, 0, RelPos(0, 0), 1, -grid[0, 0].digitToInt())
        visitCache.clear()
        contexts.clear()
        contexts.add(context)
        astar(0, 3)
        return best
    }

    fun part2(input: List<String>): Int {
        grid = CharGrid(input, '#')
        best = (grid.width + grid.height) * 9
        bestHeat = IntArray(grid.width * grid.height) { best }
        val context = LavaContext(0, 0, RelPos(0, 0), 4, -grid[0, 0].digitToInt())
        visitCache.clear()
        contexts.clear()
        contexts.add(context)
        astar(3, 10)
        return best
    }

    // test if implementation meets criteria from the description, like:
    val testInput = inlineTestInput.trim().reader().readLines()
    //val testInput = readInput("aoc2023/Day17_test")
    val testInputPart1Result = part1(testInput)
    println("Part 1 Test: $testInputPart1Result")
    val testInputPart2Result = part2(testInput)
    println("Part 2 Test: $testInputPart2Result")
    check(testInputPart1Result == 102)
    check(testInputPart2Result == 94)

    val input = readInput("aoc2023/Day17")
    part1(input).println()
    part2(input).println()
}
