--- /dev/null
+-- Writeup at https://work.njae.me.uk/2022/12/21/advent-of-code-2022-day-19/
+-- Optimised at https://work.njae.me.uk/2023/07/24/optimising-haskell-example-4/
+
+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 qualified Data.MultiSet as MS
+import Data.Sequence ((|>))
+import Data.List
+import Data.Maybe
+-- import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+import GHC.Generics (Generic)
+import Control.Parallel.Strategies
+import Control.DeepSeq
+
+data Resource = Ore | Clay | Obsidian | Geode
+ deriving (Show, Eq, Ord, Generic)
+
+instance NFData Resource
+
+type Collection = MS.MultiSet Resource
+
+
+type Blueprint = M.Map Resource Collection
+
+data TimedBlueprint = TimedBlueprint { getBlueprint :: Blueprint, getTimeLimit :: Int, getMaxRobots :: Collection}
+ deriving (Show, Eq, Ord)
+
+type BlueprintContext = Reader TimedBlueprint
+
+data SearchState = SearchState
+ { _resources :: Collection
+ , _robots :: Collection
+ , _currentTime :: Int
+ } deriving (Eq, Show, Ord)
+makeLenses ''SearchState
+
+
+instance NFData SearchState where
+ rnf (SearchState a b c) = rnf a `seq` rnf b `seq` rnf c `seq` ()
+
+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
+
+type ExploredStates = S.Set SearchState
+
+main :: IO ()
+main =
+ do dataFileName <- getDataFileName
+ text <- TIO.readFile dataFileName
+ let blueprints = successfulParse text
+ print $ part1 blueprints
+ print $ part2 blueprints
+
+part1, part2 :: [(Int, Blueprint)] -> Int
+part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- results]
+-- part1 blueprints = results
+ -- where results = fmap (scoreBlueprint 24) blueprints
+ where results = parMap rdeepseq (scoreBlueprint 24) blueprints
+
+-- part2 :: [(Int, Blueprint)] -> Int
+part2 blueprints = product [MS.occur Geode (r ^. resources) | (_, r) <- results]
+ where results = parMap rdeepseq (scoreBlueprint 32) $ take 3 blueprints
+
+robotLimits :: Blueprint -> Collection
+robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
+
+scoreBlueprint :: Int -> (Int, Blueprint) -> (Int, SearchState)
+scoreBlueprint t (n, bp) = ( n
+ , runReader searchSpace (TimedBlueprint bp t (robotLimits bp))
+ )
+
+searchSpace :: BlueprintContext SearchState
+searchSpace =
+ do agenda <- initAgenda
+ -- searchAll agenda S.empty emptySearchState
+ result <- aStar agenda S.empty
+ return $ (fromJust result) ^. current
+
+initAgenda :: BlueprintContext Agenda
+initAgenda =
+ do let startState = emptySearchState
+ b <- estimateBenefit startState
+ return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar :: Agenda -> ExploredStates -> BlueprintContext (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) ++ " benefit " ++ (show $ fst $ P.findMax agenda) ++ " : elapsed " ++ (show $ Q.length $ _trail $ snd $ 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
+ -- let cl = hashSearchState reached
+ atTimeLimit <- isTimeLimit currentAgendum
+ if atTimeLimit
+ then return (Just currentAgendum)
+ else if (reached `S.member` closed)
+ then aStar (P.deleteMax agenda) closed
+ -- else aStar newAgenda (S.insert cl closed)
+ else aStar newAgenda (S.insert reached closed)
+
+candidates :: Agendum -> ExploredStates -> BlueprintContext (Q.Seq Agendum)
+candidates agendum closed =
+ do let candidate = agendum ^. current
+ let previous = agendum ^. trail
+ -- let nextLen = Q.length previous + 1
+ let prevBenefit = agendum ^. trailBenefit
+ succs <- successors candidate
+ succAgs <- mapM (makeAgendum previous prevBenefit) succs
+ let nonloops = Q.filter (\s -> (s ^. current) `S.notMember` closed) succAgs
+ return nonloops
+
+makeAgendum :: Q.Seq SearchState -> Int -> SearchState -> BlueprintContext Agendum
+makeAgendum previous prevBenefit newState =
+ -- do predicted <- estimateBenefit newState (Q.length previous)
+ do predicted <- estimateBenefit newState
+ let newTrail = previous |> newState
+ let incurred = (MS.occur Geode (newState ^. resources))
+ return Agendum { _current = newState
+ , _trail = newTrail
+ , _trailBenefit = incurred
+ , _benefit = incurred + predicted
+ }
+
+isTimeLimit :: Agendum -> BlueprintContext Bool
+isTimeLimit agendum =
+ do timeLimit <- asks getTimeLimit
+ -- return $ Q.length (agendum ^. trail) == timeLimit
+ return $ (agendum ^. current . currentTime) >= timeLimit
+
+emptySearchState :: SearchState
+emptySearchState = SearchState { _resources = MS.empty, _robots = MS.singleton Ore, _currentTime = 0 }
+
+successors :: SearchState -> BlueprintContext (Q.Seq SearchState)
+successors state =
+ do blueprint <- asks getBlueprint
+ maxRobots <- asks getMaxRobots
+ timeLimit <- asks getTimeLimit
+
+ let robotSuccessors = Q.fromList $ catMaybes $ M.elems $ M.mapWithKey (handleRobot state maxRobots timeLimit) blueprint
+
+ let timeRemaining = timeLimit - (state ^. currentTime)
+ let gathered = MS.foldOccur (\res n acc -> MS.insertMany res (n * timeRemaining) acc)
+ MS.empty
+ (state ^. robots)
+ let delayUntilEnd = (state & currentTime .~ timeLimit
+ & resources %~ (MS.union gathered)
+ )
+ return ( robotSuccessors |> delayUntilEnd )
+
+handleRobot :: SearchState -> Collection -> Int -> Resource -> Collection -> Maybe SearchState
+handleRobot state maxRobots timeLimit robot recipe
+ | sufficientRobots robot state maxRobots = Nothing
+ -- | buildableRobot state recipe = buildRobotAndGather robot state recipe
+ | otherwise = buildWhenReady robot state recipe timeLimit
+
+-- do I already have enough of this robot?
+sufficientRobots :: Resource -> SearchState -> Collection -> Bool
+sufficientRobots robot state maxRobots =
+ (robot `MS.member` maxRobots)
+ &&
+ ((MS.occur robot (state ^. robots)) >= (MS.occur robot maxRobots))
+
+buildDelay :: SearchState -> Collection -> Maybe Int
+buildDelay state recipe
+ -- | MS.null delay = Just 0
+ | all (\r -> MS.member r rbts) (MS.distinctElems shortfall) = Just $ maximum0 $ fmap snd $ MS.toOccurList delay
+ | otherwise = Nothing
+ where shortfall = recipe `MS.difference` (state ^. resources)
+ delay = MS.foldOccur calcOneDelay MS.empty shortfall
+ rbts = state ^. robots
+ calcOneDelay resource count acc =
+ MS.insertMany resource
+ -- (count `div` (MS.occur resource rbts) + 1)
+ (ceiling $ (fromIntegral count) / (fromIntegral $ MS.occur resource rbts))
+ acc
+ maximum0 xs = if (null xs) then 0 else maximum xs
+
+buildWhenReady :: Resource -> SearchState -> Collection -> Int -> Maybe SearchState
+buildWhenReady robot state recipe timeLimit =
+ do waitDelay <- buildDelay state recipe
+ delay <- tooLate (state ^. currentTime) (waitDelay + 1) timeLimit
+ let gathered = MS.foldOccur (\res n acc -> MS.insertMany res (n * delay) acc)
+ MS.empty
+ (state ^. robots)
+ return (state & robots %~ MS.insert robot -- add the robot
+ & resources %~ (MS.union gathered)
+ & resources %~ ( `MS.difference` recipe ) -- remove the resources to build it
+ & currentTime %~ (+ delay)
+ )
+
+tooLate :: Int -> Int -> Int -> Maybe Int
+tooLate current delay timeLimit
+ | (current + delay) <= timeLimit = Just delay
+ | otherwise = Nothing
+
+
+estimateBenefit :: SearchState -> BlueprintContext Int
+estimateBenefit currentState =
+ do timeLimit <- asks getTimeLimit
+ let timeElapsed = currentState ^. currentTime
+ let timeRemaining = timeLimit - timeElapsed
+ let currentRobotsGather = (MS.occur Geode (currentState ^. robots)) * timeRemaining
+ let newRobotsGather = (timeRemaining * (timeRemaining + 1)) `div` 2
+ return $ currentRobotsGather + newRobotsGather
+
+
+-- Parse the input file
+
+blueprintsP :: Parser [(Int, Blueprint)]
+blueprintP :: Parser (Int, Blueprint)
+robotP :: Parser (Resource, Collection)
+requirementsP :: Parser Collection
+requirementP :: Parser (Resource, Int)
+resourceP, oreP, clayP, obsidianP, geodeP :: Parser Resource
+
+blueprintsP = blueprintP `sepBy` endOfLine
+blueprintP = blueprintify <$> (("Blueprint " *> decimal) <* ": ") <*> (robotP `sepBy` ". ") <* "."
+ where blueprintify n robots =
+ (n, M.fromList robots)
+robotP = (,) <$> ("Each " *> resourceP) <*> (" robot costs " *> requirementsP)
+
+requirementsP = MS.fromOccurList <$> (requirementP `sepBy` " and ")
+
+requirementP = (flip (,)) <$> (decimal <* " ") <*> resourceP
+
+resourceP = oreP <|> clayP <|> obsidianP <|> geodeP
+oreP = Ore <$ "ore"
+clayP = Clay <$ "clay"
+obsidianP = Obsidian <$ "obsidian"
+geodeP = Geode <$ "geode"
+
+successfulParse :: Text -> [(Int, Blueprint)]
+successfulParse input =
+ case parseOnly blueprintsP input of
+ Left _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
+ Right blueprints -> blueprints
\ No newline at end of file
projects / advent-of-code-22.git / blobdiff