day18_part1.kt

/*
Kotlin:
Messy - the mistake I made was not to go for memoization from the beginning which led to a lot 
of churn which is obvious from the currrent look of the code.

Run like so:
        val maze = Maze.loadFrom("prod")
        maze.scan()
        val preparedSegments = maze.prepareSegments()
        val nodes = maze.createNodeMesh(preparedSegments)
        assertEquals(3146, maze.findShortestPath(nodes))

where prod is a file with the puzzle input
*/
import java.io.File
import kotlin.math.absoluteValue

data class Pos(val row: Int, val col: Int) {
    fun dist(other: Pos) =
       (row - other.row).absoluteValue + (col - other.col).absoluteValue

    fun offset(o: Pos) = Pos(row + o.row, col + o.col)
}

data class PathInfo(val length: Int, val traversedDoors: Set<Char>)

data class Node(val pos: Pos, val connectedSegments: List<Maze.Segment>)

data class Lock(val pos: Pos, val c: Char)

data class Key(val pos: Pos, val c: Char)

data class Maze(private val data: Map<Pos, Char>, private val nRows: Int, private val nCols: Int) {
    val rawSegments = mutableListOf<Segment>()
    var locks = mutableListOf<Lock>()
    var keys = mutableListOf<Key>()
    var allKeys: MutableSet<Char> = mutableSetOf()
    var startPos: Pos? = null
    var memoized: MutableMap<Pair<Char, Set<Char>>, Int?> = mutableMapOf()

    companion object {
        fun loadFrom(fileName: String) : Maze {
            val data = mutableMapOf<Pos, Char>()
            var row = 0
            File(fileName).forEachLine { line ->
                line.trimEnd().forEachIndexed{ col,c  -> data[Pos(row, col)] = c }
                row += 1
            }
            val dim = data.keys.fold(-1 to -1){(max_row, max_col), (row, col) -> max_row.coerceAtLeast(row) to max_col.coerceAtLeast(col) }
            return Maze(data, dim.first, dim.second)
        }
    }

    data class Segment(val endA: Pos, val endB: Pos) {
        override fun toString(): String {
            return "$endA-$endB"
        }

        fun length() = (endA.row - endB.row).absoluteValue + (endA.col -endB.col).absoluteValue

        fun otherEnd(pos: Pos): Pos {
            require(pos == endA || pos == endB)
            return if (pos == endA) endB else endA
        }

        override fun equals(other: Any?): Boolean {
            if (other !is Segment) return false
            return (endA == other.endA && endB == other.endB) || (endA == other.endB && endB == other.endA)
        }

        fun horizontal() = endA.row == endB.row

        fun getIntersection(other: Segment) : Pos? {
            return if (other.endA.row > endB.row || other.endA.col > endB.col || endA.row > other.endB.row || endA.col > other.endB.col)
                null
            else if (endA.row == endB.row)
                Pos(endA.row, other.endA.col)
            else
                Pos(other.endA.row, endA.col)
        }
    }

    fun createNodeMesh(preparedSegments: List<Segment>) : Map<Pos, Node> {
        val rawNodes = preparedSegments.fold(mutableListOf<Node>(), { acc, seg -> acc.add(Node(seg.endA, mutableListOf(seg))); acc.add(Node(seg.endB, mutableListOf(seg))); acc})
        val groupedNodes = rawNodes.groupBy{it.pos}
        return groupedNodes.map { entry ->
            entry.key to Node(
                entry.key,
                entry.value.flatMap { node -> node.connectedSegments })}.toMap()
    }


    private fun splitAtKeyOrLockOrStartPosition(seg: Segment) : List<Segment> {
        val cursorInc = if (seg.horizontal()) Pos(0,1) else Pos(1,0)
        val nInc = seg.endA.dist(seg.endB)
        var segStart = seg.endA
        var cursor = seg.endA
        val newSegments = mutableListOf<Segment>()
        for (inc in 1..nInc) {
            cursor = cursor.offset(cursorInc)
            if (isNodeChar(data[cursor])) {
                newSegments.add(Segment(segStart, cursor))
                segStart = cursor
            }
        }
        if (segStart != seg.endB) {
            newSegments.add(Segment(segStart, seg.endB))
        }
        return newSegments.toList()
    }

