--- /dev/null
+-- Writeup at https://work.njae.me.uk/2021/12/16/advent-of-code-2021-day-15/
+
+import Debug.Trace
+
+
+import qualified Data.PQueue.Prio.Min as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import Data.Sequence ((<|), (|>), (><)) --, ViewR( (:>) ), ViewL( (:<) ))
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+import Data.Foldable (foldl', sum) -- (toList, foldr', foldl', all)
+-- import Data.Char
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<))
+import Data.Maybe -- (fromMaybe)
+-- import Linear (V2(..), (^+^)) --, (^-^), (*^), (^*))
+import Linear hiding (trace)
+
+
+
+pattern Empty <- (Q.viewl -> Q.EmptyL) where Empty = Q.empty
+pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<) = (Q.<|)
+pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>) = (Q.|>)
+
+
+data Amphipod = A | B | C | D deriving (Show, Read, Eq, Ord, Enum)
+
+type Coord = V2 Int -- r, c
+_r :: Lens' (V2 Int) Int
+_r = _x
+_c :: Lens' (V2 Int) Int
+_c = _y
+
+data Step = Step
+ { _destination :: Coord
+ , _distance :: Int
+ , _transits :: S.Set Coord
+ , _entryRequirement :: Maybe Amphipod
+ } deriving (Show, Eq, Ord)
+makeLenses ''Step
+
+type Steps = M.Map Coord (S.Set Step)
+
+data Burrow = Burrow
+ { _possibleSteps :: Steps
+ , _roomColumns :: M.Map Amphipod Int
+ , _hallRow :: Int
+ } deriving (Show, Eq)
+makeLenses ''Burrow
+
+type BurrowContext = Reader Burrow
+
+type MoveState = M.Map Coord Amphipod
+
+data AppliedMove = AppliedMove
+ { _afterMove :: MoveState
+ , _appliedStep :: Step
+ }
+ deriving (Show, Eq, Ord)
+makeLenses ''AppliedMove
+
+data Agendum =
+ Agendum { _current :: MoveState
+ , _trail :: Q.Seq MoveState
+ , _trailCost :: Int
+ , _cost :: Int
+ } deriving (Show, Eq)
+makeLenses ''Agendum
+
+type Agenda = P.MinPQueue Int Agendum
+
+type ExploredStates = S.Set MoveState
+
+
+
+main :: IO ()
+main =
+ do text <- readFile "data/advent23.txt"
+ -- let (burrow, initState) = mkBurrow text
+ -- print burrow
+ -- print initState
+ print $ part1 text
+ print $ part2 text
+
+
+-- part1 :: Burrow -> MoveState -> Int
+part1 text = maybe 0 _cost result
+ where
+ (burrow, initState) = mkBurrow text
+ result = runReader (searchBurrow initState) burrow
+
+part2 text = maybe 0 _cost result
+ where
+ rows = lines text
+ extraRows = [(" #D#C#B#A# " :: String), (" #D#B#A#C# " :: String)]
+ modifiedRows = (take 3 rows) ++ extraRows ++ (drop 3 rows)
+ modifiedText = unlines modifiedRows
+ (burrow, initState) = mkBurrow modifiedText
+ result = runReader (searchBurrow initState) burrow
+
+
+searchBurrow :: MoveState -> BurrowContext (Maybe Agendum)
+searchBurrow initState =
+ do agenda <- initAgenda initState
+ aStar agenda S.empty
+
+initAgenda :: MoveState -> BurrowContext Agenda
+initAgenda initState =
+ do c <- estimateCost initState
+ return $ P.singleton c Agendum { _current = initState
+ , _trail = Q.empty, _trailCost = 0, _cost = c}
+
+
+aStar :: Agenda -> ExploredStates -> BurrowContext (Maybe Agendum)
+aStar agenda closed
+ -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+ -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
+ | P.null agenda = return Nothing
+ | otherwise =
+ do let (_, currentAgendum) = P.findMin agenda
+ let reached = currentAgendum ^. current
+ nexts <- candidates currentAgendum closed
+ let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
+ reachedGoal <- isGoal reached
+ if reachedGoal
+ then return (Just currentAgendum)
+ else if reached `S.member` closed
+ then aStar (P.deleteMin agenda) closed
+ else aStar newAgenda (S.insert reached closed)
+
+
+candidates :: Agendum -> ExploredStates -> BurrowContext (Q.Seq Agendum)
+candidates agendum closed =
+ do let candidate = agendum ^. current
+ let previous = agendum ^. trail
+ let prevCost = agendum ^. trailCost
+ succs <- successors candidate
+ let nonloops = S.filter (\s -> (s ^. afterMove) `S.notMember` closed) succs
+ let nonloopsQ = Q.fromList $ S.toList nonloops
+ mapM (makeAgendum previous prevCost) nonloopsQ
+
+
+makeAgendum :: Q.Seq MoveState -> Int -> AppliedMove -> BurrowContext Agendum
+makeAgendum previous prevCost newPosition =
+ do predicted <- estimateCost (newPosition ^. afterMove)
+ -- grid <- asks _grid
+ let newTrail = previous |> (newPosition ^. afterMove)
+ let newPositionCost = stepCost newPosition
+ let incurred = prevCost + newPositionCost
+ return Agendum { _current = newPosition ^. afterMove
+ , _trail = newTrail
+ , _trailCost = incurred
+ , _cost = incurred + predicted
+ }
+
+-- class (Eq s, Ord s, Show s) => SearchState s where
+-- emptySearchState :: MoveState
+-- successors :: MoveState -> BurrowContext (Q.Seq MoveState)
+-- estimateCost :: MoveState -> BurrowContext Int
+-- isGoal :: MoveState -> BurrowContext Bool
+-- entryCost :: MoveState -> BurrowContext Int
+
+
+-- instance SearchState Position where
+
+-- emptySearchState = Position (V2 0 0)
+
+successors :: MoveState -> BurrowContext (S.Set AppliedMove)
+successors moveState =
+ do steps <- asks _possibleSteps
+ let succs = M.foldrWithKey' (legalSteps steps moveState) S.empty moveState
+ return succs
+
+legalSteps :: Steps -> MoveState -> Coord -> Amphipod -> S.Set AppliedMove -> S.Set AppliedMove
+legalSteps steps state here amphipod acc = S.union appliedSteps acc
+ where allSteps = steps ! here
+ freeSteps = S.filter freeSpaces allSteps
+ freeSpaces st = S.null $ S.intersection (M.keysSet state) (st ^. transits)
+ validTargetSteps = S.filter (\st -> fromMaybe amphipod (st ^. entryRequirement) == amphipod) freeSteps
+ openRoomSteps = S.filter (openRoom state) validTargetSteps
+ appliedSteps = S.map (\s -> AppliedMove
+ { _afterMove = (applyStep state here s)
+ , _appliedStep = s
+ }
+ ) openRoomSteps
+
+openRoom :: MoveState -> Step -> Bool
+openRoom state step
+ | isNothing e = True
+ | otherwise = M.null roomBlockers
+ where e = step ^. entryRequirement
+ je = fromJust e
+ tc = step ^. destination . _c
+ roomBlockers = M.filterWithKey (\(V2 _ ac) a -> a /= je && ac == tc) state
+
+applyStep :: MoveState -> Coord -> Step -> MoveState
+applyStep moveState here step = moveState''
+ where moveState' = M.delete here moveState
+ moveState'' = M.insert (step ^. destination) (moveState ! here) moveState'
+
+singleStepCost :: Amphipod -> Int
+singleStepCost A = 1
+singleStepCost B = 10
+singleStepCost C = 100
+singleStepCost D = 1000
+
+estimateCost :: MoveState -> BurrowContext Int
+estimateCost state =
+ do rCols <- asks _roomColumns
+ hRow <- asks _hallRow
+ let amphipodCosts = M.mapWithKey (estimateACost rCols hRow) state
+ return $ sum $ M.elems amphipodCosts
+
+estimateACost :: M.Map Amphipod Int -> Int -> Coord -> Amphipod -> Int
+estimateACost rCols hRow (V2 r c) amphipod = (singleStepCost amphipod) * dist
+ where targetCol = rCols ! amphipod
+ dist = if c == targetCol
+ then 0
+ else (r - hRow) + (abs (c - targetCol)) + 1
+
+stepCost :: AppliedMove -> Int
+stepCost aStep = (singleStepCost amphipod) * (S.size $ aStep ^. appliedStep . transits)
+ where dest = aStep ^. appliedStep . destination
+ amphipod = (aStep ^. afterMove) ! dest
+
+isGoal :: MoveState -> BurrowContext Bool
+isGoal state =
+ do rCols <- asks _roomColumns
+ let misplaced = M.filterWithKey (inWrongRoom rCols) state
+ return $ M.null misplaced
+
+inWrongRoom rCols (V2 _ c) amphipod = c /= rightCol
+ where rightCol = rCols ! amphipod
+
+
+------------------------------
+
+
+mkBurrow :: String -> (Burrow, MoveState)
+-- mkBurrow :: String -> ((S.Set Coord, M.Map Coord Amphipod), MoveState)
+mkBurrow text = (burrow, initState) -- (burrow, initState)
+ where rows = lines text
+ hall = mkHall (rows!!1)
+ rooms = mkRooms $ drop 2 rows
+ roomCols = S.map (^. _c) $ M.keysSet rooms
+ hall' = S.filter ((`S.notMember` roomCols) . (^. _c)) hall
+ routes = mkRoutes hall' rooms
+ roomColMap = M.fromList $ zip [A .. D] $ S.toAscList roomCols
+ burrow = Burrow { _possibleSteps = routes, _roomColumns = roomColMap, _hallRow = 1}
+ initState = mkInitialState rows
+
+
+mkHall :: String -> S.Set Coord
+mkHall text = S.fromList hallCoords
+ where hallCols = filter ((/= '#') . snd) $ zip [0..] text
+ hallCoords = map ((V2 1) . fst) hallCols
+
+mkRooms :: [String] -> M.Map Coord Amphipod
+mkRooms text = M.unions rooms
+ where rooms = map mkRoom $ zip [2..] text
+
+mkRoom :: (Int, String) -> M.Map Coord Amphipod
+mkRoom (r, text) = M.fromList roomCoords
+ where roomCols = filter ((`elem` ("ABCD." :: String)) . snd) $ zip [0..] text
+ roomCoords = zip (map ((V2 r) . fst) roomCols) [A .. D]
+
+-- invertRooms rooms = M.fromList [(a, M.keysSet $ M.filter (== a) rooms) | a <- [A .. D]]
+
+mkRoutes :: S.Set Coord -> M.Map Coord Amphipod -> Steps
+mkRoutes halls rooms = M.unionsWith (S.union) [hallRoutes, roomHallRoutes, roomRoomRoutes]
+ where hallRoutes = S.foldr' (mkHallRoute rooms) M.empty halls
+ roomHallRoutes = S.foldr' (mkRoomHallRoute halls) M.empty (M.keysSet rooms)
+ roomRoomRoutes = S.foldr' (mkRoomRoomRoute hallRow rooms) M.empty (M.keysSet rooms)
+ hallRow = (S.findMin halls) ^. _r
+
+mkHallRoute :: M.Map Coord Amphipod -> Coord -> Steps -> Steps
+-- mkHallRoute rooms here routes | trace ("mkHR " ++ (show here) ++ " " ++ (show routes)) False = undefined
+mkHallRoute rooms here routes = M.foldrWithKey' (mkHallRoute1 here) routes rooms
+
+mkHallRoute1 :: Coord -> Coord -> Amphipod -> Steps -> Steps
+-- mkHallRoute1 here there entry routes | trace ("mkHR1 " ++ (show here) ++ " " ++ (show there) ++ (show routes)) False = undefined
+mkHallRoute1 here@(V2 hr hc) there@(V2 tr tc) entry routes = M.insert here (S.insert step existingRoutes) routes
+ -- | trace ("mkHR1 " ++ (show here) ++ " " ++ (show there) ++ (show routes) ++ " > " ++ show res) False = undefined
+ -- | otherwise = res
+ where step = Step { _destination = there
+ , _distance = (S.size transits)
+ , _transits = transits
+ , _entryRequirement = Just entry
+ }
+ cMin = min hc tc
+ cMax = max hc tc
+ transits = S.delete here $ S.fromList $ [V2 hr c | c <- [cMin..cMax]] ++ [V2 r tc | r <- [hr..tr]]
+ existingRoutes = M.findWithDefault S.empty here routes
+ -- res = M.insert here (S.insert step existingRoutes) routes
+
+mkRoomHallRoute :: S.Set Coord -> Coord -> Steps -> Steps
+mkRoomHallRoute halls here routes = S.foldr' (mkRoomHallRoute1 here) routes halls
+
+mkRoomHallRoute1 :: Coord -> Coord -> Steps -> Steps
+mkRoomHallRoute1 here@(V2 hr hc) there@(V2 tr tc) routes = M.insert here (S.insert step existingRoutes) routes
+ where step = Step { _destination = there
+ , _distance = (S.size transits)
+ , _transits = transits
+ , _entryRequirement = Nothing
+ }
+ cMin = min hc tc
+ cMax = max hc tc
+ transits = S.delete here $ S.fromList $ [V2 r hc | r <- [tr..hr]] ++ [V2 tr c | c <- [cMin..cMax]]
+ existingRoutes = M.findWithDefault S.empty here routes
+
+mkRoomRoomRoute :: Int -> M.Map Coord Amphipod -> Coord -> Steps -> Steps
+mkRoomRoomRoute hallRow rooms here routes = M.foldrWithKey' (mkRoomRoomRoute1 hallRow here) routes rooms
+
+mkRoomRoomRoute1 :: Int -> Coord -> Coord -> Amphipod -> Steps -> Steps
+-- mkRoomRoomRoute1 _hallRow here there entry routes | trace ("mkRR1 " ++ (show here) ++ " " ++ (show there) ++ (show routes)) False = undefined
+mkRoomRoomRoute1 hallRow here@(V2 hr hc) there@(V2 tr tc) entry routes
+ | here == there = routes
+ | otherwise = M.insert here (S.insert step existingRoutes) routes
+ where step = Step { _destination = there
+ , _distance = (S.size transits)
+ , _transits = transits
+ , _entryRequirement = Just entry
+ }
+ cMin = min hc tc
+ cMax = max hc tc
+ transitUp = S.fromList [V2 r hc | r <- [hallRow..hr]]
+ transitAcross = S.fromList [V2 hallRow c | c <- [cMin..cMax]]
+ transitDown = S.fromList [V2 r tc | r <- [hallRow..tr]]
+ transits = S.delete here $ S.unions [transitUp, transitAcross, transitDown]
+ existingRoutes = M.findWithDefault S.empty here routes
+
+
+mkInitialState :: [String] -> MoveState
+mkInitialState rows =
+ M.fromList [ (V2 r c, read [(rows!!r)!!c])
+ | r <- [0..maxR], c <- [0..maxC]
+ , isAmphipod ((rows!!r)!!c)
+ ]
+ where maxR = length rows - 1
+ maxC = (length $ head rows) - 1
+ isAmphipod c = c `elem` ("ABCD" :: String)
+