--- /dev/null
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+-- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>)))
+import Data.Sequence ( (|>), Seq((:|>)) )
+import Data.List
+import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+import Control.Parallel.Strategies
+
+
+type RoomID = String
+
+data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
+ deriving (Eq, Show, Ord)
+makeLenses ''Tunnel
+
+data Room = Room
+ { _flowRate :: Int
+ , _tunnels :: S.Set Tunnel
+ } deriving (Eq, Show, Ord)
+makeLenses ''Room
+
+type Cave = M.Map RoomID Room
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
+
+type CaveContext = Reader TimedCave
+
+data SearchState = SearchState
+ { _currentRoom :: RoomID
+ , _currentTime :: Int
+ , _openValves :: S.Set RoomID
+ } deriving (Eq, Show, Ord)
+makeLenses ''SearchState
+
+data Agendum =
+ Agendum { _current :: SearchState
+ , _trail :: Q.Seq SearchState
+ , _trailBenefit :: Int
+ , _benefit :: Int
+ } deriving (Show, Eq, Ord)
+makeLenses ''Agendum
+
+type Agenda = P.MaxPQueue Int Agendum
+
+-- state, total flowed so far
+type ExploredStates = S.Set (SearchState, Int)
+
+type PartSolutions = M.Map (S.Set RoomID) Int
+
+
+main :: IO ()
+main =
+ do dataFileName <- getDataFileName
+ text <- TIO.readFile dataFileName
+ let expandedCave = successfulParse text
+ -- print cave
+ -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
+ -- print $ compress cave
+ -- putStrLn $ dotify expandedCave
+ let cave = compress expandedCave
+ print $ part1 cave
+ print $ part2 cave
+
+-- dotify cave = "graph G {\n" ++ (unlines $ concat $ M.elems $ M.mapWithKey showCRoom cave) ++ "\n}\n"
+-- where showCRoom roomID room = filter (not . null) ((showCRoomShape roomID room) : (showCRoomLinks roomID room))
+
+-- showCRoomShape roomID room
+-- | room ^. flowRate > 0 = roomID ++ " [fillcolor=grey label=\"" ++ roomID ++ ": " ++ (show $ room ^. flowRate) ++ "\"];"
+-- | otherwise = ""
+
+-- showCRoomLinks roomID room = [roomID ++ " -- " ++ (t ^. tunnelTo) ++ ";" | t <- S.toList $ room ^. tunnels, (t ^. tunnelTo) > roomID ]
+
+part1, part2 :: Cave -> Int
+-- part1 :: Cave -> Int
+part1 cave = runSearch 30 cave
+part2 cave = maximum (fmap maximum chunkSolns `using` parList rdeepseq)
+ where rawSolutions = runSearchAll 26 cave
+ solutionList = M.toList rawSolutions
+ combinations = [ fp + fe
+ | (p, fp) <- solutionList
+ , (e, fe) <- solutionList
+ , p < e
+ , S.disjoint p e
+ ]
+ chunkSolns = chunksOf 10000 combinations
+
+includeAgendum :: PartSolutions -> Agendum -> CaveContext PartSolutions
+includeAgendum results agendum =
+ do cf <- currentFlow (agendum ^. current)
+ timeLimit <- asks getTimeLimit
+ let timeLeft = timeLimit - timeOf (agendum ^. current)
+ let remainingFlow = cf * timeLeft
+ let totalFlow = remainingFlow + agendum ^. trailBenefit
+ let visitedSet = agendum ^. current . openValves
+ let currentBest = M.findWithDefault 0 visitedSet results
+ if totalFlow > currentBest
+ then return (M.insert visitedSet totalFlow results)
+ else return results
+
+runSearch :: Int -> Cave -> Int
+runSearch timeLimit cave = maybe 0 _benefit result
+ where result = runReader (searchCave "AA") (TimedCave cave timeLimit sortedRooms)
+ sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+
+runSearchAll :: Int -> Cave -> PartSolutions
+runSearchAll timeLimit cave = result
+ where result = runReader (searchCaveAll "AA") (TimedCave cave timeLimit sortedRooms)
+ sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+
+
+searchCave :: String -> CaveContext (Maybe Agendum)
+searchCave startRoom =
+ do agenda <- initAgenda startRoom
+ aStar agenda S.empty
+
+searchCaveAll :: String -> CaveContext PartSolutions
+searchCaveAll startRoom =
+ do agenda <- initAgenda startRoom
+ allSolutions agenda S.empty M.empty
+
+initAgenda :: String -> CaveContext Agenda
+initAgenda startID =
+ do let startState = emptySearchState startID
+ b <- estimateBenefit startState
+ return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar :: Agenda -> ExploredStates -> CaveContext (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.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
+ -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ " : trail " ++ (show $ _trail $ snd $ P.findMax agenda) ++ " : closed " ++ (show closed)) False = undefined
+ -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
+ | P.null agenda = return Nothing
+ | otherwise =
+ do let (_, currentAgendum) = P.findMax agenda
+ let reached = currentAgendum ^. current
+ nexts <- candidates currentAgendum closed
+ let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+ reachedGoal <- isGoal currentAgendum
+ let cl = (reached, currentAgendum ^. trailBenefit)
+ if reachedGoal
+ then return (Just currentAgendum)
+ else if (cl `S.member` closed)
+ then aStar (P.deleteMax agenda) closed
+ else aStar newAgenda (S.insert cl closed)
+
+allSolutions :: Agenda -> ExploredStates -> PartSolutions -> CaveContext PartSolutions
+allSolutions agenda closed foundSolutions
+ | P.null agenda = return foundSolutions
+ | otherwise =
+ do let (_, currentAgendum) = P.findMax agenda
+ let reached = currentAgendum ^. current
+ nexts <- candidates currentAgendum closed
+ let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+ reachedGoal <- isGoal currentAgendum
+ let cl = (reached, currentAgendum ^. trailBenefit)
+ newFoundSolutions <- includeAgendum foundSolutions currentAgendum
+ if reachedGoal
+ then allSolutions (P.deleteMax agenda) closed newFoundSolutions
+ else if (cl `S.member` closed)
+ then allSolutions (P.deleteMax agenda) closed foundSolutions
+ else allSolutions newAgenda (S.insert cl closed) newFoundSolutions
+
+
+candidates :: Agendum -> ExploredStates -> CaveContext (Q.Seq Agendum)
+candidates agendum closed =
+ do let candidate = agendum ^. current
+ let previous = agendum ^. trail
+ let prevBenefit = agendum ^. trailBenefit
+ succs <- successors candidate
+ succAgs <- mapM (makeAgendum previous prevBenefit) succs
+ let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
+ return nonloops
+
+emptySearchState :: RoomID -> SearchState
+emptySearchState startID = SearchState
+ { _currentRoom = startID
+ , _currentTime = 0
+ , _openValves = S.empty
+ }
+
+currentFlow :: SearchState -> CaveContext Int
+currentFlow state =
+ do cave <- asks getCave
+ let valves = state ^. openValves
+ let presentRooms = cave `M.restrictKeys` valves
+ return $ sumOf (folded . flowRate) presentRooms
+
+timeOf :: SearchState -> Int
+timeOf state = state ^. currentTime
+
+successors :: SearchState -> CaveContext (Q.Seq SearchState)
+successors state =
+ do isFF <- isFullFlow state
+ cave <- asks getCave
+ timeLimit <- asks getTimeLimit
+ let here = state ^. currentRoom
+ let opened = state ^. openValves
+ let now = state ^. currentTime
+ let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
+ let moves = [ SearchState
+ { _currentRoom = (t ^. tunnelTo)
+ , _currentTime = now + (t ^. tunnelLength)
+ , _openValves = opened
+ }
+ | t <- remaining
+ , now + (t ^. tunnelLength) <= timeLimit
+ ]
+ let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0 && now < timeLimit
+ then [SearchState { _currentRoom = here, _currentTime = now + 1, _openValves = S.insert here opened }]
+ else []
+ let nexts = if null opens then moves else opens
+ let nexts' = if null nexts
+ then [ SearchState
+ { _currentRoom = here
+ , _currentTime = timeLimit
+ , _openValves = opened
+ } ]
+ else nexts
+ let succs = Q.fromList nexts'
+ if isFF || (Q.null succs)
+ then return $ Q.singleton (state & currentTime .~ timeLimit)
+ else return succs
+
+
+estimateBenefit :: SearchState -> CaveContext Int
+estimateBenefit here =
+ do cave <- asks getCave
+ timeLimit <- asks getTimeLimit
+ let timeRemaining = timeLimit - (timeOf here)
+ cf <- currentFlow here
+ sortedValves <- asks getSortedRooms
+ let opened = here ^. openValves
+ let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
+ let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+ return $ (cf * timeRemaining) + otherValveFlows
+
+makeAgendum :: Q.Seq SearchState -> Int -> SearchState -> CaveContext Agendum
+makeAgendum previous prevBenefit newState =
+ do predicted <- estimateBenefit newState -- (Q.length previous)
+ -- cf <- currentFlow newState
+ oldFlow <- lastFlow previous (timeOf newState)
+ let newTrail = previous |> newState
+ let incurred = prevBenefit + oldFlow
+ return Agendum { _current = newState
+ , _trail = newTrail
+ , _trailBenefit = incurred
+ , _benefit = incurred + predicted
+ }
+
+lastFlow :: Q.Seq SearchState -> Int -> CaveContext Int
+lastFlow Q.Empty _ = return 0
+lastFlow (_ :|> previous) newTime =
+ do cf <- currentFlow previous
+ let dt = newTime - (timeOf previous)
+ return (cf * dt)
+
+isGoal :: Agendum -> CaveContext Bool
+isGoal agendum =
+ do timeLimit <- asks getTimeLimit
+ let s = agendum ^. current
+ return $ (timeOf s) == timeLimit
+
+isFullFlow :: SearchState -> CaveContext Bool
+isFullFlow state =
+ do cave <- asks getCave
+ cf <- currentFlow state
+ let ff = sumOf (folded . flowRate) cave
+ return (cf == ff)
+
+compress :: Cave -> Cave
+compress cave = M.mapWithKey (compressRoom cave) cave
+
+compressRoom :: Cave -> RoomID -> Room -> Room
+compressRoom cave here room = room & tunnels .~ t'
+ where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
+
+reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
+reachableFrom _ [] _ routes = routes
+reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
+ | here `S.member` found = reachableFrom cave boundary found routes
+ | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
+ where exits = (cave ! here) ^. tunnels
+ exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
+ legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
+ routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
+ then routes
+ else S.insert tunnel routes
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser (S.Set Tunnel)
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+ -- where roomify v ts = Room {flowRate = v, tunnels = ts }
+tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input =
+ case parseOnly caveP input of
+ Left _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+ Right cave -> cave
\ No newline at end of file
projects / advent-of-code-22.git / blobdiff