    fun prepareSegments(): List<Segment> {
        val horizontal = rawSegments.filter { p -> p.endA.row == p.endB.row }
        val vertical = rawSegments.filter { p -> p.endA.col == p.endB.col }
        val horizontalAfterJunctionSplit = splitSegments(horizontal, vertical)
        val verticalAfterJunctionSplit = splitSegments(vertical, horizontal)
        val horizontalAfterKeyLockSplit = horizontalAfterJunctionSplit.flatMap{ seg -> splitAtKeyOrLockOrStartPosition(seg)}
        val verticalAfterKeyLockSplit = verticalAfterJunctionSplit.flatMap{seg -> splitAtKeyOrLockOrStartPosition(seg)}
        val segments = mutableListOf<Segment>()
        segments.addAll(horizontalAfterKeyLockSplit)
        segments.addAll(verticalAfterKeyLockSplit)
        return segments
    }

    private fun splitSegments(segments: List<Segment>, other: List<Segment>) : List<Segment> {
        val res: MutableList<Segment> = mutableListOf()
        for (s in segments) {
            val intersections = other.mapNotNull { os -> s.getIntersection(os) }.toMutableList()
            if (intersections.isEmpty() || intersections[0] != s.endA) intersections.add(0, s.endA)
            if (intersections.isEmpty() || intersections.last() != s.endB) intersections.add(s.endB)
            for (ind in 1 until intersections.size) {
                res.add(Segment(intersections[ind - 1], intersections[ind]))
            }
        }
        return res
    }

    fun findShortestPath(nodes: Map<Pos, Node>): Int? {
        val keyPaths: MutableMap<Char, MutableMap<Char, List<PathInfo>>> = mutableMapOf()
        for (startKey in keys) {
            keyPaths[startKey.c]  = mutableMapOf()
            for (goalKey in keys) {
                if (goalKey != startKey) {
                    val startNode = nodes[startKey.pos]!!
                    val goalNode =  nodes[goalKey.pos]!!
                    keyPaths[startKey.c]?.set(goalKey.c, findAllNodePaths(startNode, goalNode, nodes, listOf(), listOf(), setOf()))
                }
            }
        }

        require(startPos != null)
        val startNode = nodes[startPos!!]!!

        val startPaths = keys.map {
                key ->
                val keyNode = nodes[key.pos]!!
                key to findAllNodePaths(startNode, keyNode, nodes, listOf(), listOf(), setOf()) }
            .map{ p -> p.first to p.second.filter { pi -> pi.traversedDoors.isEmpty() }.minBy { pi -> pi.length }?.length }
            .filter{ p -> p.second != null}.sortedBy { p -> p.second }
        return startPaths.map{
            sp ->
            val childLength = findShortestKeyCollectionPath(sp.first.c, setOf(sp.first.c), setOf(sp.first.c.toUpperCase()), sp.second!!, keyPaths)
            if (childLength == null)
                childLength
            else
                childLength + sp.second!!
        }.fold<Int?, Int?>(null, {acc, i ->
            if (acc == null) i else if ( i==null ) acc else acc.coerceAtMost(i)
        })
    }

