advent16/MainEstSort.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
   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.Sequence ((|>), Seq((:|>)), ViewR ((:>)))
  16 import Data.Sequence ( (|>), Seq((:|>)) )
  17 import Data.List
  18 import Data.List.Split (chunksOf)
  19 import Data.Ord
  20 import Control.Monad.Reader
  21 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
  22
  23
  24 type RoomID = String
  25
  26 data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
  27     deriving (Eq, Show, Ord)
  28 makeLenses ''Tunnel
  29
  30 data Room = Room
  31     { _flowRate :: Int
  32     , _tunnels :: S.Set Tunnel
  33     } deriving (Eq, Show, Ord)
  34 makeLenses ''Room
  35
  36 type Cave = M.Map RoomID Room
  37 data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int }
  38
  39 type CaveContext = Reader TimedCave
  40
  41 data SingleSearchState = SingleSearchState
  42     { _currentRoom :: RoomID
  43     , _currentTime :: Int
  44     , _sOpenValves :: S.Set RoomID
  45     } deriving (Eq, Show, Ord)
  46 makeLenses ''SingleSearchState
  47
  48 data DoubleSearchState = DoubleSearchState
  49     { _personRoom :: RoomID
  50     , _personTime :: Int
  51     , _elephantRoom :: RoomID
  52     , _elephantTime :: Int
  53     , _dOpenValves :: S.Set RoomID
  54     } deriving (Eq, Show, Ord)
  55 makeLenses ''DoubleSearchState
  56
  57 data Agendum s =
  58     Agendum { _current :: s
  59             , _trail :: Q.Seq s
  60             , _trailBenefit :: Int
  61             , _benefit :: Int
  62             } deriving (Show, Eq, Ord)
  63 makeLenses ''Agendum
  64
  65 type Agenda s = P.MaxPQueue Int (Agendum s)
  66
  67 -- state, total flowed so far
  68 type ExploredStates s = S.Set (s, Int)
  69
  70
  71 class (Eq s, Ord s, Show s) => SearchState s where
  72   emptySearchState :: RoomID -> s
  73   currentFlow :: s -> CaveContext Int
  74   timeOf :: s -> Int
  75   successors :: s -> CaveContext (Q.Seq s)
  76   -- estimateBenefit :: s -> Int -> CaveContext Int
  77   estimateBenefit :: s -> CaveContext Int
  78
  79 instance SearchState SingleSearchState where
  80   emptySearchState startID = SingleSearchState
  81         { _currentRoom = startID
  82         , _currentTime = 0
  83         , _sOpenValves = S.empty
  84         }
  85
  86   currentFlow state =
  87     do cave <- asks getCave
  88        let valves = state ^. sOpenValves
  89        let presentRooms = cave `M.restrictKeys` valves
  90        return $ sumOf (folded . flowRate) presentRooms
  91
  92   timeOf state = state ^. currentTime
  93
  94   successors state =
  95     do isFF <- isFullFlow state
  96        -- cave <- asks getCave
  97        timeLimit <- asks getTimeLimit
  98        let here = state ^. currentRoom
  99        let opened = state ^. sOpenValves
 100        let now = state ^. currentTime
 101        succs <- agentSuccessor now opened now here
 102        let succStates = Q.fromList succs
 103        if isFF || (Q.null succStates)
 104        then return $ Q.singleton (state & currentTime .~ timeLimit)
 105        else return succStates
 106
 107   estimateBenefit here =
 108     do cave <- asks getCave
 109        timeLimit <- asks getTimeLimit
 110        let timeRemaining = timeLimit - (timeOf here)
 111        cf <- currentFlow here
 112        let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
 113        let sortedClosedValves = sortOn Down closedValves
 114        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
 115        return $ (cf * timeRemaining) + otherValveFlows
 116
 117
 118 instance SearchState DoubleSearchState where
 119   emptySearchState startID = DoubleSearchState
 120       { _personRoom = startID
 121       , _personTime = 0
 122       , _elephantRoom = startID
 123       , _elephantTime = 0
 124       , _dOpenValves = S.empty
 125       }
 126
 127   currentFlow state =
 128     do cave <- asks getCave
 129        let valves = S.toList $ state ^. dOpenValves
 130        return $ sum $ fmap (\v -> (cave ! v) ^. flowRate) valves
 131        -- let presentRooms = cave `M.restrictKeys` valves
 132        -- return $ sumOf (folded . flowRate) presentRooms
 133
 134   timeOf state = min (state ^. personTime) (state ^. elephantTime)
 135
 136   successors state =
 137     do isFF <- isFullFlow state
 138        -- cave <- asks getCave
 139        timeLimit <- asks getTimeLimit
 140        let opened = state ^. dOpenValves
 141        let pNow = state ^. personTime
 142        let eNow = state ^. elephantTime
 143        let now = min pNow eNow
 144        let pHere = state ^. personRoom
 145        let eHere = state ^. elephantRoom
 146        pNexts <- agentSuccessor now opened pNow pHere
 147        eNexts <- agentSuccessor now opened eNow eHere
 148        let nexts =  [ state & personRoom .~ (p ^. currentRoom)
 149                             & personTime .~ (p ^. currentTime)
 150                             & elephantRoom .~ (e ^. currentRoom)
 151                             & elephantTime .~ (e ^. currentTime)
 152                             & dOpenValves %~ (S.union (p ^. sOpenValves) . S.union (e ^. sOpenValves))
 153                     | p <- pNexts
 154                     , e <- eNexts
 155                     ]
 156        let dedups = if pNow == eNow && pHere == eHere
 157                     then filter (\s -> (s ^. personRoom) < (s ^. elephantRoom)) nexts
 158                     else nexts
 159        -- let succStates = trace ("Succs: in " ++ (show state) ++ " out " ++ (show dedups)) (Q.fromList dedups)
 160        let succStates = Q.fromList dedups
 161        if isFF || (Q.null succStates)
 162        then return $ Q.singleton (state & personTime .~ timeLimit & elephantTime .~ timeLimit)
 163        else return succStates
 164
 165   estimateBenefit here =
 166     do cave <- asks getCave
 167        timeLimit <- asks getTimeLimit
 168        let timeRemaining = timeLimit - (timeOf here)
 169        cf <- currentFlow here
 170        let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
 171        let sortedClosedValves = fmap sum $ chunksOf 2 $  {-# SCC estSort #-} sortOn Down closedValves
 172        -- let sortedClosedValves = fmap sum $ chunksOf 2 $ reverse $ sort closedValves -- no significant improvement
 173        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
 174        -- let otherValveFlows = timeRemaining * (sum closedValves) -- 8 minute runtime rather than 1:50
 175        return $ (cf * timeRemaining) + otherValveFlows
 176
 177
 178 main :: IO ()
 179 main =
 180   do  dataFileName <- getDataFileName
 181       text <- TIO.readFile dataFileName
 182       let expandedCave = successfulParse text
 183       -- print cave
 184       -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
 185       -- print $ compress cave
 186       let cave = compress expandedCave
 187       print $ part1 cave
 188       print $ part2 cave
 189
 190 -- part1 :: Cave -> Maybe (Agendum SingleSearchState)
 191 -- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
 192
 193 -- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
 194 -- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
 195
 196 part1, part2 :: Cave -> Int
 197 -- part1 :: Cave -> Int
 198 part1 cave = maybe 0 _benefit result
 199     where result = runReader (searchCave "AA") (TimedCave cave 30) :: Maybe (Agendum SingleSearchState)
 200 part2 cave = maybe 0 _benefit result
 201     where result = runReader (searchCave "AA") (TimedCave cave 26) :: Maybe (Agendum DoubleSearchState)
 202
 203 searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
 204 searchCave startRoom =
 205     do agenda <- initAgenda startRoom
 206        aStar agenda S.empty
 207
 208 initAgenda :: SearchState s => String -> CaveContext (Agenda s)
 209 initAgenda startID =
 210     do let startState = emptySearchState startID
 211        b <- estimateBenefit startState
 212        return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
 213
 214 aStar ::  SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
 215 aStar agenda closed
 216     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
 217     -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
 218     -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
 219     | P.null agenda = return Nothing
 220     | otherwise =
 221         do  let (_, currentAgendum) = P.findMax agenda
 222             let reached = currentAgendum ^. current
 223             nexts <- candidates currentAgendum closed
 224             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
 225             -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
 226             -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
 227             -- let beamAgenda = P.fromDescList $ P.take 1000 newAgenda -- agenda beam width
 228             reachedGoal <- isGoal currentAgendum
 229             -- let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
 230             let cl = (reached, currentAgendum ^. trailBenefit)
 231             if reachedGoal
 232             then return (Just currentAgendum)
 233             else if (cl `S.member` closed)
 234                  then aStar (P.deleteMax agenda) closed
 235                  else aStar newAgenda (S.insert cl closed)
 236                  -- else aStar beamAgenda (S.insert cl closed)
 237
 238
 239 candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
 240 candidates agendum closed =
 241     do  let candidate = agendum ^. current
 242         let previous = agendum ^. trail
 243         let prevBenefit = agendum ^. trailBenefit
 244         succs <- successors candidate
 245         succAgs <- mapM (makeAgendum previous prevBenefit) succs
 246         -- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
 247         let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
 248         return nonloops
 249
 250
 251 agentSuccessor :: Int -> S.Set RoomID -> Int -> RoomID -> CaveContext [SingleSearchState]
 252 agentSuccessor now opened aTime here
 253   | aTime /= now = return [SingleSearchState { _currentRoom = here, _currentTime = aTime, _sOpenValves = opened }]
 254   | otherwise =
 255       do  cave <- asks getCave
 256           timeLimit <- asks getTimeLimit
 257           let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
 258           let moves = [ SingleSearchState
 259                             { _currentRoom = (t ^. tunnelTo)
 260                             , _currentTime = now + (t ^. tunnelLength)
 261                             , _sOpenValves = opened
 262                             }
 263                       | t <- remaining
 264                       , now + (t ^. tunnelLength) <= timeLimit
 265                       ]
 266           let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0
 267                       then [SingleSearchState { _currentRoom = here, _currentTime = aTime + 1, _sOpenValves = S.insert here opened }]
 268                       else []
 269           -- let nexts = moves ++ opens
 270           let nexts = if null opens then moves else opens
 271           let nexts' =  if null nexts
 272                         then [ SingleSearchState
 273                                 { _currentRoom = here
 274                                 , _currentTime = timeLimit
 275                                 , _sOpenValves = opened
 276                                 } ]
 277                         else nexts
 278           return nexts'
 279
 280 makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
 281 makeAgendum previous prevBenefit newState =
 282     do predicted <- estimateBenefit newState --  (Q.length previous)
 283        -- cf <- currentFlow newState
 284        oldFlow <- lastFlow previous (timeOf newState)
 285        let newTrail = previous |> newState
 286        let incurred = prevBenefit + oldFlow
 287        return Agendum { _current = newState
 288                       , _trail = newTrail
 289                       , _trailBenefit = incurred
 290                       , _benefit = incurred + predicted
 291                       }
 292
 293 lastFlow :: SearchState s => Q.Seq s -> Int -> CaveContext Int
 294 lastFlow Q.Empty _ = return 0
 295 lastFlow (_ :|> previous) newTime =
 296   do cf <- currentFlow previous
 297      let dt = newTime - (timeOf previous)
 298      return (cf * dt)
 299
 300 isGoal :: SearchState s => Agendum s -> CaveContext Bool
 301 isGoal agendum =
 302   do timeLimit <- asks getTimeLimit
 303      let s = agendum ^. current
 304      return $ (timeOf s) == timeLimit
 305
 306 isFullFlow :: SearchState s => s -> CaveContext Bool
 307 isFullFlow state =
 308   do cave <- asks getCave
 309      cf <- currentFlow state
 310      let ff = sumOf (folded . flowRate) cave
 311      return (cf == ff)
 312
 313 compress :: Cave -> Cave
 314 compress cave = M.mapWithKey (compressRoom cave) cave
 315
 316 compressRoom :: Cave -> RoomID -> Room -> Room
 317 compressRoom cave here room = room & tunnels .~ t'
 318   where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
 319
 320 reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
 321 reachableFrom _ [] _ routes = routes
 322 reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
 323   | here `S.member` found = reachableFrom cave boundary found routes
 324   | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
 325   where exits = (cave ! here) ^. tunnels
 326         exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
 327         legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
 328         routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
 329                   then routes
 330                   else S.insert tunnel routes
 331
 332
 333 -- Parse the input file
 334
 335 caveP :: Parser Cave
 336 valveP :: Parser (RoomID, Room)
 337 roomP :: Parser Room
 338 tunnelsP :: Parser (S.Set Tunnel)
 339 tunnelTextP :: Parser Text
 340
 341 caveP = M.fromList <$> valveP `sepBy` endOfLine
 342 valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
 343 roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
 344     -- where roomify v ts = Room {flowRate = v, tunnels = ts }
 345 tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
 346 tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
 347
 348 successfulParse :: Text -> Cave
 349 successfulParse input =
 350   case parseOnly caveP input of
 351     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
 352     Right cave -> cave