+ 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
+ | 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))
+
+-- assuming I can't build this robot, how long do I have to wait for the current
+-- robots to gather enough to build it?
+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) -- add the gathered resources
+ & 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
+