    fun findShortestKeyCollectionPath(currentKey: Char, foundKeys: Set<Char>, openedDoors: Set<Char>, accLength: Int, keyPaths: Map<Char, Map<Char, List<PathInfo>>>): Int? {
        val remainingKeys = (allKeys - foundKeys) - currentKey
        val memoKey = currentKey to remainingKeys
        if (memoized.containsKey(memoKey)) return memoized[memoKey]
        if (keys.size == foundKeys.size) return 0 // Path found
        val nextPaths = keyPaths[currentKey]?.
            filter{ kv -> !foundKeys.contains(kv.key)}!!
            .map{
                    kv -> kv.key to
                    kv.value.filter{ pi -> (pi.traversedDoors - openedDoors).isEmpty() }.minBy { pi -> pi.length }?.length
            }
            .filter{ kv -> kv.second != null}
            .sortedBy { p -> p.second }
        val result = nextPaths.
            map { kv ->
                val newFoundKeys = foundKeys.toMutableSet()
                val newOpenedDoors = openedDoors.toMutableSet()
                newFoundKeys.add(kv.first)
                newOpenedDoors.add(kv.first.toUpperCase())
                val childLength = findShortestKeyCollectionPath(kv.first, newFoundKeys, newOpenedDoors, accLength  + kv.second!!, keyPaths )
                if (childLength != null) {
                    childLength + kv.second!!
                } else {
                    null
                }
            }
            .fold<Int?, Int?>(null, {acc: Int?, i: Int? ->
                if (acc == null) {
                    i
                } else if (i == null) {
                    acc
                } else {
                    acc.coerceAtMost(i)
                }
            })
        memoized[memoKey] = result
        return result
    }

    fun pathLength(path: List<Segment>) = path.fold(0, {acc, seg -> acc + seg.length()})

    fun findAllNodePaths(cursor: Node, goal: Node, nodes: Map<Pos, Node>, visits: List<Node>, rootPath: List<Segment>, traversedDoors: Set<Char>) : List<PathInfo> {
        if (cursor == goal) return listOf(PathInfo(pathLength(rootPath), traversedDoors))  // Path found
        if (visits.contains(cursor))  return listOf() // Loop
        val newVisits = visits.toMutableList()
        newVisits.add(cursor)
        val newTraversedDoors = traversedDoors.toMutableSet()
        val c = data[cursor.pos]!!
        if (c.isUpperCase()) {
            newTraversedDoors.add(c)
        }
        val nextSegments = cursor.connectedSegments
            .filter{ seg -> rootPath.isEmpty() || seg != rootPath.last()} // Don't go backwards
        if (nextSegments.isEmpty()) return listOf() // Dead end
        return nextSegments.fold(mutableListOf(),
            { acc, seg ->
                val newCursor = nodes[seg.otherEnd(cursor.pos)]!!
                val newRootPath = rootPath.toMutableList()
                newRootPath.add(seg)
                val contrib = findAllNodePaths(newCursor, goal, nodes, newVisits, newRootPath, newTraversedDoors)
                acc.addAll(contrib)
                acc
            })
    }

    private fun isPenetrable(c: Char) = c == '.' || isKeyOrLock(c) || c == '@'
    private fun isKeyOrLock(c: Char) =  c.isUpperCase() || c.isLowerCase()
    private fun isNodeChar(c: Char?) = c != null && (isKeyOrLock(c) || c == '@')

    fun scan() {
        var pathStart: Pos? = null
        for (row in 0..nRows)  {
            for (col in 1..nCols) {
                val pos = Pos(row, col)
                val c = data[pos]
                if (c != null) {
                    if (isPenetrable(c)) {
                        if (pathStart == null) pathStart = pos
                        if (isNodeChar(c)) {
                            if (c.isUpperCase()) {
                                locks.add(Lock(pos, c))
                            } else if (c.isLowerCase()) {
                                keys.add(Key(pos, c))
                                allKeys.add(c)
                            } else if (c == '@') {
                                startPos = pos
                            }
                        }
                    } else {
                        if (pathStart != null) {
                            if (pos.col - pathStart.col > 1) rawSegments.add(Segment(pathStart, Pos(pos.row, pos.col -1 )))
                            pathStart = null
                        }
                    }
                }
            }
        }
        pathStart = null
        for (col in 0..nCols)  {
            for (row in 1..nRows) {
                val pos = Pos(row, col)
                val c = data[pos]
                if (c != null) {
                    if (isPenetrable(c)) {
                        if (pathStart == null) pathStart = pos
                    } else {
                        if (pathStart != null) {
                            if (pos.row - pathStart.row > 1) rawSegments.add(Segment(pathStart, Pos(pos.row - 1, pos.col)))
                            pathStart = null
                        }
                    }
                }
            }
        }
    }
}