advent19/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/21/advent-of-code-2022-day-19/
   2
   3 -- import Debug.Trace
   4
   5 import AoC
   6 import Data.Text (Text)
   7 import qualified Data.Text.IO as TIO
   8 import Data.Attoparsec.Text hiding (take, D)
   9 import Control.Applicative
  10 import qualified Data.PQueue.Prio.Max as P
  11 import qualified Data.Set as S
  12 import qualified Data.Sequence as Q
  13 import qualified Data.Map.Strict as M
  14 import Data.Map.Strict ((!))
  15 import Data.MultiSet as MS
  16 import Data.Sequence ((|>))
  17 import Data.List
  18 import Data.Maybe
  19 -- import Data.Ord
  20 import Control.Monad.Reader
  21 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  22 import GHC.Generics (Generic)
  23 import Control.Parallel.Strategies
  24 import Control.DeepSeq
  25
  26 -- pattern Empty   <- (Q.viewl -> Q.EmptyL)  where Empty = Q.empty
  27 -- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<)  = (Q.<|)
  28 -- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>)  = (Q.|>)
  29
  30 data Resource = Ore | Clay | Obsidian | Geode
  31     deriving (Show, Eq, Ord, Generic)
  32
  33 instance NFData Resource
  34
  35 type Collection = MS.MultiSet Resource
  36
  37 type Blueprint = M.Map Resource Collection
  38
  39 data TimedBlueprint = TimedBlueprint { getBlueprint :: Blueprint, getTimeLimit :: Int, getMaxRobots :: Collection}
  40     deriving (Show, Eq, Ord)
  41
  42 type BlueprintContext = Reader TimedBlueprint
  43
  44 data SingleSearchState = SingleSearchState
  45     { _resources :: Collection
  46     , _robots :: Collection
  47     } deriving (Eq, Show, Ord)
  48 makeLenses ''SingleSearchState
  49
  50 instance NFData SingleSearchState where
  51   rnf (SingleSearchState a b) = rnf a `seq` rnf b `seq` ()
  52
  53 data Agendum s =
  54     Agendum { _current :: s
  55             , _trail :: Q.Seq s
  56             , _trailBenefit :: Int
  57             , _benefit :: Int
  58             } deriving (Show, Eq, Ord)
  59 makeLenses ''Agendum
  60
  61 type Agenda s = P.MaxPQueue Int (Agendum s)
  62
  63 type ExploredStates s = S.Set (s, Int, Int)
  64
  65
  66 class (Eq s, Ord s, Show s) => SearchState s where
  67   emptySearchState :: s
  68   successors :: s -> BlueprintContext (Q.Seq s)
  69   estimateBenefit :: s -> Int -> BlueprintContext Int
  70
  71 instance SearchState SingleSearchState where
  72   emptySearchState = SingleSearchState { _resources = MS.empty, _robots = MS.singleton Ore }
  73
  74   successors state =
  75     do blueprint <- asks getBlueprint
  76        maxRobots <- asks getMaxRobots
  77        let buildableRobots = M.keys $ M.filter (\required -> required `MS.isSubsetOf` (state ^. resources)) blueprint
  78        --- if more bots than needed for making any single bot, don't make more of that bot
  79        let usefulRobots = MS.foldOccur (\res maxNeeded rs ->
  80                                           if (MS.occur res (state ^. robots)) >= maxNeeded
  81                                           then Data.List.delete res rs
  82                                           else rs
  83                                         ) buildableRobots maxRobots
  84        let madeRobots = [ state & robots %~ MS.insert robot
  85                                 & resources %~ ( `MS.difference` (blueprint ! robot) )
  86                         | robot <- usefulRobots
  87                         ]
  88        let afterBuild = [state] ++ madeRobots
  89        let afterGather = fmap (\s -> s & resources %~ (MS.union (state ^. robots))) afterBuild
  90        return $ Q.fromList afterGather
  91
  92
  93   estimateBenefit currentState timeElapsed =
  94     do timeLimit <- asks getTimeLimit
  95        let timeRemaining = timeLimit - (timeElapsed + 1)
  96        let currentGeodes = MS.occur Geode (currentState ^. resources)
  97        let currentRobotsGather = (MS.occur Geode (currentState ^. robots)) * timeRemaining
  98        let newRobotsGather = (timeRemaining * (timeRemaining + 1)) `div` 2
  99        return $ currentGeodes + currentRobotsGather + newRobotsGather
 100
 101
 102 main :: IO ()
 103 main =
 104   do  dataFileName <- getDataFileName
 105       text <- TIO.readFile dataFileName
 106       let blueprints = successfulParse text
 107       print $ part1 blueprints
 108       print $ part2 blueprints
 109
 110 part1 :: [(Int, Blueprint)] -> Int
 111 part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- results]
 112     where results = [ (n, _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
 113                     | (n, blueprint) <- blueprints ] :: [(Int, SingleSearchState)]
 114           robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
 115 -- part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- pResults]
 116 --     where -- results = [ (n, _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
 117 --                     -- | (n, blueprint) <- blueprints ] :: [(Int, SingleSearchState)]
 118 --           -- pResults = parMap rdeepseq id results
 119 --           -- pResults = (fmap runABlueprint blueprints) `using` parList rdeepseq
 120 --           pResults = (fmap runABlueprint blueprints) `using` (parList rdeepseq)
 121 --           runABlueprint (n, blueprint) = (n, _current $ fromJust $
 122 --                                                 runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
 123 --           robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
 124
 125 part2 :: [(Int, Blueprint)] -> Int
 126 part2 blueprints = product [MS.occur Geode (r ^. resources) | r <- pResults]
 127     where results = [ _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 32 (robotLimits blueprint))
 128                     | (_, blueprint) <- (take 3 blueprints) ] :: [SingleSearchState]
 129           pResults = parMap rdeepseq id results
 130           robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
 131
 132 searchSpace :: SearchState s => BlueprintContext (Maybe (Agendum s))
 133 searchSpace =
 134     do agenda <- initAgenda
 135        -- aStar agenda S.empty
 136        res <- aStar agenda S.empty
 137        return (res `seq` res)
 138
 139 initAgenda :: SearchState s => BlueprintContext (Agenda s)
 140 initAgenda =
 141     do let startState = emptySearchState
 142        b <- estimateBenefit startState 0
 143        return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
 144
 145 aStar ::  SearchState s => Agenda s -> ExploredStates s -> BlueprintContext (Maybe (Agendum s))
 146 aStar agenda closed
 147     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 148     -- | 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
 149     | P.null agenda = return Nothing
 150     | otherwise =
 151         do  let (_, currentAgendum) = P.findMax agenda
 152             let reached = currentAgendum ^. current
 153             nexts <- candidates currentAgendum closed
 154             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
 155             -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
 156             -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
 157             reachedGoal <- isGoal currentAgendum
 158             let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
 159             if reachedGoal
 160             then return (Just currentAgendum)
 161             else if (cl `S.member` closed)
 162                  then aStar (P.deleteMax agenda) closed
 163                  -- else aStar newAgenda (S.insert cl closed)
 164                  else aStar newAgenda (S.insert cl closed)
 165
 166 candidates :: SearchState s => Agendum s -> ExploredStates s -> BlueprintContext (Q.Seq (Agendum s))
 167 candidates agendum closed =
 168     do  let candidate = agendum ^. current
 169         let previous = agendum ^. trail
 170         let prevBenefit = agendum ^. trailBenefit
 171         succs <- successors candidate
 172         succAgs <- mapM (makeAgendum previous prevBenefit) succs
 173         let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
 174         return nonloops
 175
 176 makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> BlueprintContext (Agendum s)
 177 makeAgendum previous prevBenefit newState =
 178     do predicted <- estimateBenefit newState (Q.length previous)
 179        let newTrail = previous |> newState
 180        -- let incurred = geodesFound newState
 181        let incurred = 0
 182        return Agendum { _current = newState
 183                       , _trail = newTrail
 184                       , _trailBenefit = incurred
 185                       , _benefit = incurred + predicted
 186                       }
 187
 188 isGoal :: SearchState s => Agendum s -> BlueprintContext Bool
 189 isGoal agendum =
 190   do timeLimit <- asks getTimeLimit
 191      return $ Q.length (agendum ^. trail) == timeLimit
 192
 193 -- Parse the input file
 194
 195 blueprintsP :: Parser [(Int, Blueprint)]
 196 blueprintP :: Parser (Int, Blueprint)
 197 robotP :: Parser (Resource, Collection)
 198 requirementsP :: Parser Collection
 199 requirementP :: Parser (Resource, Int)
 200 resourceP, oreP, clayP, obsidianP, geodeP :: Parser Resource
 201
 202 blueprintsP = blueprintP `sepBy` endOfLine
 203 blueprintP = blueprintify <$> (("Blueprint " *> decimal) <* ": ") <*> (robotP `sepBy` ". ") <* "."
 204   where blueprintify n robots =
 205           (n, M.fromList robots)
 206 robotP = (,) <$> ("Each " *> resourceP) <*> (" robot costs " *> requirementsP)
 207
 208 requirementsP = MS.fromOccurList <$> (requirementP `sepBy` " and ")
 209
 210 requirementP = (flip (,)) <$> (decimal <* " ") <*> resourceP
 211
 212 resourceP = oreP <|> clayP <|> obsidianP <|> geodeP
 213 oreP = Ore <$ "ore"
 214 clayP = Clay <$ "clay"
 215 obsidianP = Obsidian <$ "obsidian"
 216 geodeP = Geode <$ "geode"
 217
 218 successfulParse :: Text -> [(Int, Blueprint)]
 219 successfulParse input =
 220   case parseOnly blueprintsP input of
 221     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 222     Right blueprints -> blueprints