Day 23: A Long Walk

{-# LANGUAGE TupleSections #-}
module Main (main) where

import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Map (Map, (!))
import Data.Set (Set)
import Data.Maybe (fromJust, mapMaybe)
import Data.List (maximumBy)
import Data.Tuple (swap)
import Data.Function (on)

main :: IO ()
main = do
    input <- readFile "./input.txt"
    putStrLn $ "Part 1: " ++ show (solve1 $ parseInput input)
    putStrLn $ "Part 2: " ++ show (solve2 $ parseInput input)

data MazeTile = Wall | Open | L | R | U | D deriving (Show, Eq)
type Maze = Map Coord MazeTile
type Coord = (Int, Int)

parseInput :: String -> (Maze, Coord, Coord)
parseInput str = (maze, startPos, endPos) where
    rows = length $ lines str
    cols = length $ head (lines str)
    startPos = (0, 1)
    endPos = (rows - 1, cols - 2)
    maze = Map.fromList $ concat [[((r, c), parseTile x) | (c, x) <- zip [0..] row] | (r, row) <- zip [0..] (lines str)]

parseTile :: Char -> MazeTile
parseTile c = case c of
    '#' -> Wall
    '.' -> Open
    'v' -> D
    '^' -> U
    '<' -> L
    '>' -> R
    _   -> error "Invalid Tile"

solve1 :: (Maze, Coord, Coord) -> Int
solve1 (maze, startPos, endPos) = length . fromJust $ longestPath startPos Set.empty where
    longestPath :: Coord -> Set Coord -> Maybe [Coord]
    longestPath currPos visited
        | currPos == endPos = Just []
        | otherwise = if null pathOptions then Nothing else Just bestPath where
            nextPositions = filter (`Set.notMember` visited) (getNeighbors maze currPos False)
            pathOptions = mapMaybe (\pos -> longestPath pos (Set.insert pos visited)) nextPositions
            bestPath = currPos : maximumBy (\a b -> compare (length a) (length b)) pathOptions

getNeighbors :: Maze -> Coord -> Bool -> [Coord]
getNeighbors maze (r, c) passThrough = validNeighbors where
    neighbors = if passThrough then
        if maze ! (r, c) /= Wall then [(r+1, c), (r-1, c), (r, c+1), (r, c-1)] else []
    else
        case maze ! (r, c) of
            Open -> [(r+1, c), (r-1, c), (r, c+1), (r, c-1)]
            L    -> [(r, c-1)]
            R    -> [(r, c+1)]
            U    -> [(r-1, c)]
            D    -> [(r+1, c)]
            Wall -> error "Stuck in a wall"
    validNeighbors = filter (\n -> Map.member n maze && maze ! n /= Wall) neighbors

type JunctionID = Char
type JunctionGraph = Map JunctionID (Map JunctionID Int)

idList :: [JunctionID]
idList = ['A'..'Z'] ++ ['a'..'z']

findJunctions :: Maze -> [Coord]
findJunctions maze = junctions where
    junctions = filter (\pos -> length (getNeighbors maze pos True) > 2) (Map.keys maze)

buildJunctionGraph :: Maze -> [(JunctionID, Coord)] -> JunctionGraph
buildJunctionGraph maze junctions = junctionGraph where
    idLookup = Map.fromList $ map swap junctions
    junctionGraph = Map.fromList $ [(startID, Map.fromList $ findConnectedJunctions startCoord) | (startID, startCoord) <- junctions]

    findConnectedJunctions :: Coord -> [(JunctionID, Int)]
    findConnectedJunctions startingPos = solve (map (, 1) $ getNeighbors maze startingPos True) (Set.singleton startingPos) where
        solve :: [(Coord, Int)] -> Set Coord -> [(JunctionID, Int)]
        solve [] _ = []
        solve ((x, d):xs) visited
            | Map.member x idLookup = (idLookup ! x, d) : solve xs visited
            | otherwise = solve (neighbors ++ xs) (Set.insert x visited) where
                neighbors = map (,d+1) $ filter (`Set.notMember` visited) (getNeighbors maze x True)

solve2 :: (Maze, Coord, Coord) -> Int
solve2 (maze, startPos, endPos) = longestPath (head idList) where
    junctions = zip idList (startPos : findJunctions maze ++ [endPos])
    junctionGraph = buildJunctionGraph maze junctions
    endID = idList !! (length junctions - 1)
    longestPath :: JunctionID -> Int
    longestPath startID = snd . fromJust $ solve startID Set.empty where
        solve :: JunctionID -> Set JunctionID -> Maybe ([JunctionID], Int)
        solve currID visited
            | currID == endID = Just ([endID], 0)
            | otherwise = if null pathOptions' then Nothing else Just bestPath where
                nextPositions = filter (`Set.notMember` visited) (Map.keys $ junctionGraph ! currID)
                pathOptions = mapMaybe (\jid -> solve jid (Set.insert currID visited)) nextPositions
                pathOptions' = map (\(p, d) -> (currID : p, d + (junctionGraph ! currID ! head p))) pathOptions
                bestPath = maximumBy (compare `on` snd) pathOptions'