--- /dev/null
+-- Writeup at https://work.njae.me.uk/2022/12/21/advent-of-code-2022-day-19/
+
+-- 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.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
+
+-- 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 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 SingleSearchState = SingleSearchState
+ { _resources :: Collection
+ , _robots :: Collection
+ } deriving (Eq, Show, Ord)
+makeLenses ''SingleSearchState
+
+instance NFData SingleSearchState where
+ rnf (SingleSearchState a b) = rnf a `seq` rnf b `seq` ()
+
+data Agendum s =
+ Agendum { _current :: s
+ , _trail :: Q.Seq s
+ , _trailBenefit :: Int
+ , _benefit :: Int
+ } deriving (Show, Eq, Ord)
+makeLenses ''Agendum
+
+type Agenda s = P.MaxPQueue Int (Agendum s)
+
+type ExploredStates s = S.Set (s, Int, Int)
+
+
+class (Eq s, Ord s, Show s) => SearchState s where
+ emptySearchState :: s
+ successors :: s -> BlueprintContext (Q.Seq s)
+ estimateBenefit :: s -> Int -> BlueprintContext Int
+
+instance SearchState SingleSearchState where
+ emptySearchState = SingleSearchState { _resources = MS.empty, _robots = MS.singleton Ore }
+
+ successors state =
+ do blueprint <- asks getBlueprint
+ maxRobots <- asks getMaxRobots
+ let buildableRobots = M.keys $ M.filter (\required -> required `MS.isSubsetOf` (state ^. resources)) blueprint
+ --- if more bots than needed for making any single bot, don't make more of that bot
+ let usefulRobots = MS.foldOccur (\res maxNeeded rs ->
+ if (MS.occur res (state ^. robots)) >= maxNeeded
+ then Data.List.delete res rs
+ else rs
+ ) buildableRobots maxRobots
+ let madeRobots = [ state & robots %~ MS.insert robot
+ & resources %~ ( `MS.difference` (blueprint ! robot) )
+ | robot <- usefulRobots
+ ]
+ let afterBuild = [state] ++ madeRobots
+ let afterGather = fmap (\s -> s & resources %~ (MS.union (state ^. robots))) afterBuild
+ return $ Q.fromList afterGather
+
+
+ estimateBenefit currentState timeElapsed =
+ do timeLimit <- asks getTimeLimit
+ let timeRemaining = timeLimit - (timeElapsed + 1)
+ let currentGeodes = MS.occur Geode (currentState ^. resources)
+ let currentRobotsGather = (MS.occur Geode (currentState ^. robots)) * timeRemaining
+ let newRobotsGather = (timeRemaining * (timeRemaining + 1)) `div` 2
+ return $ currentGeodes + currentRobotsGather + newRobotsGather
+
+
+main :: IO ()
+main =
+ do dataFileName <- getDataFileName
+ text <- TIO.readFile dataFileName
+ let blueprints = successfulParse text
+ print $ part1 blueprints
+ -- print $ part2 blueprints
+
+-- part1 :: [(Int, Blueprint)] -> Int
+part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- results]
+-- part1 blueprints = pResults
+ where results = [ (n, _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
+ | (n, blueprint) <- blueprints ] :: [(Int, SingleSearchState)]
+ pResults = parMap rdeepseq id results
+ robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
+-- part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- pResults]
+-- where -- results = [ (n, _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
+-- -- | (n, blueprint) <- blueprints ] :: [(Int, SingleSearchState)]
+-- -- pResults = parMap rdeepseq id results
+-- -- pResults = (fmap runABlueprint blueprints) `using` parList rdeepseq
+-- pResults = (fmap runABlueprint blueprints) `using` (parList rdeepseq)
+-- runABlueprint (n, blueprint) = (n, _current $ fromJust $
+-- runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
+-- robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
+
+part2 :: [(Int, Blueprint)] -> Int
+part2 blueprints = product [MS.occur Geode (r ^. resources) | r <- pResults]
+ where results = [ _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 32 (robotLimits blueprint))
+ | (_, blueprint) <- (take 3 blueprints) ] :: [SingleSearchState]
+ pResults = parMap rdeepseq id results
+ robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
+
+searchSpace :: SearchState s => BlueprintContext (Maybe (Agendum s))
+searchSpace =
+ do agenda <- initAgenda
+ -- aStar agenda S.empty
+ res <- aStar agenda S.empty
+ return (res `seq` res)
+
+initAgenda :: SearchState s => BlueprintContext (Agenda s)
+initAgenda =
+ do let startState = emptySearchState
+ b <- estimateBenefit startState 0
+ return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar :: SearchState s => Agenda s -> ExploredStates s -> BlueprintContext (Maybe (Agendum s))
+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 beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
+ -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+ reachedGoal <- isGoal currentAgendum
+ let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+ 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)
+ else aStar newAgenda (S.insert cl closed)
+
+candidates :: SearchState s => Agendum s -> ExploredStates s -> BlueprintContext (Q.Seq (Agendum s))
+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, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+ return nonloops
+
+makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> BlueprintContext (Agendum s)
+makeAgendum previous prevBenefit newState =
+ do predicted <- estimateBenefit newState (Q.length previous)
+ let newTrail = previous |> newState
+ -- let incurred = geodesFound newState
+ let incurred = 0
+ return Agendum { _current = newState
+ , _trail = newTrail
+ , _trailBenefit = incurred
+ , _benefit = incurred + predicted
+ }
+
+isGoal :: SearchState s => Agendum s -> BlueprintContext Bool
+isGoal agendum =
+ do timeLimit <- asks getTimeLimit
+ return $ Q.length (agendum ^. trail) == timeLimit
+
+-- 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