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