package Y2023.D17

import Y2023.D17.Direction.*
import java.lang.IllegalStateException
import kotlin.math.min

class City(private val blocks: Map<Pair<Int, Int>, Int>) {

    fun navigate(): Int {
        val heatLossMap = mutableMapOf<Triple<Pair<Int, Int>, Direction, Int>, Step>()
        var leastHeatLoss = Int.MAX_VALUE
        val goal = blocks.maxBy { it.key.first + it.key.second }!!.key
        var options = listOf(
            (0 to 1) to Step(0 to 1, 0, Right, 1, 0),
            (1 to 0) to Step(1 to 0, 0, Down, 1, 0)
        )
        while (options.isNotEmpty()) {
            val newOptions = mutableListOf<Pair<Pair<Int, Int>, Step>>()
            options.forEach { option ->
                if (option.second.totalHeatLoss > leastHeatLoss) return@forEach
                val heatLossInCity = blocks[option.first] ?: -1
                if (heatLossInCity == -1) return@forEach

                val cacheKey = Triple(option.first, option.second.direction, option.second.straight)
                if (heatLossMap[cacheKey]?.totalHeatLoss ?: Int.MAX_VALUE <= option.second.totalHeatLoss) return@forEach
                heatLossMap[cacheKey] = option.second

                if (option.first == goal && option.second.afterTurn == 0) {
                    leastHeatLoss = min(leastHeatLoss, option.second.totalHeatLoss + heatLossInCity)
                    println("!!!")
                    println(leastHeatLoss)
                    println("!!!")
                    return@forEach
                }
                if (option.second.straight != 9) newOptions += option.first + option.second.direction.moveForward() to option.second.goStraight(heatLossInCity)
                if (option.second.afterTurn > 0) return@forEach
                val clockwise = option.second.direction.turnCW()
                val counterClockwise = option.second.direction.turnCCW()
                newOptions += option.first + clockwise.moveForward() to option.second.turn(heatLossInCity, clockwise)
                newOptions += option.first + counterClockwise.moveForward() to option.second.turn(heatLossInCity, counterClockwise)
            }
            options = newOptions
        }
        return leastHeatLoss
    }

    private fun paint(heatLossMap: MutableMap<Triple<Pair<Int, Int>, Direction, Int>, Step>, goal: Pair<Int, Int>) {
        println(" D  U  L  R \t".repeat(goal.second + 1))
        (0 .. goal.first).forEach {y ->
            (0 .. goal.second).forEach {x ->
                val directions = listOf(Down, Up, Left, Right)
                val steps = 1..3
                directions.forEach {
                    steps.forEach {step ->
                        val message = "${heatLossMap[Triple((y to x),it,step)]?.totalHeatLoss ?: ""} "
                        print(message.padStart(3, ' '))
                    }
                }
                print("\t")
            }
            println()
        }
        println()
    }
}

data class Step(val block: Pair<Int, Int>, val totalHeatLoss: Int, val direction: Direction, val straight: Int, val afterTurn: Int) {
    fun goStraight(heatLossInCity: Int) =
        Step(block + direction.moveForward(), totalHeatLoss + heatLossInCity, direction, straight + 1, if (afterTurn == 0) 0 else afterTurn-1)

    fun turn(heatLossInCity: Int, newDirection: Direction): Step {
        return Step(block + newDirection.moveForward(), totalHeatLoss + heatLossInCity, newDirection, 0, 3)
    }
}

enum class Direction {
    Start, Right, Down, Left, Up;

    fun moveForward() =
        when (this) {
            Down -> 1 to 0
            Up -> -1 to 0
            Right -> 0 to 1
            Left -> 0 to -1
            Start -> throw IllegalStateException("Moved forward on Start")
        }

    fun turnCCW() =
        when (this) {
            Down -> Right
            Up -> Left
            Right -> Up
            Left -> Down
            Start -> throw IllegalStateException("Moved forward on Start")
        }

    fun turnCW() =
        when (this) {
            Down -> Left
            Up -> Right
            Right -> Down
            Left -> Up
            Start -> throw IllegalStateException("Moved forward on Start")
        }

}

private operator fun Pair<Int, Int>.plus(other: Pair<Int, Int>) = this.first + other.first to this.second + other.second