--- /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
+hashCollection c = (MS.occur Ore c) + 200 * (MS.occur Clay c) + 200 * 200 * (MS.occur Obsidian c) + 200 * 200 * 200 * (MS.occur Geode c)
+
+
+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
+ } deriving (Eq, Show, Ord)
+makeLenses ''SearchState
+hashSearchState s = (hashCollection (s ^. resources), hashCollection (s ^. robots))
+
+instance NFData SearchState where
+ rnf (SearchState a b) = rnf a `seq` rnf b `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, Int, Int)
+-- type ExploredStates = S.Set ((Int, Int), Int, Int)
+type ExploredStates = S.Set ((Int, Int), Int)
+
+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]
+ -- 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
+ , _current $ fromJust $ runReader searchSpace (TimedBlueprint bp t (robotLimits bp))
+ )
+
+searchSpace :: BlueprintContext (Maybe Agendum)
+searchSpace =
+ do agenda <- initAgenda
+ aStar agenda S.empty 0
+
+initAgenda :: BlueprintContext Agenda
+initAgenda =
+ do let startState = emptySearchState
+ b <- estimateBenefit startState 0
+ return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar :: Agenda -> ExploredStates -> Int -> BlueprintContext (Maybe Agendum)
+aStar agenda closed bestFound
+ -- | 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
+ -- | trace ((show bestFound) ++ " " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ ": " ++ (show $ _current $ snd $ P.findMax agenda)) False = undefined
+ | P.null agenda = return Nothing
+ | otherwise =
+ do let (_, currentAgendum) = P.findMax agenda
+ let reached = currentAgendum ^. current
+ let bestFound' = max bestFound (currentAgendum ^. trailBenefit)
+ nexts <- candidates currentAgendum closed bestFound'
+ 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)
+ -- let cl = (hashSearchState reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+ let cl = (hashSearchState reached, Q.length $ currentAgendum ^. trail)
+ if reachedGoal
+ then return (Just currentAgendum)
+ else if (cl `S.member` closed)
+ then aStar (P.deleteMax agenda) closed bestFound'
+ -- else aStar newAgenda (S.insert cl closed)
+ else aStar newAgenda (S.insert cl closed) bestFound'
+
+candidates :: Agendum -> ExploredStates -> Int -> BlueprintContext (Q.Seq Agendum)
+candidates agendum closed bestFound =
+ 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 succAgs' = Q.filter (\a -> a ^. benefit >= bestFound) succAgs
+ -- let nonloops = Q.filter (\s -> (hashSearchState $ s ^. current, s ^. trailBenefit, nextLen) `S.notMember` closed) succAgs'
+ let nonloops = Q.filter (\s -> (hashSearchState $ s ^. current, nextLen) `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)
+ let newTrail = previous |> newState
+ let incurred = (MS.occur Geode (newState ^. resources))
+ -- let incurred = 0
+ return Agendum { _current = newState
+ , _trail = newTrail
+ , _trailBenefit = incurred
+ , _benefit = incurred + predicted
+ }
+
+isGoal :: Agendum -> BlueprintContext Bool
+isGoal agendum =
+ do timeLimit <- asks getTimeLimit
+ return $ Q.length (agendum ^. trail) == timeLimit
+
+emptySearchState :: SearchState
+emptySearchState = SearchState { _resources = MS.empty, _robots = MS.singleton Ore }
+
+successors :: SearchState -> BlueprintContext (Q.Seq SearchState)
+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 :: SearchState -> Int -> BlueprintContext Int
+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
+ 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