Done day 23
[advent-of-code-21.git] / advent23 / Main.hs
diff --git a/advent23/Main.hs b/advent23/Main.hs
new file mode 100644 (file)
index 0000000..2666769
--- /dev/null
@@ -0,0 +1,343 @@
+-- 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)
+