aoc24e.hs

{-# LANGUAGE TupleSections #-}
{-# OPTIONS_GHC -Wall #-}

import Control.Monad ( guard, when )
import Data.List (sortBy)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import System.Environment (getArgs)

-- import Debug.Trace (trace)

neigh :: Int -> Int -> M.Map (Int, Int) Char -> ((Int, Int), Int) -> [((Int, Int), Int)]
neigh pitchHeight pitchWidth grid ((row, col), t) = do
  (row', col') <- [(row + r, col + c) | (r, c) <- [(0, 0), (-1, 0), (1, 0), (0, -1), (0, 1)]]
  guard $ '#' /= M.findWithDefault '#' (row, col) grid
  when ((1 <= row) && (row <= pitchHeight)) $ do
    guard $ '<' /= grid M.! (row', 1 + ((col' + t) `mod` pitchWidth))
    guard $ '>' /= grid M.! (row', 1 + ((col' - t - 2) `mod` pitchWidth))
    guard $ '^' /= grid M.! (1 + ((row' + t) `mod` pitchHeight), col')
    guard $ 'v' /= grid M.! (1 + ((row' - t - 2) `mod` pitchHeight), col')
  pure ((row', col'), t + 1)

estimate :: (Int, Int) -> ((Int, Int), a) -> Int
estimate goal ((row, col), _) = abs (fst goal - row) + abs (snd goal - col)

-- Apply technique mentioned in https://www.reddit.com/r/haskell/comments/p2tm2n/comment/h8muqzd/,
-- only modified slightly as all steps are size 1

data Worky a = Done a | Worky (Worky a) | Fail deriving (Show, Eq)

instance Semigroup (Worky a) where
  (Done x) <> _ = Done x
  _ <> (Done x) = Done x
  Fail <> x = x
  x <> Fail = x
  (Worky a) <> (Worky b) = Worky (a <> b)

instance Monoid (Worky a) where
  mempty = Fail

finish :: Worky a -> Maybe a
finish (Done x) = Just x
finish (Worky w) = finish w
finish Fail = Nothing

type State = ((Int, Int), Int)

workGoal :: (State -> [State]) -> (State -> Int) -> (State -> Bool) -> State -> Worky State
workGoal nbF estF goalF = go . S.singleton
 where
  go :: S.Set State -> Worky State
  go iStates =
   let dones = S.toList $ S.filter goalF iStates
   in
      case dones of
        (x:_) -> Done x
        _ -> Worky $ go $ S.fromList $ concatMap nbF $ S.toList iStates 

main :: IO ()
main = do
  args <- getArgs
  let filename =
        if null args
          then "aoc24.in"
          else head args
  s <- lines <$> readFile filename
  let grid = M.fromList [((row, col), ch) | (row, line) <- zip [0 ..] s, (col, ch) <- zip [0 ..] line]
  let start = head $ map ((0,) . fst) $ filter (\(_, ch) -> ch == '.') $ zip [0 ..] (head s)
  let goal = head $ map ((length s - 1,) . fst) $ filter (\(_, ch) -> ch == '.') $ zip [0 ..] (last s)
  let pitchWidth = length (head s) - 2
  let pitchHeight = length s - 2

  let Just part1 = finish  $ workGoal (neigh pitchHeight pitchWidth grid) (estimate goal) ((== goal) . fst) (start, 0)
  print $ snd part1

  let Just part2a = finish  $ workGoal (neigh pitchHeight pitchWidth grid) (estimate start) ((== start) . fst) part1
  let Just part2b = finish  $ workGoal (neigh pitchHeight pitchWidth grid) (estimate goal) ((== goal) . fst) part2a
  print $ snd part2b