Reworking day 16
[advent-of-code-22.git] / advent16 / MainOriginalNoBeam.hs
1 -- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
2
3 -- import Debug.Trace
4
5 import AoC
6 import Data.Text (Text)
7 import qualified Data.Text.IO as TIO
8 import Data.Attoparsec.Text hiding (take, D)
9 import Control.Applicative
10 import qualified Data.PQueue.Prio.Max as P
11 import qualified Data.Set as S
12 import qualified Data.Sequence as Q
13 import qualified Data.Map.Strict as M
14 import Data.Map.Strict ((!))
15 import Data.Sequence ((|>))
16 import Data.List
17 import Data.List.Split (chunksOf)
18 import Data.Ord
19 import Control.Monad.Reader
20 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
21
22 -- pattern Empty <- (Q.viewl -> Q.EmptyL) where Empty = Q.empty
23 -- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<) = (Q.<|)
24 -- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>) = (Q.|>)
25
26 type RoomID = String
27
28 data Room = Room
29 { _flowRate :: Int
30 , _tunnels :: [RoomID]
31 } deriving (Eq, Show, Ord)
32 makeLenses ''Room
33
34 type Cave = M.Map RoomID Room
35 data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int}
36
37 type CaveContext = Reader TimedCave
38
39 data SingleSearchState = SingleSearchState
40 { _currentRoom :: RoomID
41 , _sOpenValves :: S.Set RoomID
42 } deriving (Eq, Show, Ord)
43 makeLenses ''SingleSearchState
44
45 data DoubleSearchState = DoubleSearchState
46 { _personRoom :: RoomID
47 , _elephantRoom :: RoomID
48 , _dOpenValves :: S.Set RoomID
49 } deriving (Eq, Show, Ord)
50 makeLenses ''DoubleSearchState
51
52 data Agendum s =
53 Agendum { _current :: s
54 , _trail :: Q.Seq s
55 , _trailBenefit :: Int
56 , _benefit :: Int
57 } deriving (Show, Eq, Ord)
58 makeLenses ''Agendum
59
60 type Agenda s = P.MaxPQueue Int (Agendum s)
61
62 type ExploredStates s = S.Set (s, Int, Int)
63
64
65 class (Eq s, Ord s, Show s) => SearchState s where
66 emptySearchState :: RoomID -> s
67 currentFlow :: s -> CaveContext Int
68 successors :: s -> CaveContext (Q.Seq s)
69 estimateBenefit :: s -> Int -> CaveContext Int
70
71 instance SearchState SingleSearchState where
72 emptySearchState startID = SingleSearchState { _currentRoom = startID, _sOpenValves = S.empty }
73
74 currentFlow state =
75 do cave <- asks getCave
76 let valves = state ^. sOpenValves
77 let presentRooms = cave `M.restrictKeys` valves
78 return $ sumOf (folded . flowRate) presentRooms
79
80 successors state =
81 do isFF <- isFullFlow state
82 let here = state ^. currentRoom
83 let opened = state ^. sOpenValves
84 succPairs <- personSuccessor here opened
85 let succStates =
86 [ SingleSearchState
87 { _currentRoom = r
88 , _sOpenValves = o
89 }
90 | (r, o) <- succPairs
91 ]
92 if isFF
93 then return $ Q.singleton state
94 else return $ Q.fromList succStates
95
96 estimateBenefit here timeElapsed =
97 do cave <- asks getCave
98 timeLimit <- asks getTimeLimit
99 let timeRemaining = timeLimit - (timeElapsed + 2)
100 cf <- currentFlow here
101 let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
102 let sortedClosedValves = sortOn Down closedValves
103 let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
104 return $ (cf * timeRemaining) + otherValveFlows
105
106
107 instance SearchState DoubleSearchState where
108 emptySearchState startID = DoubleSearchState
109 { _personRoom = startID
110 , _elephantRoom = startID
111 , _dOpenValves = S.empty
112 }
113
114 currentFlow state =
115 do cave <- asks getCave
116 let valves = state ^. dOpenValves
117 let presentRooms = cave `M.restrictKeys` valves
118 return $ sumOf (folded . flowRate) presentRooms
119
120 successors state =
121 do isFF <- isFullFlow state
122 let pHere = state ^. personRoom
123 let eHere = state ^. elephantRoom
124 let opened = state ^. dOpenValves
125 pSuccPairs <- personSuccessor pHere opened
126 eSuccPairs <- personSuccessor eHere opened
127 let succStates =
128 [ DoubleSearchState
129 { _personRoom = p
130 , _elephantRoom = e
131 , _dOpenValves = S.union po eo
132 }
133 | (p, po) <- pSuccPairs
134 , (e, eo) <- eSuccPairs
135 ]
136 if isFF
137 then return $ Q.singleton state
138 else return $ Q.fromList succStates
139
140 estimateBenefit here timeElapsed =
141 do cave <- asks getCave
142 timeLimit <- asks getTimeLimit
143 let timeRemaining = timeLimit - (timeElapsed + 2)
144 cf <- currentFlow here
145 let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
146 let sortedClosedValves = fmap sum $ chunksOf 2 $ sortOn Down closedValves
147 let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
148 return $ (cf * timeRemaining) + otherValveFlows
149
150
151 main :: IO ()
152 main =
153 do dataFileName <- getDataFileName
154 text <- TIO.readFile dataFileName
155 let cave = successfulParse text
156 -- print cave
157 print $ part1 cave
158 print $ part2 cave
159
160 -- part1 :: Cave -> Maybe (Agendum SingleSearchState)
161 -- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
162
163 -- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
164 -- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
165
166 part1, part2 :: Cave -> Int
167 part1 cave = maybe 0 _benefit result
168 where result = runReader (searchCave "AA") (TimedCave cave 30) :: Maybe (Agendum SingleSearchState)
169 part2 cave = maybe 0 _benefit result
170 where result = runReader (searchCave "AA") (TimedCave cave 26) :: Maybe (Agendum DoubleSearchState)
171
172 searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
173 searchCave startRoom =
174 do agenda <- initAgenda startRoom
175 aStar agenda S.empty
176
177 initAgenda :: SearchState s => String -> CaveContext (Agenda s)
178 initAgenda startID =
179 do let startState = emptySearchState startID
180 b <- estimateBenefit startState 0
181 return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
182
183 aStar :: SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
184 aStar agenda closed
185 -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
186 -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : len " ++ (show $ Q.length $ _trail $ snd $ P.findMax agenda)) False = undefined
187 | P.null agenda = return Nothing
188 | otherwise =
189 do let (_, currentAgendum) = P.findMax agenda
190 let reached = currentAgendum ^. current
191 nexts <- candidates currentAgendum closed
192 let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
193 -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
194 -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
195 reachedGoal <- isGoal currentAgendum
196 let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
197 if reachedGoal
198 then return (Just currentAgendum)
199 else if (cl `S.member` closed)
200 then aStar (P.deleteMax agenda) closed
201 else aStar newAgenda (S.insert cl closed)
202 -- else aStar beamAgenda (S.insert cl closed)
203
204
205 candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
206 candidates agendum closed =
207 do let candidate = agendum ^. current
208 let previous = agendum ^. trail
209 let prevBenefit = agendum ^. trailBenefit
210 succs <- successors candidate
211 succAgs <- mapM (makeAgendum previous prevBenefit) succs
212 let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
213 return nonloops
214
215 personSuccessor, openValveSuccessor, walkSuccessor :: RoomID -> S.Set RoomID -> CaveContext [(RoomID, S.Set RoomID)]
216 personSuccessor here opened =
217 do ovs <- openValveSuccessor here opened
218 ws <- walkSuccessor here opened
219 return (ovs ++ ws)
220
221 openValveSuccessor here opened
222 | here `S.member` opened = return []
223 | otherwise = return [(here, S.insert here opened)]
224
225 walkSuccessor here opened =
226 do cave <- asks getCave
227 let neighbours = (cave ! here) ^. tunnels
228 return [(n, opened) | n <- neighbours]
229
230 makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
231 makeAgendum previous prevBenefit newState =
232 do predicted <- estimateBenefit newState (Q.length previous)
233 cf <- currentFlow newState
234 let newTrail = previous |> newState
235 let incurred = prevBenefit + cf
236 return Agendum { _current = newState
237 , _trail = newTrail
238 , _trailBenefit = incurred
239 , _benefit = incurred + predicted
240 }
241
242
243 isGoal :: SearchState s => Agendum s -> CaveContext Bool
244 isGoal agendum =
245 do timeLimit <- asks getTimeLimit
246 return $ Q.length (agendum ^. trail) == (timeLimit - 1)
247
248 isFullFlow :: SearchState s => s -> CaveContext Bool
249 isFullFlow state =
250 do cave <- asks getCave
251 cf <- currentFlow state
252 let ff = sumOf (folded . flowRate) cave
253 return (cf == ff)
254
255
256 -- Parse the input file
257
258 caveP :: Parser Cave
259 valveP :: Parser (RoomID, Room)
260 roomP :: Parser Room
261 tunnelsP :: Parser [RoomID]
262 tunnelTextP :: Parser Text
263
264 caveP = M.fromList <$> valveP `sepBy` endOfLine
265 valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
266 roomP = roomify <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
267 where roomify v ts = Room {_flowRate = v, _tunnels = ts }
268 tunnelsP = (many1 letter) `sepBy` ", "
269 tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
270
271 successfulParse :: Text -> Cave
272 successfulParse input =
273 case parseOnly caveP input of
274 Left _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
275 Right cave -> cave