advent20/src/advent20.hs

   1 import Debug.Trace
   2
   3 -- import qualified Data.Text.IO as TIO
   4
   5 -- import qualified Data.Map.Strict as M
   6 -- import Data.Map.Strict ((!))
   7 import qualified Data.PQueue.Prio.Min as P
   8 import qualified Data.Set as S
   9 import qualified Data.Sequence as Q
  10 import Data.Sequence ((<|)) -- , (|>), (><))
  11 import Data.Foldable (foldl', sum) -- (toList, foldr', foldl', all)
  12 import Data.Char
  13 import Control.Monad.Reader
  14 import Control.Lens hiding ((<|), (|>))
  15 import Data.Maybe (fromMaybe)
  16
  17
  18
  19 type Position = (Int, Int) -- r, c
  20
  21
  22 data Location = Inner | Outer deriving (Eq, Ord, Show)
  23 data Portal = Portal { _label :: String
  24                      , _position :: Position
  25                      , _location :: Location
  26                      } deriving (Eq, Ord, Show)
  27 makeLenses ''Portal
  28
  29 type Portals = S.Set Portal
  30
  31
  32 type ExpandedMaze = S.Set Position
  33 data MazeComplex = MazeComplex
  34     { _mazeE :: ExpandedMaze
  35     , _portalsE :: Portals
  36     , _portalLocs :: S.Set Position
  37     } deriving (Eq, Ord, Show)
  38 makeLenses ''MazeComplex
  39
  40 type EdgeConnects = (Portal, Portal)
  41 data EdgeType = Walk | Teleport deriving (Eq, Ord, Show)
  42
  43 data Edge = Edge { _connects :: EdgeConnects
  44                  , _edgeType :: EdgeType
  45                  , _distance :: Int
  46                  } deriving (Eq, Ord, Show)
  47 makeLenses ''Edge
  48
  49 type Edges = S.Set Edge
  50
  51 -- type Maze = S.Set Edge
  52 data Maze = Maze { _maze :: Edges
  53                  , _costPerLevel :: Int
  54                  , _costToFinish :: Int
  55                  } deriving (Eq, Ord, Show)
  56 -- makeLenses ''Maze
  57
  58 type MazeContext = Reader Maze
  59
  60 class (Eq s, Ord s, Show s) => SearchState s where
  61     successors :: s -> MazeContext (Q.Seq (s, Edge))
  62     estimateCost :: s -> MazeContext Int
  63     emptySearchState :: Portal -> s
  64     isGoal :: s -> MazeContext Bool
  65
  66 data LevelledSearchState = LevelledSearchState
  67     { _portalS :: Portal
  68     , _levelS :: Int
  69     } deriving (Eq, Ord, Show)
  70 makeLenses ''LevelledSearchState
  71
  72
  73 type ExploredStates s = S.Set s
  74
  75 data Agendum s =
  76     Agendum { _current :: s
  77             , _trail :: Q.Seq Edge
  78             , _cost :: Int
  79             } deriving (Show, Eq)
  80 makeLenses ''Agendum
  81
  82 type Agenda s = P.MinPQueue Int (Agendum s)
  83
  84
  85 main :: IO ()
  86 main = do
  87         maze <- setup
  88         -- print maze
  89         putStrLn $ showContracted maze
  90         print $ part1 maze
  91         print $ part2 maze
  92
  93
  94 setup = do
  95         text <- readFile "data/advent20.txt"
  96         let mc = buildComplex text
  97         -- print mc
  98         -- print maze
  99         return $ contractMaze mc
 100
 101
 102 part1 :: Maze -> Int
 103 -- part1 :: Maze -> Maybe (Agendum Portal)
 104 part1 maze = maybe 0 _cost result
 105     where result = runReader searchMaze maze :: Maybe (Agendum Portal)
 106
 107 part2 :: Maze -> Int
 108 -- part2 :: Maze -> Maybe (Agendum LevelledSearchState)
 109 part2 maze = maybe 0 _cost result
 110     where result = runReader searchMaze maze :: Maybe (Agendum LevelledSearchState)
 111
 112
 113 buildComplex :: String -> MazeComplex
 114 buildComplex text = mc & portalLocs .~ pLocs & portalsE .~ portals'
 115     where mc = foldl' (buildMazeRow rows) mc0 [0..l]
 116           mc0 = MazeComplex {_mazeE = S.empty, _portalsE = S.empty, _portalLocs = S.empty}
 117           rows = lines text
 118           l = length rows - 1
 119           minR = 2
 120           maxR = l - 2
 121           minC = 2
 122           maxC = length (rows!!2) - 3
 123           pLocs = S.map _position (mc ^. portalsE)
 124           portals = mc ^. portalsE
 125           portals' = S.map (classifiyPortal minR maxR minC maxC) portals
 126
 127 classifiyPortal :: Int -> Int -> Int -> Int -> Portal -> Portal
 128 classifiyPortal minR maxR minC maxC portal = portal & location .~ loc
 129     where (r, c) = portal ^. position
 130           loc = if (r == minR) || (r == maxR) || (c == minC) || (c == maxC)
 131                 then Outer
 132                 else Inner
 133
 134 buildMazeRow :: [String] -> MazeComplex -> Int -> MazeComplex
 135 buildMazeRow rows mc r = foldl' (buildMazeCell rows r) mc [0..l]
 136     where l = length (rows!!r) - 1
 137
 138
 139 buildMazeCell :: [String] -> Int -> MazeComplex -> Int -> MazeComplex
 140 buildMazeCell rows r mc c
 141     | char == '.' = mc'
 142     | isUpper char = mc & portalsE .~ portals'
 143     | otherwise = mc
 144     where char = (rows!!r)!!c
 145           mc' = mc & mazeE %~ (S.insert here)
 146           here = (r, c)
 147           portals' = makePortal (mc ^. portalsE) rows char r c
 148
 149
 150 makePortal portals rows hc r c
 151     | isUpper rc = if pr == '.'
 152                    then S.insert (Portal { _label = [hc, rc], _position = (r, c + 2), _location = Outer } ) portals
 153                    else S.insert (Portal { _label = [hc, rc], _position = (r, c - 1), _location = Outer } ) portals
 154     | isUpper dc = if pd == '.'
 155                    then S.insert (Portal { _label = [hc, dc], _position = (r + 2, c), _location = Outer } ) portals
 156                    else S.insert (Portal { _label = [hc, dc], _position = (r - 1, c), _location = Outer } ) portals
 157     | otherwise = portals
 158     where -- lc = charAt rows r (c - 1)
 159           rc = charAt rows r (c + 1)
 160           -- uc = charAt rows (r - 1) c
 161           dc = charAt rows (r + 1) c
 162           -- pu = charAt rows (r - 1) c
 163           pd = charAt rows (r + 2) c
 164           -- pl = charAt rows r (c - 1)
 165           pr = charAt rows r (c + 2)
 166
 167
 168 charAt :: [String] -> Int -> Int -> Char
 169 charAt rows r c = atDef ' ' (atDef "" rows r) c
 170
 171 atDef :: a -> [a] -> Int -> a
 172 atDef x xs i = fromMaybe x $ atMaybe xs i
 173 -- atDef x = (fromMaybe x) . atMaybe
 174
 175 atMaybe :: [a] -> Int -> Maybe a
 176 atMaybe xs i
 177     | i < 0 = Nothing
 178     | i >= (length xs) = Nothing
 179     | otherwise = Just (xs!!i)
 180
 181
 182 contractMaze :: MazeComplex -> Maze
 183 contractMaze expanded = Maze
 184     { _maze = S.union mazeW mazeP
 185     , _costPerLevel = cpl
 186     , _costToFinish = ctf
 187     }
 188     where starts = expanded ^. portalsE
 189           mazeP = S.foldr (contractFrom expanded) S.empty starts
 190           mazeW = S.foldr (addWarp starts) S.empty starts
 191           cpl = minimum $ map (^. distance) $ S.toList $ S.filter (\e -> e ^. edgeType == Walk) mazeP
 192           ctf = minimum $ map (^. distance) $ S.toList $ S.filter (edgeTouches term) mazeP
 193           term = S.findMin $ S.filter (\p -> p ^. label == "ZZ") starts
 194
 195 contractFrom :: MazeComplex -> Portal -> Edges -> Edges
 196 contractFrom expanded start maze0 = S.union maze0 reachables
 197     where startPos = start ^. position
 198           reachables = reachableFrom [(startPos, 0)] S.empty expanded' start
 199           expanded' = expanded & portalsE %~ (S.delete start)
 200                                & portalLocs %~ (S.delete startPos)
 201                                -- & mazeE %~ (S.delete startPos)
 202
 203 reachableFrom :: [(Position, Int)] -> (S.Set Position) -> MazeComplex -> Portal -> Edges
 204 reachableFrom [] _closed _expanded _start = S.empty
 205 reachableFrom ((here, dist):boundary) closed expanded start
 206     | here `S.member` closed = reachableFrom boundary closed expanded start
 207     | here `S.member` ps = S.insert edge $ reachableFrom boundary closed' expanded start
 208     | otherwise = reachableFrom boundary' closed' expanded start
 209     where nbrs0 = S.intersection (expanded ^. mazeE) $ possibleNeighbours here
 210           nbrs = S.difference nbrs0 closed
 211           closed' = S.insert here closed
 212           ps = expanded ^. portalLocs
 213           other = S.findMin $ S.filter (\p -> p ^. position == here) $ expanded ^. portalsE
 214           edge = Edge { _connects = mkConnection start other, _edgeType = Walk, _distance = dist }
 215           neighbours = S.map (\n -> (n, dist + 1)) nbrs
 216           boundary' = boundary ++ (S.toAscList neighbours)
 217
 218
 219 addWarp :: Portals -> Portal -> Edges -> Edges
 220 addWarp portals portal warps
 221     | S.null others = warps
 222     | otherwise = S.insert warp warps
 223     where others = S.filter (portalsConnect lab pos) portals
 224           lab = portal ^. label
 225           pos = portal ^. position
 226           other = S.findMin others
 227           warp = Edge {_connects = mkConnection portal other, _edgeType = Teleport, _distance = 1}
 228
 229 portalsConnect :: String -> Position -> Portal -> Bool
 230 portalsConnect lab pos portal = (pLabel == lab) && (pPos /= pos)
 231     where pLabel = portal ^. label
 232           pPos = portal ^. position
 233
 234
 235 mkConnection :: Portal -> Portal -> EdgeConnects
 236 mkConnection a b = if a < b then (a, b) else (b, a)
 237
 238
 239 edgeTouches :: Portal -> Edge -> Bool
 240 edgeTouches p e
 241     | p == a = True
 242     | p == b = True
 243     | otherwise = False
 244     where (a, b) = e ^. connects
 245
 246 edgeOther :: Portal -> Edge -> Portal
 247 edgeOther x e
 248     | x == a = b
 249     | otherwise = a
 250     where (a, b) = e ^. connects
 251
 252 possibleNeighbours :: Position -> S.Set Position
 253 possibleNeighbours (r, c) = S.fromList [(r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)]
 254
 255
 256 mazePortals edges = S.foldr' mps S.empty edges
 257     where mps e ps = let (p1, p2) = e ^. connects
 258                      in S.insert p1 $ S.insert p2 ps
 259
 260
 261 searchMaze ::  SearchState s => MazeContext (Maybe (Agendum s))
 262 searchMaze =
 263     do agenda <- initAgenda
 264        aStar agenda S.empty
 265
 266 initAgenda ::  SearchState s => MazeContext (Agenda s)
 267 initAgenda =
 268     do edges <- asks _maze
 269        let portals = mazePortals edges
 270        let portal = S.findMin $ S.filter (\p -> p ^. label == "AA") portals
 271        let ss = emptySearchState portal
 272        c <- estimateCost ss
 273        return $ P.singleton c Agendum { _current = ss, _trail = Q.empty, _cost = c}
 274
 275
 276 aStar ::  SearchState s => Agenda s -> ExploredStates s -> MazeContext (Maybe (Agendum s))
 277 aStar agenda closed
 278     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 279     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMin agenda) ) False = undefined
 280     | P.null agenda = return Nothing
 281     | otherwise =
 282         do  let (_, currentAgendum) = P.findMin agenda
 283             let reached = _current currentAgendum
 284             nexts <- candidates currentAgendum closed
 285             let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
 286             reachedGoal <- isGoal reached
 287             if reachedGoal
 288             then return (Just currentAgendum)
 289             else if reached `S.member` closed
 290                  then aStar (P.deleteMin agenda) closed
 291                  else aStar newAgenda (S.insert reached closed)
 292
 293
 294 candidates ::  SearchState s => Agendum s -> ExploredStates s -> MazeContext (Q.Seq (Agendum s))
 295 candidates agendum closed =
 296     do  let candidate = agendum ^. current
 297         let previous = agendum ^. trail
 298         -- let prevCost = agendum ^. cost
 299         succs <- successors candidate
 300         let nonloops = Q.filter (\s -> not $ (fst s) `S.member` closed) succs
 301         mapM (makeAgendum previous) nonloops
 302
 303 makeAgendum ::  SearchState s => (Q.Seq Edge) -> (s, Edge) -> MazeContext (Agendum s)
 304 makeAgendum previous (newP, newE) =
 305     do predicted <- estimateCost newP
 306        let newTrail = newE <| previous
 307        let incurred = sum $ fmap (^. distance) newTrail
 308        return Agendum { _current = newP
 309                       , _trail = newTrail
 310                       , _cost = incurred + predicted
 311                       }
 312
 313
 314 instance SearchState Portal where
 315
 316     emptySearchState portal = portal
 317
 318     -- successors :: Portal -> MazeContext (Q.Seq (Portal, Edge))
 319     successors portal =
 320         do maze <- asks _maze
 321            let edges = S.filter (edgeTouches portal) maze
 322            let locations = S.map (\e -> (edgeOther portal e, e)) edges
 323            let succs = S.foldr' (<|) Q.empty locations
 324            return succs
 325
 326     -- estimateCost :: Portal -> MazeContext Int
 327     estimateCost _portal = return 0
 328
 329     -- isGoal :: Portal -> MazeContext Bool
 330     isGoal portal = return $ portal ^. label == "ZZ"
 331
 332 instance SearchState LevelledSearchState where
 333     emptySearchState portal = LevelledSearchState {_portalS = portal, _levelS = 0}
 334
 335     -- successors :: LevelledSearchState -> MazeContext (Q.Seq (LevelledSearchState, Edge))
 336     successors ss =
 337         do maze <- asks _maze
 338            let lvl = ss ^. levelS
 339            let portal = ss ^. portalS
 340            let edges = S.filter (edgeTouches portal) maze
 341            let lvlEdges = S.filter (edgeAtLevel portal lvl) edges
 342            let locations = S.map (\e -> (newLSS portal lvl e, e)) lvlEdges
 343            let locations' = S.filter (\(s, _) -> (s ^. levelS) >= 0) locations
 344            let succs = S.foldr' (<|) Q.empty locations'
 345            return succs
 346
 347     -- estimateCost :: Portal -> MazeContext Int
 348     estimateCost ss = -- return 0
 349         do let lvl = ss ^. levelS
 350            cpl <- asks _costPerLevel
 351            ctf <- asks _costToFinish
 352            let cplT = if ss ^. portalS . location == Outer
 353                       then cpl * (lvl - 1) + 1
 354                       else cpl * lvl
 355            if isTerminal (ss ^. portalS)
 356            then return 0
 357            else return (cplT + ctf)
 358
 359     -- isGoal :: LevelledSearchState -> MazeContext Bool
 360     isGoal ss = return $ ss ^. portalS . label == "ZZ"
 361
 362 edgeAtLevel portal lvl edge
 363     -- | (lvl == 0) && (isTerminal other) && (et == Walk) = True
 364     | (lvl /= 0) && (isTerminal other) && (et == Walk) = False
 365     | (lvl == 0) && (not $ isTerminal other) && (et == Walk) && ((other ^. location) == Outer) = False
 366     | otherwise = True
 367     where other = edgeOther portal edge
 368           et = edge ^. edgeType
 369
 370 isTerminal p = (p ^. label == "AA") || (p ^. label == "ZZ")
 371
 372 newLSS :: Portal -> Int -> Edge -> LevelledSearchState
 373 newLSS portal lvl edge
 374     | et == Teleport && pl == Outer = LevelledSearchState { _portalS = otherPortal, _levelS = lvl - 1 }
 375     | et == Teleport && pl == Inner = LevelledSearchState { _portalS = otherPortal, _levelS = lvl + 1 }
 376     | otherwise = LevelledSearchState { _portalS = otherPortal, _levelS = lvl } -- et == Walk
 377     where pl = portal ^. location
 378           et = edge ^. edgeType
 379           otherPortal = edgeOther portal edge
 380
 381
 382 showContracted :: Maze -> String
 383 showContracted m = "graph Maze {\n" ++ bulk ++ "\n}"
 384     where   bulk = S.foldr (\e s -> (showEdge e) ++ s) "" (_maze m)
 385
 386 showEdge :: Edge -> String
 387 showEdge e = (showPortal h) ++ " -- " ++ (showPortal t) ++ " [ label = \"" ++ edgeLabel ++ "\" style = \"" ++ style ++ "\"];\n"
 388     where -- edgeLabel e = (show (e ^. edgeType)) ++ ", " ++ (show (e ^. distance))
 389           (edgeLabel, style) =
 390                 if (e ^. edgeType) == Walk
 391                 then (show (e ^. distance), "solid")
 392                 else ("", "dashed")
 393           (h, t) = e ^. connects
 394           -- showPortal p = p ^. label ++ (show (fst (p ^. position))) ++ "c" ++ (show (snd (p ^. position))) ++ (take 1 $ show (p ^. location))
 395           showPortal p = p ^. label ++ "_" ++ (take 1 $ show (p ^. location))