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