advent19/MainOriginal.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 -- part1 blueprints = pResults
 113     where results = [ (n, _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
 114                     | (n, blueprint) <- blueprints ] :: [(Int, SingleSearchState)]
 115           pResults = parMap rdeepseq id results
 116           robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
 117 -- part1 blueprints = sum [n * (MS.occur Geode (r ^. resources)) | (n, r) <- pResults]
 118 --     where -- results = [ (n, _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
 119 --                     -- | (n, blueprint) <- blueprints ] :: [(Int, SingleSearchState)]
 120 --           -- pResults = parMap rdeepseq id results
 121 --           -- pResults = (fmap runABlueprint blueprints) `using` parList rdeepseq
 122 --           pResults = (fmap runABlueprint blueprints) `using` (parList rdeepseq)
 123 --           runABlueprint (n, blueprint) = (n, _current $ fromJust $
 124 --                                                 runReader searchSpace (TimedBlueprint blueprint 24 (robotLimits blueprint)) )
 125 --           robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
 126
 127 part2 :: [(Int, Blueprint)] -> Int
 128 part2 blueprints = product [MS.occur Geode (r ^. resources) | r <- pResults]
 129     where results = [ _current $ fromJust $ runReader searchSpace (TimedBlueprint blueprint 32 (robotLimits blueprint))
 130                     | (_, blueprint) <- (take 3 blueprints) ] :: [SingleSearchState]
 131           pResults = parMap rdeepseq id results
 132           robotLimits bp = M.foldl' MS.maxUnion MS.empty bp
 133
 134 searchSpace :: SearchState s => BlueprintContext (Maybe (Agendum s))
 135 searchSpace =
 136     do agenda <- initAgenda
 137        -- aStar agenda S.empty
 138        res <- aStar agenda S.empty
 139        return (res `seq` res)
 140
 141 initAgenda :: SearchState s => BlueprintContext (Agenda s)
 142 initAgenda =
 143     do let startState = emptySearchState
 144        b <- estimateBenefit startState 0
 145        return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
 146
 147 aStar ::  SearchState s => Agenda s -> ExploredStates s -> BlueprintContext (Maybe (Agendum s))
 148 aStar agenda closed
 149     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 150     -- | 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
 151     | P.null agenda = return Nothing
 152     | otherwise =
 153         do  let (_, currentAgendum) = P.findMax agenda
 154             let reached = currentAgendum ^. current
 155             nexts <- candidates currentAgendum closed
 156             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
 157             -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
 158             -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
 159             reachedGoal <- isGoal currentAgendum
 160             let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
 161             if reachedGoal
 162             then return (Just currentAgendum)
 163             else if (cl `S.member` closed)
 164                  then aStar (P.deleteMax agenda) closed
 165                  -- else aStar newAgenda (S.insert cl closed)
 166                  else aStar newAgenda (S.insert cl closed)
 167
 168 candidates :: SearchState s => Agendum s -> ExploredStates s -> BlueprintContext (Q.Seq (Agendum s))
 169 candidates agendum closed =
 170     do  let candidate = agendum ^. current
 171         let previous = agendum ^. trail
 172         let prevBenefit = agendum ^. trailBenefit
 173         succs <- successors candidate
 174         succAgs <- mapM (makeAgendum previous prevBenefit) succs
 175         let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
 176         return nonloops
 177
 178 makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> BlueprintContext (Agendum s)
 179 makeAgendum previous prevBenefit newState =
 180     do predicted <- estimateBenefit newState (Q.length previous)
 181        let newTrail = previous |> newState
 182        -- let incurred = geodesFound newState
 183        let incurred = 0
 184        return Agendum { _current = newState
 185                       , _trail = newTrail
 186                       , _trailBenefit = incurred
 187                       , _benefit = incurred + predicted
 188                       }
 189
 190 isGoal :: SearchState s => Agendum s -> BlueprintContext Bool
 191 isGoal agendum =
 192   do timeLimit <- asks getTimeLimit
 193      return $ Q.length (agendum ^. trail) == timeLimit
 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