[Language: Haskell] Day 17 :-> 2023

module Main where

import Data.List
import Data.List.Split
import qualified Data.Map as M
import Data.PSQueue (Binding ((:->)))
import qualified Data.PSQueue as PQ

type Pos = (Int, Int)
type Prio = (Int, Int, Int)
type State = (Pos, (Int, Dir))

data Dir = U | R | D | L | None deriving (Show, Eq, Ord)

(#+) :: Pos -> Dir -> Pos
(x, y) #+ dir = case dir of
  U -> (x, y - 1)
  R -> (x + 1, y)
  D -> (x, y + 1)
  L -> (x - 1, y)
  None -> (x, y)

opposite :: Dir -> Dir
opposite U = D
opposite D = U
opposite R = L
opposite L = R
opposite None = None

readInt :: String -> Int
readInt = read

main :: IO ()
main = do
  ls <- lines <$> readFile "input.txt"
  let m = M.fromList . concatMap (\(y, line) -> zipWith (\x c -> ((x, y), readInt [c])) [0 ..] line) . zip [0 ..] $ ls
  let lines = part1 m
  putStrLn $ "Part 1: " ++ (show . part1 $ m)
  putStrLn $ "Part 2: " ++ (show . part2 $ m)

shortestPath :: M.Map Pos Int -> Pos -> (State -> Bool) -> (State -> [Dir]) -> (Pos -> Int) -> Int
shortestPath m start isFinished getNextDirs heuristic =
    let initialH = heuristic start
    in sp (PQ.singleton (start, (0, None)) (initialH, 0, initialH)) M.empty
  where
    sp :: PQ.PSQ State Prio -> M.Map State Int -> Int
    sp open closed =
      case PQ.minView open of
        Nothing -> error "Couldn't find a way. :frown:"
        Just (state :-> (f, g, h), restOpen) ->
            if isFinished state
            then g
            else let neighbors = map (posDirToBinding g) . filter (`M.notMember` closed) . filter ((`M.member` m) . fst) $ getNext state
                     newOpen = foldl' updateQueue restOpen neighbors
                 in sp newOpen (M.insert state g closed)

    getNext :: State -> [State]
    getNext state@(pos, (n, dir)) = map f . getNextDirs $ state
      where f :: Dir -> (Pos, (Int, Dir))
            f d | d == dir = (pos #+ d, (n + 1, d)) | otherwise = (pos #+ d, (1, d))

    updateQueue :: PQ.PSQ State Prio -> PQ.Binding State Prio -> PQ.PSQ State Prio
    updateQueue pq (key :-> prio) = PQ.alter (updateBinding prio) key pq

    updateBinding :: Prio -> Maybe Prio -> Maybe Prio
    updateBinding prio Nothing = Just prio
    updateBinding prio@(f, g, h) oldPrio@(Just (oldF, oldG, oldH)) = if g < oldG then Just prio else oldPrio

    posDirToBinding :: Int -> State -> PQ.Binding State Prio
    posDirToBinding g pd@(pos, (n, dir)) = pd :-> (g + cost + h, g + cost, h)
        where cost = m M.! pos
              h = heuristic pos

manhattan :: Pos -> Pos -> Int
manhattan (x1, y1) (x2, y2) = abs (x2 - x1) + abs (y2 - y1)

part1 :: M.Map Pos Int -> Int
part1 m = shortestPath m (0, 0) finishRule getNextDirs heuristic
  where finish :: Pos
        (finish, _) = M.findMax m

        heuristic :: Pos -> Int
        heuristic = manhattan finish

        finishRule :: State -> Bool
        finishRule (pos, _) = pos == finish

        getNextDirs :: State -> [Dir]
        getNextDirs (_, (n, dir))
          | n >= 3     = [U, R, D, L] \\ [opposite dir, dir]
          | otherwise  = [U, R, D, L] \\ [opposite dir]

part2 :: M.Map Pos Int -> Int
part2 m = shortestPath m (0, 0) finishRule getNextDirs heuristic
  where finish :: Pos
        (finish, _) = M.findMax m

        heuristic :: Pos -> Int
        heuristic = manhattan finish

        finishRule :: State -> Bool
        finishRule (pos, (n, _)) = pos == finish && n >= 4

        getNextDirs :: State -> [Dir]
        getNextDirs state@(_, (n, dir))
          | n == 0      = [U, R, D, L]
          | n < 4       = [dir]
          | n >= 10     = [U, R, D, L] \\ [opposite dir, dir]
          | otherwise   = [U, R, D, L] \\ [opposite dir]