Reworking day 16
authorNeil Smith <NeilNjae@users.noreply.github.com>
Tue, 18 Jul 2023 14:22:56 +0000 (15:22 +0100)
committerNeil Smith <NeilNjae@users.noreply.github.com>
Fri, 21 Jul 2023 13:53:18 +0000 (14:53 +0100)
advent-of-code22.cabal
advent16/Main.hs
advent16/MainBeam.hs [new file with mode: 0644]
advent16/MainCustomClosed.hs [new file with mode: 0644]
advent16/MainEstSort.hs [new file with mode: 0644]
advent16/MainOriginal.hs [new file with mode: 0644]
advent16/MainOriginalNoBeam.hs [new file with mode: 0644]
advent16/MainSPar.hs [new file with mode: 0644]
advent16/MainSubsets.hs [new file with mode: 0644]
advent16/a16-solution.dot.png [new file with mode: 0644]
advent16/a16.dot [new file with mode: 0644]

index 5a5315da9d503c3ac3eb6a676e9454c160f7885d..8a1b109a3c095313d03dd5acb3ff1e13ed866435 100644 (file)
@@ -206,6 +206,63 @@ executable advent15prof
                        -eventlog
                        -rtsopts "-with-rtsopts=-N -p -s -hT -ls"
 
+executable advent16original
+  import: common-extensions, build-directives
+  main-is: advent16/MainOriginal.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split
+
+executable advent16originalnobeam
+  import: common-extensions, build-directives
+  main-is: advent16/MainOriginalNoBeam.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split
+
+executable advent16sort
+  import: common-extensions, build-directives
+  main-is: advent16/MainEstSort.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split
+
+executable advent16beam
+  import: common-extensions, build-directives
+  main-is: advent16/MainBeam.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split
+
+executable advent16customclosed
+  import: common-extensions, build-directives
+  main-is: advent16/MainCustomClosed.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split
+
+executable advent16spar
+  import: common-extensions, build-directives
+  main-is: advent16/MainSPar.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split, parallel, deepseq
+
+executable advent16sparprof
+  import: common-extensions, build-directives
+  main-is: advent16/MainSPar.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split, parallel, deepseq
+  ghc-options:         -O2 
+                       -Wall 
+                       -threaded 
+                       -eventlog
+                       -fprof-auto
+                       -rtsopts "-with-rtsopts=-N -p -s -hT -ls"
+
+executable advent16subsets
+  import: common-extensions, build-directives
+  main-is: advent16/MainSubsets.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split
+
+executable advent16subsetsprof
+  import: common-extensions, build-directives
+  main-is: advent16/MainSubsets.hs
+  build-depends: text, attoparsec, containers, pqueue, mtl, lens, split
+  ghc-options:         -O2 
+                       -Wall 
+                       -threaded 
+                       -eventlog
+                       -fprof-auto
+                       -rtsopts "-with-rtsopts=-N -p -s -hT -ls"
+
 executable advent16
   import: common-extensions, build-directives
   main-is: advent16/Main.hs
index 7db601d30bafb2aeb9906e281296994366cde433..30d0843b1c2239add994e76b6bc19c60605aa8f4 100644 (file)
@@ -1,6 +1,6 @@
 -- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
 
--- import Debug.Trace
+import Debug.Trace
 
 import AoC
 import Data.Text (Text)
@@ -12,39 +12,44 @@ import qualified Data.Set as S
 import qualified Data.Sequence as Q
 import qualified Data.Map.Strict as M
 import Data.Map.Strict ((!))
-import Data.Sequence ((|>)) 
+-- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>))) 
+import Data.Sequence ( (|>), Seq((:|>)) )
 import Data.List
 import Data.List.Split (chunksOf)
 import Data.Ord
 import Control.Monad.Reader
 import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
 
--- pattern Empty   <- (Q.viewl -> Q.EmptyL)  where Empty = Q.empty
--- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<)  = (Q.<|) 
--- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>)  = (Q.|>) 
 
 type RoomID = String
 
+data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
+    deriving (Eq, Show, Ord)
+makeLenses ''Tunnel
+
 data Room = Room 
     { _flowRate :: Int
-    , _tunnels :: [RoomID]
+    , _tunnels :: S.Set Tunnel
     } deriving (Eq, Show, Ord)
 makeLenses ''Room
 
 type Cave = M.Map RoomID Room
-data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int}
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
 
 type CaveContext = Reader TimedCave
 
 data SingleSearchState = SingleSearchState
     { _currentRoom :: RoomID
+    , _currentTime :: Int
     , _sOpenValves :: S.Set RoomID
     } deriving (Eq, Show, Ord)
 makeLenses ''SingleSearchState
 
 data DoubleSearchState = DoubleSearchState
     { _personRoom :: RoomID
+    , _personTime :: Int
     , _elephantRoom :: RoomID
+    , _elephantTime :: Int
     , _dOpenValves :: S.Set RoomID
     } deriving (Eq, Show, Ord)
 makeLenses ''DoubleSearchState
@@ -59,17 +64,24 @@ makeLenses ''Agendum
 
 type Agenda s = P.MaxPQueue Int (Agendum s)
 
-type ExploredStates s = S.Set (s, Int, Int)
+-- state, total flowed so far
+type ExploredStates s = S.Set (s, Int)
 
 
 class (Eq s, Ord s, Show s) => SearchState s where
   emptySearchState :: RoomID -> s
   currentFlow :: s -> CaveContext Int
+  timeOf :: s -> Int
   successors :: s -> CaveContext (Q.Seq s)
-  estimateBenefit :: s -> Int -> CaveContext Int
+  -- estimateBenefit :: s -> Int -> CaveContext Int
+  estimateBenefit :: s -> CaveContext Int
 
 instance SearchState SingleSearchState where
-  emptySearchState startID = SingleSearchState { _currentRoom = startID, _sOpenValves = S.empty }
+  emptySearchState startID = SingleSearchState 
+        { _currentRoom = startID
+        , _currentTime = 0
+        , _sOpenValves = S.empty 
+        }
 
   currentFlow state =
     do cave <- asks getCave
@@ -77,29 +89,31 @@ instance SearchState SingleSearchState where
        let presentRooms = cave `M.restrictKeys` valves
        return $ sumOf (folded . flowRate) presentRooms
 
+  timeOf state = state ^. currentTime
+
   successors state = 
     do isFF <- isFullFlow state
+       -- cave <- asks getCave
+       timeLimit <- asks getTimeLimit
        let here = state ^. currentRoom
        let opened = state ^. sOpenValves
-       succPairs <- personSuccessor here opened
-       let succStates = 
-              [ SingleSearchState 
-                  { _currentRoom = r
-                  , _sOpenValves = o
-                  }
-              | (r, o) <- succPairs
-              ]
-       if isFF
-       then return $ Q.singleton state
-       else return $ Q.fromList succStates
-
-  estimateBenefit here timeElapsed = 
+       let now = state ^. currentTime
+       succs <- agentSuccessor now opened now here
+       let succStates = Q.fromList succs
+       if isFF || (Q.null succStates)
+       then return $ Q.singleton (state & currentTime .~ timeLimit)
+       else return succStates
+
+  estimateBenefit here =  
     do cave <- asks getCave
        timeLimit <- asks getTimeLimit
-       let timeRemaining = timeLimit - (timeElapsed + 2)
+       let timeRemaining = timeLimit - (timeOf here)
        cf <- currentFlow here
-       let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
-       let sortedClosedValves = sortOn Down closedValves
+       -- let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
+       -- let sortedClosedValves = sortOn Down closedValves
+       sortedValves <- asks getSortedRooms
+       let opened = here ^. sOpenValves
+       let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
        return $ (cf * timeRemaining) + otherValveFlows
 
@@ -107,44 +121,64 @@ instance SearchState SingleSearchState where
 instance SearchState DoubleSearchState where
   emptySearchState startID = DoubleSearchState 
       { _personRoom = startID
+      , _personTime = 0
       , _elephantRoom = startID
+      , _elephantTime = 0
       , _dOpenValves = S.empty 
       }
 
   currentFlow state =
     do cave <- asks getCave
-       let valves = state ^. dOpenValves
-       let presentRooms = cave `M.restrictKeys` valves
-       return $ sumOf (folded . flowRate) presentRooms
+       let valves = S.toList $ state ^. dOpenValves
+       return $ sum $ fmap (\v -> (cave ! v) ^. flowRate) valves
+       -- let presentRooms = cave `M.restrictKeys` valves
+       -- return $ sumOf (folded . flowRate) presentRooms
+
+  timeOf state = min (state ^. personTime) (state ^. elephantTime)
 
   successors state = 
     do isFF <- isFullFlow state
+       -- cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let opened = state ^. dOpenValves
+       let pNow = state ^. personTime
+       let eNow = state ^. elephantTime
+       let now = min pNow eNow
        let pHere = state ^. personRoom
        let eHere = state ^. elephantRoom
-       let opened = state ^. dOpenValves
-       pSuccPairs <- personSuccessor pHere opened
-       eSuccPairs <- personSuccessor eHere opened
-       let succStates = 
-              [ DoubleSearchState 
-                  { _personRoom = p
-                  , _elephantRoom = e
-                  , _dOpenValves = S.union po eo
-                  }
-              | (p, po) <- pSuccPairs
-              , (e, eo) <- eSuccPairs
-              ]
-       if isFF
-       then return $ Q.singleton state
-       else return $ Q.fromList succStates
-
-  estimateBenefit here timeElapsed = 
+       pNexts <- agentSuccessor now opened pNow pHere
+       eNexts <- agentSuccessor now opened eNow eHere
+       let nexts =  [ state & personRoom .~ (p ^. currentRoom)
+                            & personTime .~ (p ^. currentTime)
+                            & elephantRoom .~ (e ^. currentRoom)
+                            & elephantTime .~ (e ^. currentTime)
+                            & dOpenValves %~ (S.union (p ^. sOpenValves) . S.union (e ^. sOpenValves))
+                    | p <- pNexts
+                    , e <- eNexts
+                    ]
+       let dedups = if pNow == eNow && pHere == eHere
+                    then filter (\s -> (s ^. personRoom) < (s ^. elephantRoom)) nexts
+                    -- else nexts
+                    else filter (\s -> (s ^. personRoom) /= (s ^. elephantRoom)) nexts
+       -- let succStates = trace ("Succs: in " ++ (show state) ++ " out " ++ (show dedups)) (Q.fromList dedups)
+       let succStates = Q.fromList dedups
+       if isFF || (Q.null succStates)
+       then return $ Q.singleton (state & personTime .~ timeLimit & elephantTime .~ timeLimit)
+       else return succStates
+
+  estimateBenefit here =  
     do cave <- asks getCave
        timeLimit <- asks getTimeLimit
-       let timeRemaining = timeLimit - (timeElapsed + 2)
+       let timeRemaining = timeLimit - (timeOf here)
        cf <- currentFlow here
-       let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
-       let sortedClosedValves = fmap sum $ chunksOf 2 $ sortOn Down closedValves
+       -- let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
+       -- let sortedClosedValves = fmap sum $ chunksOf 2 $  {-# SCC estSort #-} sortOn Down closedValves
+       -- let sortedClosedValves = fmap sum $ chunksOf 2 $ reverse $ sort closedValves -- no significant improvement
+       sortedValves <- asks getSortedRooms
+       let opened = here ^. dOpenValves
+       let sortedClosedValves = fmap sum $ chunksOf 2 $ [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
        let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       -- let otherValveFlows = timeRemaining * (sum closedValves) -- 8 minute runtime rather than 1:50
        return $ (cf * timeRemaining) + otherValveFlows
 
 
@@ -152,8 +186,11 @@ main :: IO ()
 main = 
   do  dataFileName <- getDataFileName
       text <- TIO.readFile dataFileName
-      let cave = successfulParse text
+      let expandedCave = successfulParse text
       -- print cave
+      -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
+      -- print $ compress cave
+      let cave = compress expandedCave
       print $ part1 cave
       print $ part2 cave
 
@@ -164,10 +201,14 @@ main =
 -- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
 
 part1, part2 :: Cave -> Int
+-- part1 :: Cave -> Int
 part1 cave = maybe 0 _benefit result
-    where result = runReader (searchCave "AA") (TimedCave cave 30) :: Maybe (Agendum SingleSearchState)
+    where result = runReader (searchCave "AA") (TimedCave cave 30 sortedRooms) :: Maybe (Agendum SingleSearchState)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
 part2 cave = maybe 0 _benefit result
-    where result = runReader (searchCave "AA") (TimedCave cave 26) :: Maybe (Agendum DoubleSearchState)
+    where result = runReader (searchCave "AA") (TimedCave cave 26 sortedRooms) :: Maybe (Agendum DoubleSearchState)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+          -- sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys cave
 
 searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
 searchCave startRoom = 
@@ -177,13 +218,15 @@ searchCave startRoom =
 initAgenda :: SearchState s => String -> CaveContext (Agenda s)
 initAgenda startID = 
     do let startState = emptySearchState startID
-       b <- estimateBenefit startState 0
+       b <- estimateBenefit startState 
        return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
 
 aStar ::  SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
 aStar agenda closed 
     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
-    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : len " ++ (show $ Q.length $ _trail $ snd $ P.findMax agenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ " : trail " ++ (show $ _trail $ snd $ P.findMax agenda) ++ " : closed " ++ (show closed)) False = undefined
+    -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
     | P.null agenda = return Nothing
     | otherwise = 
         do  let (_, currentAgendum) = P.findMax agenda
@@ -191,15 +234,17 @@ aStar agenda closed
             nexts <- candidates currentAgendum closed
             let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
             -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
-            let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+            -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+            -- let beamAgenda = P.fromDescList $ P.take 1000 newAgenda -- agenda beam width
             reachedGoal <- isGoal currentAgendum
-            let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+            -- let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+            let cl = (reached, currentAgendum ^. trailBenefit)
             if reachedGoal
             then return (Just currentAgendum)
             else if (cl `S.member` closed)
                  then aStar (P.deleteMax agenda) closed
-                 -- else aStar newAgenda (S.insert cl closed)
-                 else aStar beamAgenda (S.insert cl closed)
+                 else aStar newAgenda (S.insert cl closed)
+                 -- else aStar beamAgenda (S.insert cl closed)
 
 
 candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
@@ -209,41 +254,65 @@ candidates agendum closed =
         let prevBenefit = agendum ^. trailBenefit
         succs <- successors candidate
         succAgs <- mapM (makeAgendum previous prevBenefit) succs
-        let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+        -- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+        let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
         return nonloops
 
-personSuccessor, openValveSuccessor, walkSuccessor ::  RoomID -> S.Set RoomID -> CaveContext [(RoomID, S.Set RoomID)]
-personSuccessor here opened =
-  do ovs <- openValveSuccessor here opened
-     ws <- walkSuccessor here opened
-     return (ovs ++ ws)
 
-openValveSuccessor here opened
-  | here `S.member` opened = return []
-  | otherwise = return [(here, S.insert here opened)]
-
-walkSuccessor here opened = 
-  do cave <- asks getCave
-     let neighbours = (cave ! here) ^. tunnels
-     return [(n, opened) | n <- neighbours]
+agentSuccessor :: Int -> S.Set RoomID -> Int -> RoomID -> CaveContext [SingleSearchState]
+agentSuccessor now opened aTime here 
+  | aTime /= now = return [SingleSearchState { _currentRoom = here, _currentTime = aTime, _sOpenValves = opened }]
+  | otherwise = 
+      do  cave <- asks getCave 
+          timeLimit <- asks getTimeLimit
+          let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
+          let moves = [ SingleSearchState 
+                            { _currentRoom = (t ^. tunnelTo)
+                            , _currentTime = now + (t ^. tunnelLength)
+                            , _sOpenValves = opened
+                            }
+                      | t <- remaining
+                      , now + (t ^. tunnelLength) <= timeLimit
+                      ]
+          let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0
+                      then [SingleSearchState { _currentRoom = here, _currentTime = aTime + 1, _sOpenValves = S.insert here opened }]
+                      else []
+          -- let nexts = moves ++ opens
+          let nexts = if null opens then moves else opens
+          let nexts' =  if null nexts
+                        then [ SingleSearchState 
+                                { _currentRoom = here
+                                , _currentTime = timeLimit
+                                , _sOpenValves = opened
+                                } ]
+                        else nexts
+          return nexts'
 
 makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
 makeAgendum previous prevBenefit newState = 
-    do predicted <- estimateBenefit newState (Q.length previous)
-       cf <- currentFlow newState
+    do predicted <- estimateBenefit newState --  (Q.length previous)
+       -- cf <- currentFlow newState
+       oldFlow <- lastFlow previous (timeOf newState)
        let newTrail = previous |> newState
-       let incurred = prevBenefit + cf
+       let incurred = prevBenefit + oldFlow
        return Agendum { _current = newState
                       , _trail = newTrail
                       , _trailBenefit = incurred
                       , _benefit = incurred + predicted
                       }
 
+lastFlow :: SearchState s => Q.Seq s -> Int -> CaveContext Int
+lastFlow Q.Empty _ = return 0
+lastFlow (_ :|> previous) newTime = 
+  do cf <- currentFlow previous
+     let dt = newTime - (timeOf previous)
+     return (cf * dt)
 
 isGoal :: SearchState s => Agendum s -> CaveContext Bool
 isGoal agendum = 
   do timeLimit <- asks getTimeLimit
-     return $ Q.length (agendum ^. trail) == (timeLimit - 1)
+     let s = agendum ^. current
+     return $ (timeOf s) == timeLimit
 
 isFullFlow :: SearchState s => s -> CaveContext Bool
 isFullFlow state = 
@@ -252,21 +321,40 @@ isFullFlow state =
      let ff = sumOf (folded . flowRate) cave
      return (cf == ff)
 
+compress :: Cave -> Cave
+compress cave = M.mapWithKey (compressRoom cave) cave
+
+compressRoom :: Cave -> RoomID -> Room -> Room
+compressRoom cave here room = room & tunnels .~ t'
+  where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
+
+reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
+reachableFrom _ [] _ routes = routes
+reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
+  | here `S.member` found = reachableFrom cave boundary found routes
+  | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
+  where exits = (cave ! here) ^. tunnels
+        exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
+        legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
+        routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
+                  then routes
+                  else S.insert tunnel routes
+
 
 -- Parse the input file
 
 caveP :: Parser Cave
 valveP :: Parser (RoomID, Room)
 roomP :: Parser Room
-tunnelsP :: Parser [RoomID]
-turnnelTextP :: Parser Text
+tunnelsP :: Parser (S.Set Tunnel)
+tunnelTextP :: Parser Text
 
 caveP = M.fromList <$> valveP `sepBy` endOfLine
 valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
-roomP = roomify <$> (" has flow rate=" *> decimal) <*> (turnnelTextP *> tunnelsP)
-    where roomify v ts = Room {_flowRate = v, _tunnels = ts }
-tunnelsP = (many1 letter) `sepBy` ", "
-turnnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+    -- where roomify v ts = Room {flowRate = v, tunnels = ts }
+tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
 
 successfulParse :: Text -> Cave
 successfulParse input = 
diff --git a/advent16/MainBeam.hs b/advent16/MainBeam.hs
new file mode 100644 (file)
index 0000000..f82731c
--- /dev/null
@@ -0,0 +1,362 @@
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+-- import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+-- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>))) 
+import Data.Sequence ( (|>), Seq((:|>)) )
+import Data.List
+import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+
+
+type RoomID = String
+
+data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
+    deriving (Eq, Show, Ord)
+makeLenses ''Tunnel
+
+data Room = Room 
+    { _flowRate :: Int
+    , _tunnels :: S.Set Tunnel
+    } deriving (Eq, Show, Ord)
+makeLenses ''Room
+
+type Cave = M.Map RoomID Room
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
+
+type CaveContext = Reader TimedCave
+
+data SingleSearchState = SingleSearchState
+    { _currentRoom :: RoomID
+    , _currentTime :: Int
+    , _sOpenValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''SingleSearchState
+
+data DoubleSearchState = DoubleSearchState
+    { _personRoom :: RoomID
+    , _personTime :: Int
+    , _elephantRoom :: RoomID
+    , _elephantTime :: Int
+    , _dOpenValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''DoubleSearchState
+
+data Agendum s = 
+    Agendum { _current :: s
+            , _trail :: Q.Seq s
+            , _trailBenefit :: Int
+            , _benefit :: Int
+            } deriving (Show, Eq, Ord)
+makeLenses ''Agendum   
+
+type Agenda s = P.MaxPQueue Int (Agendum s)
+
+-- state, total flowed so far
+type ExploredStates s = S.Set (s, Int)
+
+
+class (Eq s, Ord s, Show s) => SearchState s where
+  emptySearchState :: RoomID -> s
+  currentFlow :: s -> CaveContext Int
+  timeOf :: s -> Int
+  successors :: s -> CaveContext (Q.Seq s)
+  -- estimateBenefit :: s -> Int -> CaveContext Int
+  estimateBenefit :: s -> CaveContext Int
+
+instance SearchState SingleSearchState where
+  emptySearchState startID = SingleSearchState 
+        { _currentRoom = startID
+        , _currentTime = 0
+        , _sOpenValves = S.empty 
+        }
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. sOpenValves
+       let presentRooms = cave `M.restrictKeys` valves
+       return $ sumOf (folded . flowRate) presentRooms
+
+  timeOf state = state ^. currentTime
+
+  successors state = 
+    do isFF <- isFullFlow state
+       -- cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let here = state ^. currentRoom
+       let opened = state ^. sOpenValves
+       let now = state ^. currentTime
+       succs <- agentSuccessor now opened now here
+       let succStates = Q.fromList succs
+       if isFF || (Q.null succStates)
+       then return $ Q.singleton (state & currentTime .~ timeLimit)
+       else return succStates
+
+  estimateBenefit here =  
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeOf here)
+       cf <- currentFlow here
+       -- let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
+       -- let sortedClosedValves = sortOn Down closedValves
+       sortedValves <- asks getSortedRooms
+       let opened = here ^. sOpenValves
+       let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+instance SearchState DoubleSearchState where
+  emptySearchState startID = DoubleSearchState 
+      { _personRoom = startID
+      , _personTime = 0
+      , _elephantRoom = startID
+      , _elephantTime = 0
+      , _dOpenValves = S.empty 
+      }
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = S.toList $ state ^. dOpenValves
+       return $ sum $ fmap (\v -> (cave ! v) ^. flowRate) valves
+       -- let presentRooms = cave `M.restrictKeys` valves
+       -- return $ sumOf (folded . flowRate) presentRooms
+
+  timeOf state = min (state ^. personTime) (state ^. elephantTime)
+
+  successors state = 
+    do isFF <- isFullFlow state
+       -- cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let opened = state ^. dOpenValves
+       let pNow = state ^. personTime
+       let eNow = state ^. elephantTime
+       let now = min pNow eNow
+       let pHere = state ^. personRoom
+       let eHere = state ^. elephantRoom
+       pNexts <- agentSuccessor now opened pNow pHere
+       eNexts <- agentSuccessor now opened eNow eHere
+       let nexts =  [ state & personRoom .~ (p ^. currentRoom)
+                            & personTime .~ (p ^. currentTime)
+                            & elephantRoom .~ (e ^. currentRoom)
+                            & elephantTime .~ (e ^. currentTime)
+                            & dOpenValves %~ (S.union (p ^. sOpenValves) . S.union (e ^. sOpenValves))
+                    | p <- pNexts
+                    , e <- eNexts
+                    ]
+       let dedups = if pNow == eNow && pHere == eHere
+                    then filter (\s -> (s ^. personRoom) < (s ^. elephantRoom)) nexts
+                    -- else nexts
+                    else filter (\s -> (s ^. personRoom) /= (s ^. elephantRoom)) nexts
+       -- let succStates = trace ("Succs: in " ++ (show state) ++ " out " ++ (show dedups)) (Q.fromList dedups)
+       let succStates = Q.fromList dedups
+       if isFF || (Q.null succStates)
+       then return $ Q.singleton (state & personTime .~ timeLimit & elephantTime .~ timeLimit)
+       else return succStates
+
+  estimateBenefit here =  
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeOf here)
+       cf <- currentFlow here
+       -- let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
+       -- let sortedClosedValves = fmap sum $ chunksOf 2 $  {-# SCC estSort #-} sortOn Down closedValves
+       -- let sortedClosedValves = fmap sum $ chunksOf 2 $ reverse $ sort closedValves -- no significant improvement
+       sortedValves <- asks getSortedRooms
+       let opened = here ^. dOpenValves
+       let sortedClosedValves = fmap sum $ chunksOf 2 $ [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       -- let otherValveFlows = timeRemaining * (sum closedValves) -- 8 minute runtime rather than 1:50
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+main :: IO ()
+main = 
+  do  dataFileName <- getDataFileName
+      text <- TIO.readFile dataFileName
+      let expandedCave = successfulParse text
+      -- print cave
+      -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
+      -- print $ compress cave
+      let cave = compress expandedCave
+      print $ part1 cave
+      print $ part2 cave
+
+-- part1 :: Cave -> Maybe (Agendum SingleSearchState)
+-- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
+
+-- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
+-- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
+
+part1, part2 :: Cave -> Int
+-- part1 :: Cave -> Int
+part1 cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave 30 sortedRooms) :: Maybe (Agendum SingleSearchState)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+part2 cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave 26 sortedRooms) :: Maybe (Agendum DoubleSearchState)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+          -- sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys cave
+
+searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
+searchCave startRoom = 
+    do agenda <- initAgenda startRoom
+       aStar agenda S.empty
+
+initAgenda :: SearchState s => String -> CaveContext (Agenda s)
+initAgenda startID = 
+    do let startState = emptySearchState startID
+       b <- estimateBenefit startState 
+       return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar ::  SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
+aStar agenda closed 
+    -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
+    | P.null agenda = return Nothing
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
+            -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+            let beamAgenda = P.fromDescList $ P.take 1000 newAgenda -- agenda beam width
+            reachedGoal <- isGoal currentAgendum
+            -- let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+            let cl = (reached, currentAgendum ^. trailBenefit)
+            if reachedGoal
+            then return (Just currentAgendum)
+            else if (cl `S.member` closed)
+                 then aStar (P.deleteMax agenda) closed
+                 -- else aStar newAgenda (S.insert cl closed)
+                 else aStar beamAgenda (S.insert cl closed)
+
+
+candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
+candidates agendum closed = 
+    do  let candidate = agendum ^. current
+        let previous = agendum ^. trail
+        let prevBenefit = agendum ^. trailBenefit
+        succs <- successors candidate
+        succAgs <- mapM (makeAgendum previous prevBenefit) succs
+        -- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+        let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
+        return nonloops
+
+
+agentSuccessor :: Int -> S.Set RoomID -> Int -> RoomID -> CaveContext [SingleSearchState]
+agentSuccessor now opened aTime here 
+  | aTime /= now = return [SingleSearchState { _currentRoom = here, _currentTime = aTime, _sOpenValves = opened }]
+  | otherwise = 
+      do  cave <- asks getCave 
+          timeLimit <- asks getTimeLimit
+          let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
+          let moves = [ SingleSearchState 
+                            { _currentRoom = (t ^. tunnelTo)
+                            , _currentTime = now + (t ^. tunnelLength)
+                            , _sOpenValves = opened
+                            }
+                      | t <- remaining
+                      , now + (t ^. tunnelLength) <= timeLimit
+                      ]
+          let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0
+                      then [SingleSearchState { _currentRoom = here, _currentTime = aTime + 1, _sOpenValves = S.insert here opened }]
+                      else []
+          -- let nexts = moves ++ opens
+          let nexts = if null opens then moves else opens
+          let nexts' =  if null nexts
+                        then [ SingleSearchState 
+                                { _currentRoom = here
+                                , _currentTime = timeLimit
+                                , _sOpenValves = opened
+                                } ]
+                        else nexts
+          return nexts'
+
+makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
+makeAgendum previous prevBenefit newState = 
+    do predicted <- estimateBenefit newState --  (Q.length previous)
+       -- cf <- currentFlow newState
+       oldFlow <- lastFlow previous (timeOf newState)
+       let newTrail = previous |> newState
+       let incurred = prevBenefit + oldFlow
+       return Agendum { _current = newState
+                      , _trail = newTrail
+                      , _trailBenefit = incurred
+                      , _benefit = incurred + predicted
+                      }
+
+lastFlow :: SearchState s => Q.Seq s -> Int -> CaveContext Int
+lastFlow Q.Empty _ = return 0
+lastFlow (_ :|> previous) newTime = 
+  do cf <- currentFlow previous
+     let dt = newTime - (timeOf previous)
+     return (cf * dt)
+
+isGoal :: SearchState s => Agendum s -> CaveContext Bool
+isGoal agendum = 
+  do timeLimit <- asks getTimeLimit
+     let s = agendum ^. current
+     return $ (timeOf s) == timeLimit
+
+isFullFlow :: SearchState s => s -> CaveContext Bool
+isFullFlow state = 
+  do cave <- asks getCave
+     cf <- currentFlow state
+     let ff = sumOf (folded . flowRate) cave
+     return (cf == ff)
+
+compress :: Cave -> Cave
+compress cave = M.mapWithKey (compressRoom cave) cave
+
+compressRoom :: Cave -> RoomID -> Room -> Room
+compressRoom cave here room = room & tunnels .~ t'
+  where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
+
+reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
+reachableFrom _ [] _ routes = routes
+reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
+  | here `S.member` found = reachableFrom cave boundary found routes
+  | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
+  where exits = (cave ! here) ^. tunnels
+        exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
+        legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
+        routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
+                  then routes
+                  else S.insert tunnel routes
+
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser (S.Set Tunnel)
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+    -- where roomify v ts = Room {flowRate = v, tunnels = ts }
+tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input = 
+  case parseOnly caveP input of
+    Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+    Right cave -> cave
\ No newline at end of file
diff --git a/advent16/MainCustomClosed.hs b/advent16/MainCustomClosed.hs
new file mode 100644 (file)
index 0000000..63fdf60
--- /dev/null
@@ -0,0 +1,401 @@
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+-- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>))) 
+import Data.Sequence ( (<|), (|>), Seq((:|>)) )
+import Data.List
+import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+
+
+type RoomID = String
+
+data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
+    deriving (Eq, Show, Ord)
+makeLenses ''Tunnel
+
+data Room = Room 
+    { _flowRate :: Int
+    , _tunnels :: S.Set Tunnel
+    } deriving (Eq, Show, Ord)
+makeLenses ''Room
+
+type Cave = M.Map RoomID Room
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
+
+type CaveContext = Reader TimedCave
+
+data SingleSearchState = SingleSearchState
+    { _currentRoom :: RoomID
+    , _currentTime :: Int
+    , _sOpenValves :: [RoomID]
+    } deriving (Eq, Show, Ord)
+makeLenses ''SingleSearchState
+
+data DoubleSearchState = DoubleSearchState
+    { _personRoom :: RoomID
+    , _personTime :: Int
+    , _elephantRoom :: RoomID
+    , _elephantTime :: Int
+    , _dOpenValves :: [RoomID]
+    } deriving (Eq, Show, Ord)
+makeLenses ''DoubleSearchState
+
+data Agendum s = 
+    Agendum { _current :: s
+            , _trail :: Q.Seq s
+            , _trailBenefit :: Int
+            , _benefit :: Int
+            } deriving (Show, Eq, Ord)
+makeLenses ''Agendum   
+
+type Agenda s = P.MaxPQueue Int (Agendum s)
+
+-- state, total flowed so far
+-- type ExploredStates s = S.Set (s, Int)
+
+
+class (Eq s, Ord s, Show s) => SearchState s where
+  emptySearchState :: RoomID -> s
+  currentFlow :: s -> CaveContext Int
+  timeOf :: s -> Int
+  successors :: s -> CaveContext (Q.Seq s)
+  -- estimateBenefit :: s -> Int -> CaveContext Int
+  estimateBenefit :: s -> CaveContext Int
+
+  data ExploredStateKey s
+  -- type ExploredStates s
+
+  mkExploredKey :: s -> (ExploredStateKey s)
+
+-- type ExploredStates s = M.Map (ExploredStateKey s) Int -- room/valves to time
+type ExploredStates s = S.Set ((ExploredStateKey s), Int) -- room & valves, trail benefit
+
+instance SearchState SingleSearchState where
+  emptySearchState startID = SingleSearchState 
+        { _currentRoom = startID
+        , _currentTime = 0
+        , _sOpenValves = []
+        }
+
+  data ExploredStateKey SingleSearchState = SingleExploredStateKey RoomID [RoomID] -- current room and open valves
+    deriving (Show, Eq, Ord)
+
+  mkExploredKey s = SingleExploredStateKey (s ^. currentRoom) (s ^. sOpenValves)
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. sOpenValves
+       let presentRooms = cave `M.restrictKeys` (S.fromList valves)
+       -- let presentRooms = M.filter (\v -> v `elem` valves) cave
+       return $ sumOf (folded . flowRate) presentRooms
+
+  timeOf state = state ^. currentTime
+
+  successors state = 
+    do isFF <- isFullFlow state
+       -- cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let here = state ^. currentRoom
+       let opened = state ^. sOpenValves
+       let now = state ^. currentTime
+       succs <- agentSuccessor now opened now here
+       let succStates = Q.fromList succs
+       if isFF || (Q.null succStates)
+       then return $ Q.singleton (state & currentTime .~ timeLimit)
+       else return succStates
+
+  estimateBenefit here =  
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeOf here)
+       cf <- currentFlow here
+       -- let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
+       -- let sortedClosedValves = sortOn Down closedValves
+       sortedValves <- asks getSortedRooms
+       let opened = here ^. sOpenValves
+       let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `notElem` opened]
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+instance SearchState DoubleSearchState where
+  emptySearchState startID = DoubleSearchState 
+      { _personRoom = startID
+      , _personTime = 0
+      , _elephantRoom = startID
+      , _elephantTime = 0
+      , _dOpenValves = []
+      }
+
+  data ExploredStateKey DoubleSearchState = DoubleExploredStateKey RoomID RoomID [RoomID] -- current room (person, elephant) and open valves
+    -- deriving (Show)
+    deriving (Show, Eq, Ord)
+  -- type ExploredStates DoubleSearchState = M.Map (DoubleExploredStateKey DoubleSearchState) Int -- room/valves to time
+
+  mkExploredKey s = DoubleExploredStateKey minRoom maxRoom (s ^. dOpenValves)
+    where minRoom = min (s ^. personRoom) (s ^. elephantRoom)
+          maxRoom = max (s ^. personRoom) (s ^. elephantRoom)
+
+  currentFlow state =
+    do cave <- asks getCave
+       -- let valves = S.toList $ state ^. dOpenValves
+       let valves = state ^. dOpenValves
+       return $ sum $ fmap (\v -> (cave ! v) ^. flowRate) valves
+       -- let presentRooms = cave `M.restrictKeys` valves
+       -- return $ sumOf (folded . flowRate) presentRooms
+
+  timeOf state = min (state ^. personTime) (state ^. elephantTime)
+
+  successors state = 
+    do isFF <- isFullFlow state
+       -- cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let opened = state ^. dOpenValves
+       let pNow = state ^. personTime
+       let eNow = state ^. elephantTime
+       let now = min pNow eNow
+       let pHere = state ^. personRoom
+       let eHere = state ^. elephantRoom
+       pNexts <- agentSuccessor now opened pNow pHere
+       eNexts <- agentSuccessor now opened eNow eHere
+       let nexts =  [ state & personRoom .~ (p ^. currentRoom)
+                            & personTime .~ (p ^. currentTime)
+                            & elephantRoom .~ (e ^. currentRoom)
+                            & elephantTime .~ (e ^. currentTime)
+                            -- & dOpenValves %~ (S.union (p ^. sOpenValves) . S.union (e ^. sOpenValves))
+                            & dOpenValves .~ (union (union opened (p ^. sOpenValves)) (e ^. sOpenValves))
+                    | p <- pNexts
+                    , e <- eNexts
+                    ]
+       let dedups = if pNow == eNow && pHere == eHere
+                    then filter (\s -> (s ^. personRoom) < (s ^. elephantRoom)) nexts
+                    else nexts
+       -- let succStates = trace ("Succs: in " ++ (show state) ++ " out " ++ (show dedups)) (Q.fromList dedups)
+       let succStates = Q.fromList dedups
+       if isFF || (Q.null succStates)
+       then return $ Q.singleton (state & personTime .~ timeLimit & elephantTime .~ timeLimit)
+       else return succStates
+
+  estimateBenefit here =  
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeOf here)
+       cf <- currentFlow here
+       -- let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
+       -- let sortedClosedValves = fmap sum $ chunksOf 2 $  {-# SCC estSort #-} sortOn Down closedValves
+       -- let sortedClosedValves = fmap sum $ chunksOf 2 $ reverse $ sort closedValves -- no significant improvement
+       sortedValves <- asks getSortedRooms
+       let opened = here ^. dOpenValves
+       let sortedClosedValves = fmap sum $ chunksOf 2 $ [(cave ! v) ^. flowRate | v <- sortedValves, v `notElem` opened]
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       -- let otherValveFlows = timeRemaining * (sum closedValves) -- 8 minute runtime rather than 1:50
+       return $ (cf * timeRemaining) + otherValveFlows
+
+-- instance Eq (ExploredStateKey DoubleSearchState) where 
+--     (DoubleExploredStateKey r1a r1b v1) == (DoubleExploredStateKey r2a r2b v2) = 
+--       -- ((r1a == r2a && r1b == r2b) || (r1a == r2b && r1b == r2a)) && v1 == v2
+--       ((min r1a r1b), (max r1a r1b), v1) == ((min r2a r2b), (max r2a r2b), v2)
+--   -- data instance Ord DoubleExploredStateKey where 
+-- instance Ord (ExploredStateKey DoubleSearchState) where 
+--     (DoubleExploredStateKey r1a r1b v1) `compare` (DoubleExploredStateKey r2a r2b v2) = 
+--       ((min r1a r1b), (max r1a r1b), v1) `compare` ((min r2a r2b), (max r2a r2b), v2)
+
+
+main :: IO ()
+main = 
+  do  dataFileName <- getDataFileName
+      text <- TIO.readFile dataFileName
+      let expandedCave = successfulParse text
+      -- print cave
+      -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
+      -- print $ compress cave
+      let cave = compress expandedCave
+      print $ part1 cave
+      print $ part2 cave
+
+-- part1 :: Cave -> Maybe (Agendum SingleSearchState)
+-- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
+
+-- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
+-- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
+
+part1, part2 :: Cave -> Int
+-- part1 :: Cave -> Int
+part1 cave = maybe 0 _benefit result
+-- part1 cave = result
+    where result = runReader (searchCave "AA") (TimedCave cave 30 sortedRooms) :: Maybe (Agendum SingleSearchState)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+part2 cave = maybe 0 _benefit result
+-- part2 cave = result
+    where result = runReader (searchCave "AA") (TimedCave cave 26 sortedRooms) :: Maybe (Agendum DoubleSearchState)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+          -- sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys cave
+
+searchCave :: ((Ord (ExploredStateKey s)), (Show (ExploredStateKey s)), SearchState s) => String -> CaveContext (Maybe (Agendum s))
+searchCave startRoom = 
+    do agenda <- initAgenda startRoom
+       aStar agenda S.empty
+
+initAgenda :: ((Ord (ExploredStateKey s)), (Show (ExploredStateKey s)), SearchState s) => String -> CaveContext (Agenda s)
+initAgenda startID = 
+    do let startState = emptySearchState startID
+       b <- estimateBenefit startState 
+       return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar :: ((Ord (ExploredStateKey s)), (Show (ExploredStateKey s)), SearchState s) => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
+aStar agenda closed 
+       -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ " : trail " ++ (show $ _trail $ snd $ P.findMax agenda) ++ " : closed " ++ (show closed)) False = undefined
+       -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
+    | P.null agenda = return Nothing
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
+            -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+            -- let beamAgenda = P.fromDescList $ P.take 1000 newAgenda -- agenda beam width
+            reachedGoal <- isGoal currentAgendum
+            -- let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+            -- let cl = (reached, currentAgendum ^. trailBenefit)
+            let cl = (mkExploredKey reached, currentAgendum ^. trailBenefit)
+            if reachedGoal
+            then return (Just currentAgendum)
+            else if (cl `elem` closed)
+                 then aStar (P.deleteMax agenda) closed
+                 else aStar newAgenda (S.insert cl closed)
+
+
+candidates :: ((Ord (ExploredStateKey s)), SearchState s) => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
+candidates agendum closed = 
+    do  let candidate = agendum ^. current
+        let previous = agendum ^. trail
+        let prevBenefit = agendum ^. trailBenefit
+        succs <- successors candidate
+        succAgs <- mapM (makeAgendum previous prevBenefit) succs
+        -- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+        -- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
+        let nonloops = Q.filter (\l -> ((mkExploredKey (l ^. current)), l ^. trailBenefit) `notElem` closed) succAgs
+        return nonloops
+
+
+agentSuccessor :: Int -> [RoomID] -> Int -> RoomID -> CaveContext [SingleSearchState]
+agentSuccessor now opened aTime here 
+  | aTime /= now = return [SingleSearchState { _currentRoom = here, _currentTime = aTime, _sOpenValves = opened }]
+  | otherwise = 
+      do  cave <- asks getCave 
+          timeLimit <- asks getTimeLimit
+          -- let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
+          let remaining = [ t 
+                          | t <- (S.toList ((cave ! here) ^. tunnels))
+                          , (t ^. tunnelTo) `notElem` opened
+                          ]
+          let moves = [ SingleSearchState 
+                            { _currentRoom = (t ^. tunnelTo)
+                            , _currentTime = now + (t ^. tunnelLength)
+                            , _sOpenValves = opened
+                            }
+                      | t <- remaining
+                      , now + (t ^. tunnelLength) <= timeLimit
+                      ]
+          let moves' = ( SingleSearchState 
+                                { _currentRoom = here
+                                , _currentTime = timeLimit
+                                , _sOpenValves = opened
+                                } 
+                       : moves)   
+          let opens = if here `notElem` opened && (cave ! here) ^. flowRate > 0
+                      then [SingleSearchState { _currentRoom = here, _currentTime = aTime + 1, _sOpenValves = opened ++ [here] }]
+                      else []
+          -- let nexts = moves ++ opens
+          let nexts = if null opens then moves' else opens
+          return nexts
+
+makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
+makeAgendum previous prevBenefit newState = 
+    do predicted <- estimateBenefit newState --  (Q.length previous)
+       -- cf <- currentFlow newState
+       oldFlow <- lastFlow previous (timeOf newState)
+       let newTrail = previous |> newState
+       let incurred = prevBenefit + oldFlow
+       return Agendum { _current = newState
+                      , _trail = newTrail
+                      , _trailBenefit = incurred
+                      , _benefit = incurred + predicted
+                      }
+
+lastFlow :: SearchState s => Q.Seq s -> Int -> CaveContext Int
+lastFlow Q.Empty _ = return 0
+lastFlow (_ :|> previous) newTime = 
+  do cf <- currentFlow previous
+     let dt = newTime - (timeOf previous)
+     return (cf * dt)
+
+isGoal :: SearchState s => Agendum s -> CaveContext Bool
+isGoal agendum = 
+  do timeLimit <- asks getTimeLimit
+     let s = agendum ^. current
+     return $ (timeOf s) == timeLimit
+
+isFullFlow :: SearchState s => s -> CaveContext Bool
+isFullFlow state = 
+  do cave <- asks getCave
+     cf <- currentFlow state
+     let ff = sumOf (folded . flowRate) cave
+     return (cf == ff)
+
+compress :: Cave -> Cave
+compress cave = M.mapWithKey (compressRoom cave) cave
+
+compressRoom :: Cave -> RoomID -> Room -> Room
+compressRoom cave here room = room & tunnels .~ t'
+  where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
+
+reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
+reachableFrom _ [] _ routes = routes
+reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
+  | here `S.member` found = reachableFrom cave boundary found routes
+  | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
+  where exits = (cave ! here) ^. tunnels
+        exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
+        legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
+        routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
+                  then routes
+                  else S.insert tunnel routes
+
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser (S.Set Tunnel)
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+    -- where roomify v ts = Room {flowRate = v, tunnels = ts }
+tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input = 
+  case parseOnly caveP input of
+    Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+    Right cave -> cave
\ No newline at end of file
diff --git a/advent16/MainEstSort.hs b/advent16/MainEstSort.hs
new file mode 100644 (file)
index 0000000..459a10b
--- /dev/null
@@ -0,0 +1,352 @@
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+-- import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+-- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>))) 
+import Data.Sequence ( (|>), Seq((:|>)) )
+import Data.List
+import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+
+
+type RoomID = String
+
+data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
+    deriving (Eq, Show, Ord)
+makeLenses ''Tunnel
+
+data Room = Room 
+    { _flowRate :: Int
+    , _tunnels :: S.Set Tunnel
+    } deriving (Eq, Show, Ord)
+makeLenses ''Room
+
+type Cave = M.Map RoomID Room
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int }
+
+type CaveContext = Reader TimedCave
+
+data SingleSearchState = SingleSearchState
+    { _currentRoom :: RoomID
+    , _currentTime :: Int
+    , _sOpenValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''SingleSearchState
+
+data DoubleSearchState = DoubleSearchState
+    { _personRoom :: RoomID
+    , _personTime :: Int
+    , _elephantRoom :: RoomID
+    , _elephantTime :: Int
+    , _dOpenValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''DoubleSearchState
+
+data Agendum s = 
+    Agendum { _current :: s
+            , _trail :: Q.Seq s
+            , _trailBenefit :: Int
+            , _benefit :: Int
+            } deriving (Show, Eq, Ord)
+makeLenses ''Agendum   
+
+type Agenda s = P.MaxPQueue Int (Agendum s)
+
+-- state, total flowed so far
+type ExploredStates s = S.Set (s, Int)
+
+
+class (Eq s, Ord s, Show s) => SearchState s where
+  emptySearchState :: RoomID -> s
+  currentFlow :: s -> CaveContext Int
+  timeOf :: s -> Int
+  successors :: s -> CaveContext (Q.Seq s)
+  -- estimateBenefit :: s -> Int -> CaveContext Int
+  estimateBenefit :: s -> CaveContext Int
+
+instance SearchState SingleSearchState where
+  emptySearchState startID = SingleSearchState 
+        { _currentRoom = startID
+        , _currentTime = 0
+        , _sOpenValves = S.empty 
+        }
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. sOpenValves
+       let presentRooms = cave `M.restrictKeys` valves
+       return $ sumOf (folded . flowRate) presentRooms
+
+  timeOf state = state ^. currentTime
+
+  successors state = 
+    do isFF <- isFullFlow state
+       -- cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let here = state ^. currentRoom
+       let opened = state ^. sOpenValves
+       let now = state ^. currentTime
+       succs <- agentSuccessor now opened now here
+       let succStates = Q.fromList succs
+       if isFF || (Q.null succStates)
+       then return $ Q.singleton (state & currentTime .~ timeLimit)
+       else return succStates
+
+  estimateBenefit here =  
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeOf here)
+       cf <- currentFlow here
+       let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
+       let sortedClosedValves = sortOn Down closedValves
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+instance SearchState DoubleSearchState where
+  emptySearchState startID = DoubleSearchState 
+      { _personRoom = startID
+      , _personTime = 0
+      , _elephantRoom = startID
+      , _elephantTime = 0
+      , _dOpenValves = S.empty 
+      }
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = S.toList $ state ^. dOpenValves
+       return $ sum $ fmap (\v -> (cave ! v) ^. flowRate) valves
+       -- let presentRooms = cave `M.restrictKeys` valves
+       -- return $ sumOf (folded . flowRate) presentRooms
+
+  timeOf state = min (state ^. personTime) (state ^. elephantTime)
+
+  successors state = 
+    do isFF <- isFullFlow state
+       -- cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let opened = state ^. dOpenValves
+       let pNow = state ^. personTime
+       let eNow = state ^. elephantTime
+       let now = min pNow eNow
+       let pHere = state ^. personRoom
+       let eHere = state ^. elephantRoom
+       pNexts <- agentSuccessor now opened pNow pHere
+       eNexts <- agentSuccessor now opened eNow eHere
+       let nexts =  [ state & personRoom .~ (p ^. currentRoom)
+                            & personTime .~ (p ^. currentTime)
+                            & elephantRoom .~ (e ^. currentRoom)
+                            & elephantTime .~ (e ^. currentTime)
+                            & dOpenValves %~ (S.union (p ^. sOpenValves) . S.union (e ^. sOpenValves))
+                    | p <- pNexts
+                    , e <- eNexts
+                    ]
+       let dedups = if pNow == eNow && pHere == eHere
+                    then filter (\s -> (s ^. personRoom) < (s ^. elephantRoom)) nexts
+                    else nexts
+       -- let succStates = trace ("Succs: in " ++ (show state) ++ " out " ++ (show dedups)) (Q.fromList dedups)
+       let succStates = Q.fromList dedups
+       if isFF || (Q.null succStates)
+       then return $ Q.singleton (state & personTime .~ timeLimit & elephantTime .~ timeLimit)
+       else return succStates
+
+  estimateBenefit here =  
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeOf here)
+       cf <- currentFlow here
+       let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
+       let sortedClosedValves = fmap sum $ chunksOf 2 $  {-# SCC estSort #-} sortOn Down closedValves
+       -- let sortedClosedValves = fmap sum $ chunksOf 2 $ reverse $ sort closedValves -- no significant improvement
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       -- let otherValveFlows = timeRemaining * (sum closedValves) -- 8 minute runtime rather than 1:50
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+main :: IO ()
+main = 
+  do  dataFileName <- getDataFileName
+      text <- TIO.readFile dataFileName
+      let expandedCave = successfulParse text
+      -- print cave
+      -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
+      -- print $ compress cave
+      let cave = compress expandedCave
+      print $ part1 cave
+      print $ part2 cave
+
+-- part1 :: Cave -> Maybe (Agendum SingleSearchState)
+-- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
+
+-- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
+-- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
+
+part1, part2 :: Cave -> Int
+-- part1 :: Cave -> Int
+part1 cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave 30) :: Maybe (Agendum SingleSearchState)
+part2 cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave 26) :: Maybe (Agendum DoubleSearchState)
+
+searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
+searchCave startRoom = 
+    do agenda <- initAgenda startRoom
+       aStar agenda S.empty
+
+initAgenda :: SearchState s => String -> CaveContext (Agenda s)
+initAgenda startID = 
+    do let startState = emptySearchState startID
+       b <- estimateBenefit startState 
+       return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar ::  SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
+aStar agenda closed 
+    -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
+    | P.null agenda = return Nothing
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
+            -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+            -- let beamAgenda = P.fromDescList $ P.take 1000 newAgenda -- agenda beam width
+            reachedGoal <- isGoal currentAgendum
+            -- let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+            let cl = (reached, currentAgendum ^. trailBenefit)
+            if reachedGoal
+            then return (Just currentAgendum)
+            else if (cl `S.member` closed)
+                 then aStar (P.deleteMax agenda) closed
+                 else aStar newAgenda (S.insert cl closed)
+                 -- else aStar beamAgenda (S.insert cl closed)
+
+
+candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
+candidates agendum closed = 
+    do  let candidate = agendum ^. current
+        let previous = agendum ^. trail
+        let prevBenefit = agendum ^. trailBenefit
+        succs <- successors candidate
+        succAgs <- mapM (makeAgendum previous prevBenefit) succs
+        -- let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+        let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
+        return nonloops
+
+
+agentSuccessor :: Int -> S.Set RoomID -> Int -> RoomID -> CaveContext [SingleSearchState]
+agentSuccessor now opened aTime here 
+  | aTime /= now = return [SingleSearchState { _currentRoom = here, _currentTime = aTime, _sOpenValves = opened }]
+  | otherwise = 
+      do  cave <- asks getCave 
+          timeLimit <- asks getTimeLimit
+          let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
+          let moves = [ SingleSearchState 
+                            { _currentRoom = (t ^. tunnelTo)
+                            , _currentTime = now + (t ^. tunnelLength)
+                            , _sOpenValves = opened
+                            }
+                      | t <- remaining
+                      , now + (t ^. tunnelLength) <= timeLimit
+                      ]
+          let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0
+                      then [SingleSearchState { _currentRoom = here, _currentTime = aTime + 1, _sOpenValves = S.insert here opened }]
+                      else []
+          -- let nexts = moves ++ opens
+          let nexts = if null opens then moves else opens
+          let nexts' =  if null nexts
+                        then [ SingleSearchState 
+                                { _currentRoom = here
+                                , _currentTime = timeLimit
+                                , _sOpenValves = opened
+                                } ]
+                        else nexts
+          return nexts'
+
+makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
+makeAgendum previous prevBenefit newState = 
+    do predicted <- estimateBenefit newState --  (Q.length previous)
+       -- cf <- currentFlow newState
+       oldFlow <- lastFlow previous (timeOf newState)
+       let newTrail = previous |> newState
+       let incurred = prevBenefit + oldFlow
+       return Agendum { _current = newState
+                      , _trail = newTrail
+                      , _trailBenefit = incurred
+                      , _benefit = incurred + predicted
+                      }
+
+lastFlow :: SearchState s => Q.Seq s -> Int -> CaveContext Int
+lastFlow Q.Empty _ = return 0
+lastFlow (_ :|> previous) newTime = 
+  do cf <- currentFlow previous
+     let dt = newTime - (timeOf previous)
+     return (cf * dt)
+
+isGoal :: SearchState s => Agendum s -> CaveContext Bool
+isGoal agendum = 
+  do timeLimit <- asks getTimeLimit
+     let s = agendum ^. current
+     return $ (timeOf s) == timeLimit
+
+isFullFlow :: SearchState s => s -> CaveContext Bool
+isFullFlow state = 
+  do cave <- asks getCave
+     cf <- currentFlow state
+     let ff = sumOf (folded . flowRate) cave
+     return (cf == ff)
+
+compress :: Cave -> Cave
+compress cave = M.mapWithKey (compressRoom cave) cave
+
+compressRoom :: Cave -> RoomID -> Room -> Room
+compressRoom cave here room = room & tunnels .~ t'
+  where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
+
+reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
+reachableFrom _ [] _ routes = routes
+reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
+  | here `S.member` found = reachableFrom cave boundary found routes
+  | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
+  where exits = (cave ! here) ^. tunnels
+        exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
+        legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
+        routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
+                  then routes
+                  else S.insert tunnel routes
+
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser (S.Set Tunnel)
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+    -- where roomify v ts = Room {flowRate = v, tunnels = ts }
+tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input = 
+  case parseOnly caveP input of
+    Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+    Right cave -> cave
\ No newline at end of file
diff --git a/advent16/MainOriginal.hs b/advent16/MainOriginal.hs
new file mode 100644 (file)
index 0000000..3000671
--- /dev/null
@@ -0,0 +1,275 @@
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+-- import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+import Data.Sequence ((|>)) 
+import Data.List
+import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+
+-- pattern Empty   <- (Q.viewl -> Q.EmptyL)  where Empty = Q.empty
+-- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<)  = (Q.<|) 
+-- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>)  = (Q.|>) 
+
+type RoomID = String
+
+data Room = Room 
+    { _flowRate :: Int
+    , _tunnels :: [RoomID]
+    } deriving (Eq, Show, Ord)
+makeLenses ''Room
+
+type Cave = M.Map RoomID Room
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int}
+
+type CaveContext = Reader TimedCave
+
+data SingleSearchState = SingleSearchState
+    { _currentRoom :: RoomID
+    , _sOpenValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''SingleSearchState
+
+data DoubleSearchState = DoubleSearchState
+    { _personRoom :: RoomID
+    , _elephantRoom :: RoomID
+    , _dOpenValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''DoubleSearchState
+
+data Agendum s = 
+    Agendum { _current :: s
+            , _trail :: Q.Seq s
+            , _trailBenefit :: Int
+            , _benefit :: Int
+            } deriving (Show, Eq, Ord)
+makeLenses ''Agendum   
+
+type Agenda s = P.MaxPQueue Int (Agendum s)
+
+type ExploredStates s = S.Set (s, Int, Int)
+
+
+class (Eq s, Ord s, Show s) => SearchState s where
+  emptySearchState :: RoomID -> s
+  currentFlow :: s -> CaveContext Int
+  successors :: s -> CaveContext (Q.Seq s)
+  estimateBenefit :: s -> Int -> CaveContext Int
+
+instance SearchState SingleSearchState where
+  emptySearchState startID = SingleSearchState { _currentRoom = startID, _sOpenValves = S.empty }
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. sOpenValves
+       let presentRooms = cave `M.restrictKeys` valves
+       return $ sumOf (folded . flowRate) presentRooms
+
+  successors state = 
+    do isFF <- isFullFlow state
+       let here = state ^. currentRoom
+       let opened = state ^. sOpenValves
+       succPairs <- personSuccessor here opened
+       let succStates = 
+              [ SingleSearchState 
+                  { _currentRoom = r
+                  , _sOpenValves = o
+                  }
+              | (r, o) <- succPairs
+              ]
+       if isFF
+       then return $ Q.singleton state
+       else return $ Q.fromList succStates
+
+  estimateBenefit here timeElapsed = 
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeElapsed + 2)
+       cf <- currentFlow here
+       let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
+       let sortedClosedValves = sortOn Down closedValves
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+instance SearchState DoubleSearchState where
+  emptySearchState startID = DoubleSearchState 
+      { _personRoom = startID
+      , _elephantRoom = startID
+      , _dOpenValves = S.empty 
+      }
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. dOpenValves
+       let presentRooms = cave `M.restrictKeys` valves
+       return $ sumOf (folded . flowRate) presentRooms
+
+  successors state = 
+    do isFF <- isFullFlow state
+       let pHere = state ^. personRoom
+       let eHere = state ^. elephantRoom
+       let opened = state ^. dOpenValves
+       pSuccPairs <- personSuccessor pHere opened
+       eSuccPairs <- personSuccessor eHere opened
+       let succStates = 
+              [ DoubleSearchState 
+                  { _personRoom = p
+                  , _elephantRoom = e
+                  , _dOpenValves = S.union po eo
+                  }
+              | (p, po) <- pSuccPairs
+              , (e, eo) <- eSuccPairs
+              ]
+       if isFF
+       then return $ Q.singleton state
+       else return $ Q.fromList succStates
+
+  estimateBenefit here timeElapsed = 
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeElapsed + 2)
+       cf <- currentFlow here
+       let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
+       let sortedClosedValves = fmap sum $ chunksOf 2 $ sortOn Down closedValves
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+main :: IO ()
+main = 
+  do  dataFileName <- getDataFileName
+      text <- TIO.readFile dataFileName
+      let cave = successfulParse text
+      -- print cave
+      print $ part1 cave
+      print $ part2 cave
+
+-- part1 :: Cave -> Maybe (Agendum SingleSearchState)
+-- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
+
+-- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
+-- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
+
+part1, part2 :: Cave -> Int
+part1 cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave 30) :: Maybe (Agendum SingleSearchState)
+part2 cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave 26) :: Maybe (Agendum DoubleSearchState)
+
+searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
+searchCave startRoom = 
+    do agenda <- initAgenda startRoom
+       aStar agenda S.empty
+
+initAgenda :: SearchState s => String -> CaveContext (Agenda s)
+initAgenda startID = 
+    do let startState = emptySearchState startID
+       b <- estimateBenefit startState 0
+       return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar ::  SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
+aStar agenda closed 
+    -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : len " ++ (show $ Q.length $ _trail $ snd $ P.findMax agenda)) False = undefined
+    | P.null agenda = return Nothing
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
+            let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+            reachedGoal <- isGoal currentAgendum
+            let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+            if reachedGoal
+            then return (Just currentAgendum)
+            else if (cl `S.member` closed)
+                 then aStar (P.deleteMax agenda) closed
+                 -- else aStar newAgenda (S.insert cl closed)
+                 else aStar beamAgenda (S.insert cl closed)
+
+
+candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
+candidates agendum closed = 
+    do  let candidate = agendum ^. current
+        let previous = agendum ^. trail
+        let prevBenefit = agendum ^. trailBenefit
+        succs <- successors candidate
+        succAgs <- mapM (makeAgendum previous prevBenefit) succs
+        let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+        return nonloops
+
+personSuccessor, openValveSuccessor, walkSuccessor ::  RoomID -> S.Set RoomID -> CaveContext [(RoomID, S.Set RoomID)]
+personSuccessor here opened =
+  do ovs <- openValveSuccessor here opened
+     ws <- walkSuccessor here opened
+     return (ovs ++ ws)
+
+openValveSuccessor here opened
+  | here `S.member` opened = return []
+  | otherwise = return [(here, S.insert here opened)]
+
+walkSuccessor here opened = 
+  do cave <- asks getCave
+     let neighbours = (cave ! here) ^. tunnels
+     return [(n, opened) | n <- neighbours]
+
+makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
+makeAgendum previous prevBenefit newState = 
+    do predicted <- estimateBenefit newState (Q.length previous)
+       cf <- currentFlow newState
+       let newTrail = previous |> newState
+       let incurred = prevBenefit + cf
+       return Agendum { _current = newState
+                      , _trail = newTrail
+                      , _trailBenefit = incurred
+                      , _benefit = incurred + predicted
+                      }
+
+
+isGoal :: SearchState s => Agendum s -> CaveContext Bool
+isGoal agendum = 
+  do timeLimit <- asks getTimeLimit
+     return $ Q.length (agendum ^. trail) == (timeLimit - 1)
+
+isFullFlow :: SearchState s => s -> CaveContext Bool
+isFullFlow state = 
+  do cave <- asks getCave
+     cf <- currentFlow state
+     let ff = sumOf (folded . flowRate) cave
+     return (cf == ff)
+
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser [RoomID]
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = roomify <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+    where roomify v ts = Room {_flowRate = v, _tunnels = ts }
+tunnelsP = (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input = 
+  case parseOnly caveP input of
+    Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+    Right cave -> cave
\ No newline at end of file
diff --git a/advent16/MainOriginalNoBeam.hs b/advent16/MainOriginalNoBeam.hs
new file mode 100644 (file)
index 0000000..b074d93
--- /dev/null
@@ -0,0 +1,275 @@
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+-- import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+import Data.Sequence ((|>)) 
+import Data.List
+import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+
+-- pattern Empty   <- (Q.viewl -> Q.EmptyL)  where Empty = Q.empty
+-- pattern x :< xs <- (Q.viewl -> x Q.:< xs) where (:<)  = (Q.<|) 
+-- pattern xs :> x <- (Q.viewr -> xs Q.:> x) where (:>)  = (Q.|>) 
+
+type RoomID = String
+
+data Room = Room 
+    { _flowRate :: Int
+    , _tunnels :: [RoomID]
+    } deriving (Eq, Show, Ord)
+makeLenses ''Room
+
+type Cave = M.Map RoomID Room
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int}
+
+type CaveContext = Reader TimedCave
+
+data SingleSearchState = SingleSearchState
+    { _currentRoom :: RoomID
+    , _sOpenValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''SingleSearchState
+
+data DoubleSearchState = DoubleSearchState
+    { _personRoom :: RoomID
+    , _elephantRoom :: RoomID
+    , _dOpenValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''DoubleSearchState
+
+data Agendum s = 
+    Agendum { _current :: s
+            , _trail :: Q.Seq s
+            , _trailBenefit :: Int
+            , _benefit :: Int
+            } deriving (Show, Eq, Ord)
+makeLenses ''Agendum   
+
+type Agenda s = P.MaxPQueue Int (Agendum s)
+
+type ExploredStates s = S.Set (s, Int, Int)
+
+
+class (Eq s, Ord s, Show s) => SearchState s where
+  emptySearchState :: RoomID -> s
+  currentFlow :: s -> CaveContext Int
+  successors :: s -> CaveContext (Q.Seq s)
+  estimateBenefit :: s -> Int -> CaveContext Int
+
+instance SearchState SingleSearchState where
+  emptySearchState startID = SingleSearchState { _currentRoom = startID, _sOpenValves = S.empty }
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. sOpenValves
+       let presentRooms = cave `M.restrictKeys` valves
+       return $ sumOf (folded . flowRate) presentRooms
+
+  successors state = 
+    do isFF <- isFullFlow state
+       let here = state ^. currentRoom
+       let opened = state ^. sOpenValves
+       succPairs <- personSuccessor here opened
+       let succStates = 
+              [ SingleSearchState 
+                  { _currentRoom = r
+                  , _sOpenValves = o
+                  }
+              | (r, o) <- succPairs
+              ]
+       if isFF
+       then return $ Q.singleton state
+       else return $ Q.fromList succStates
+
+  estimateBenefit here timeElapsed = 
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeElapsed + 2)
+       cf <- currentFlow here
+       let closedValves = (cave `M.withoutKeys` (here ^. sOpenValves)) ^.. folded . flowRate
+       let sortedClosedValves = sortOn Down closedValves
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+instance SearchState DoubleSearchState where
+  emptySearchState startID = DoubleSearchState 
+      { _personRoom = startID
+      , _elephantRoom = startID
+      , _dOpenValves = S.empty 
+      }
+
+  currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. dOpenValves
+       let presentRooms = cave `M.restrictKeys` valves
+       return $ sumOf (folded . flowRate) presentRooms
+
+  successors state = 
+    do isFF <- isFullFlow state
+       let pHere = state ^. personRoom
+       let eHere = state ^. elephantRoom
+       let opened = state ^. dOpenValves
+       pSuccPairs <- personSuccessor pHere opened
+       eSuccPairs <- personSuccessor eHere opened
+       let succStates = 
+              [ DoubleSearchState 
+                  { _personRoom = p
+                  , _elephantRoom = e
+                  , _dOpenValves = S.union po eo
+                  }
+              | (p, po) <- pSuccPairs
+              , (e, eo) <- eSuccPairs
+              ]
+       if isFF
+       then return $ Q.singleton state
+       else return $ Q.fromList succStates
+
+  estimateBenefit here timeElapsed = 
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeElapsed + 2)
+       cf <- currentFlow here
+       let closedValves = (cave `M.withoutKeys` (here ^. dOpenValves)) ^.. folded . flowRate
+       let sortedClosedValves = fmap sum $ chunksOf 2 $ sortOn Down closedValves
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+
+main :: IO ()
+main = 
+  do  dataFileName <- getDataFileName
+      text <- TIO.readFile dataFileName
+      let cave = successfulParse text
+      -- print cave
+      print $ part1 cave
+      print $ part2 cave
+
+-- part1 :: Cave -> Maybe (Agendum SingleSearchState)
+-- part1 cave = runReader (searchCave "AA") (TimedCave cave 30)
+
+-- part2 :: Cave -> Maybe (Agendum DoubleSearchState)
+-- part2 cave = runReader (searchCave "AA") (TimedCave cave 26)
+
+part1, part2 :: Cave -> Int
+part1 cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave 30) :: Maybe (Agendum SingleSearchState)
+part2 cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave 26) :: Maybe (Agendum DoubleSearchState)
+
+searchCave :: SearchState s => String -> CaveContext (Maybe (Agendum s))
+searchCave startRoom = 
+    do agenda <- initAgenda startRoom
+       aStar agenda S.empty
+
+initAgenda :: SearchState s => String -> CaveContext (Agenda s)
+initAgenda startID = 
+    do let startState = emptySearchState startID
+       b <- estimateBenefit startState 0
+       return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar ::  SearchState s => Agenda s -> ExploredStates s -> CaveContext (Maybe (Agendum s))
+aStar agenda closed 
+    -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : len " ++ (show $ Q.length $ _trail $ snd $ P.findMax agenda)) False = undefined
+    | P.null agenda = return Nothing
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            -- let beamAgenda = P.fromDescList $ P.take 10000 newAgenda -- agenda beam width
+            -- let beamAgenda = P.fromDescList $ P.take 5000 newAgenda -- agenda beam width
+            reachedGoal <- isGoal currentAgendum
+            let cl = (reached, currentAgendum ^. trailBenefit, Q.length $ currentAgendum ^. trail)
+            if reachedGoal
+            then return (Just currentAgendum)
+            else if (cl `S.member` closed)
+                 then aStar (P.deleteMax agenda) closed
+                 else aStar newAgenda (S.insert cl closed)
+                 -- else aStar beamAgenda (S.insert cl closed)
+
+
+candidates :: SearchState s => Agendum s -> ExploredStates s -> CaveContext (Q.Seq (Agendum s))
+candidates agendum closed = 
+    do  let candidate = agendum ^. current
+        let previous = agendum ^. trail
+        let prevBenefit = agendum ^. trailBenefit
+        succs <- successors candidate
+        succAgs <- mapM (makeAgendum previous prevBenefit) succs
+        let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit, Q.length $ s ^. trail) `S.notMember` closed) succAgs
+        return nonloops
+
+personSuccessor, openValveSuccessor, walkSuccessor ::  RoomID -> S.Set RoomID -> CaveContext [(RoomID, S.Set RoomID)]
+personSuccessor here opened =
+  do ovs <- openValveSuccessor here opened
+     ws <- walkSuccessor here opened
+     return (ovs ++ ws)
+
+openValveSuccessor here opened
+  | here `S.member` opened = return []
+  | otherwise = return [(here, S.insert here opened)]
+
+walkSuccessor here opened = 
+  do cave <- asks getCave
+     let neighbours = (cave ! here) ^. tunnels
+     return [(n, opened) | n <- neighbours]
+
+makeAgendum :: SearchState s => Q.Seq s -> Int -> s -> CaveContext (Agendum s)
+makeAgendum previous prevBenefit newState = 
+    do predicted <- estimateBenefit newState (Q.length previous)
+       cf <- currentFlow newState
+       let newTrail = previous |> newState
+       let incurred = prevBenefit + cf
+       return Agendum { _current = newState
+                      , _trail = newTrail
+                      , _trailBenefit = incurred
+                      , _benefit = incurred + predicted
+                      }
+
+
+isGoal :: SearchState s => Agendum s -> CaveContext Bool
+isGoal agendum = 
+  do timeLimit <- asks getTimeLimit
+     return $ Q.length (agendum ^. trail) == (timeLimit - 1)
+
+isFullFlow :: SearchState s => s -> CaveContext Bool
+isFullFlow state = 
+  do cave <- asks getCave
+     cf <- currentFlow state
+     let ff = sumOf (folded . flowRate) cave
+     return (cf == ff)
+
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser [RoomID]
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = roomify <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+    where roomify v ts = Room {_flowRate = v, _tunnels = ts }
+tunnelsP = (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input = 
+  case parseOnly caveP input of
+    Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+    Right cave -> cave
\ No newline at end of file
diff --git a/advent16/MainSPar.hs b/advent16/MainSPar.hs
new file mode 100644 (file)
index 0000000..d3be5ec
--- /dev/null
@@ -0,0 +1,323 @@
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+-- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>))) 
+import Data.Sequence ( (|>), Seq((:|>)) )
+import Data.List
+import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+import Control.Parallel.Strategies
+
+
+type RoomID = String
+
+data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
+    deriving (Eq, Show, Ord)
+makeLenses ''Tunnel
+
+data Room = Room 
+    { _flowRate :: Int
+    , _tunnels :: S.Set Tunnel
+    } deriving (Eq, Show, Ord)
+makeLenses ''Room
+
+type Cave = M.Map RoomID Room
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
+
+type CaveContext = Reader TimedCave
+
+data SearchState = SearchState
+    { _currentRoom :: RoomID
+    , _currentTime :: Int
+    , _openValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''SearchState
+
+data Agendum = 
+    Agendum { _current :: SearchState
+            , _trail :: Q.Seq SearchState
+            , _trailBenefit :: Int
+            , _benefit :: Int
+            } deriving (Show, Eq, Ord)
+makeLenses ''Agendum   
+
+type Agenda = P.MaxPQueue Int Agendum
+
+-- state, total flowed so far
+type ExploredStates = S.Set (SearchState, Int)
+
+type PartSolutions = M.Map (S.Set RoomID) Int
+
+
+main :: IO ()
+main = 
+  do  dataFileName <- getDataFileName
+      text <- TIO.readFile dataFileName
+      let expandedCave = successfulParse text
+      -- print cave
+      -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
+      -- print $ compress cave
+      -- putStrLn $ dotify expandedCave
+      let cave = compress expandedCave
+      print $ part1 cave
+      print $ part2 cave
+
+-- dotify cave = "graph G {\n" ++ (unlines $ concat $ M.elems $ M.mapWithKey showCRoom cave) ++ "\n}\n"
+--   where showCRoom roomID room = filter (not . null) ((showCRoomShape roomID room) : (showCRoomLinks roomID room))
+
+-- showCRoomShape roomID room 
+--   | room ^. flowRate > 0 = roomID ++ " [fillcolor=grey label=\"" ++ roomID ++ ": " ++ (show $ room ^. flowRate) ++ "\"];"
+--   | otherwise = ""
+
+-- showCRoomLinks roomID room = [roomID ++ " -- " ++ (t ^. tunnelTo) ++ ";" | t <- S.toList $ room ^. tunnels, (t ^. tunnelTo) > roomID ]
+
+part1, part2 :: Cave -> Int
+-- part1 :: Cave -> Int
+part1 cave = runSearch 30 cave
+part2 cave = maximum (fmap maximum chunkSolns `using` parList rdeepseq)
+  where rawSolutions = runSearchAll 26 cave
+        solutionList = M.toList rawSolutions
+        combinations = [ fp + fe 
+                       | (p, fp) <- solutionList
+                       , (e, fe) <- solutionList
+                       , p < e
+                       , S.disjoint p e
+                       ]
+        chunkSolns = chunksOf 10000 combinations
+
+includeAgendum :: PartSolutions -> Agendum -> CaveContext PartSolutions
+includeAgendum results agendum =
+  do cf <- currentFlow (agendum ^. current)
+     timeLimit <- asks getTimeLimit
+     let timeLeft = timeLimit - timeOf (agendum ^. current)
+     let remainingFlow = cf * timeLeft
+     let totalFlow = remainingFlow + agendum ^. trailBenefit
+     let visitedSet = agendum ^. current . openValves
+     let currentBest = M.findWithDefault 0 visitedSet results
+     if totalFlow > currentBest
+     then return (M.insert visitedSet totalFlow results)
+     else return results
+
+runSearch :: Int -> Cave -> Int
+runSearch timeLimit cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave timeLimit sortedRooms)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+
+runSearchAll :: Int -> Cave -> PartSolutions
+runSearchAll timeLimit cave = result
+    where result = runReader (searchCaveAll "AA") (TimedCave cave timeLimit sortedRooms)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+
+
+searchCave :: String -> CaveContext (Maybe Agendum)
+searchCave startRoom = 
+    do agenda <- initAgenda startRoom
+       aStar agenda S.empty
+
+searchCaveAll :: String -> CaveContext PartSolutions
+searchCaveAll startRoom = 
+    do agenda <- initAgenda startRoom
+       allSolutions agenda S.empty M.empty
+
+initAgenda :: String -> CaveContext Agenda
+initAgenda startID = 
+    do let startState = emptySearchState startID
+       b <- estimateBenefit startState 
+       return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar ::  Agenda -> ExploredStates -> CaveContext (Maybe Agendum)
+aStar agenda closed 
+    -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ " : trail " ++ (show $ _trail $ snd $ P.findMax agenda) ++ " : closed " ++ (show closed)) False = undefined
+    -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
+    | P.null agenda = return Nothing
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            reachedGoal <- isGoal currentAgendum
+            let cl = (reached, currentAgendum ^. trailBenefit)
+            if reachedGoal
+            then return (Just currentAgendum)
+            else if (cl `S.member` closed)
+                 then aStar (P.deleteMax agenda) closed
+                 else aStar newAgenda (S.insert cl closed)
+
+allSolutions ::  Agenda -> ExploredStates -> PartSolutions -> CaveContext PartSolutions
+allSolutions agenda closed foundSolutions
+    | P.null agenda = return foundSolutions
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            reachedGoal <- isGoal currentAgendum
+            let cl = (reached, currentAgendum ^. trailBenefit)
+            newFoundSolutions <- includeAgendum foundSolutions currentAgendum
+            if reachedGoal
+            then allSolutions (P.deleteMax agenda) closed newFoundSolutions 
+            else if (cl `S.member` closed)
+                 then allSolutions (P.deleteMax agenda) closed foundSolutions 
+                 else allSolutions newAgenda (S.insert cl closed) newFoundSolutions 
+
+
+candidates :: Agendum -> ExploredStates -> CaveContext (Q.Seq Agendum)
+candidates agendum closed = 
+    do  let candidate = agendum ^. current
+        let previous = agendum ^. trail
+        let prevBenefit = agendum ^. trailBenefit
+        succs <- successors candidate
+        succAgs <- mapM (makeAgendum previous prevBenefit) succs
+        let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
+        return nonloops
+
+emptySearchState :: RoomID -> SearchState
+emptySearchState startID = SearchState 
+    { _currentRoom = startID
+    , _currentTime = 0
+    , _openValves = S.empty 
+    }
+
+currentFlow :: SearchState -> CaveContext Int
+currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. openValves
+       let presentRooms = cave `M.restrictKeys` valves
+       return $ sumOf (folded . flowRate) presentRooms
+
+timeOf :: SearchState -> Int
+timeOf state = state ^. currentTime
+
+successors :: SearchState -> CaveContext (Q.Seq SearchState)
+successors state = 
+    do isFF <- isFullFlow state
+       cave <- asks getCave 
+       timeLimit <- asks getTimeLimit
+       let here = state ^. currentRoom
+       let opened = state ^. openValves
+       let now = state ^. currentTime
+       let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
+       let moves = [ SearchState 
+                         { _currentRoom = (t ^. tunnelTo)
+                         , _currentTime = now + (t ^. tunnelLength)
+                         , _openValves = opened
+                         }
+                   | t <- remaining
+                   , now + (t ^. tunnelLength) <= timeLimit
+                   ]
+       let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0 && now < timeLimit
+                   then [SearchState { _currentRoom = here, _currentTime = now + 1, _openValves = S.insert here opened }]
+                   else []
+       let nexts = if null opens then moves else opens
+       let nexts' =  if null nexts
+                     then [ SearchState 
+                             { _currentRoom = here
+                             , _currentTime = timeLimit
+                             , _openValves = opened
+                             } ]
+                     else nexts
+       let succs = Q.fromList nexts'
+       if isFF || (Q.null succs)
+          then return $ Q.singleton (state & currentTime .~ timeLimit)
+          else return succs
+
+
+estimateBenefit :: SearchState -> CaveContext Int
+estimateBenefit here =  
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeOf here)
+       cf <- currentFlow here
+       sortedValves <- asks getSortedRooms
+       let opened = here ^. openValves
+       let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+makeAgendum :: Q.Seq SearchState -> Int -> SearchState -> CaveContext Agendum
+makeAgendum previous prevBenefit newState = 
+    do predicted <- estimateBenefit newState --  (Q.length previous)
+       -- cf <- currentFlow newState
+       oldFlow <- lastFlow previous (timeOf newState)
+       let newTrail = previous |> newState
+       let incurred = prevBenefit + oldFlow
+       return Agendum { _current = newState
+                      , _trail = newTrail
+                      , _trailBenefit = incurred
+                      , _benefit = incurred + predicted
+                      }
+
+lastFlow :: Q.Seq SearchState -> Int -> CaveContext Int
+lastFlow Q.Empty _ = return 0
+lastFlow (_ :|> previous) newTime = 
+  do cf <- currentFlow previous
+     let dt = newTime - (timeOf previous)
+     return (cf * dt)
+
+isGoal :: Agendum -> CaveContext Bool
+isGoal agendum = 
+  do timeLimit <- asks getTimeLimit
+     let s = agendum ^. current
+     return $ (timeOf s) == timeLimit
+
+isFullFlow :: SearchState -> CaveContext Bool
+isFullFlow state = 
+  do cave <- asks getCave
+     cf <- currentFlow state
+     let ff = sumOf (folded . flowRate) cave
+     return (cf == ff)
+
+compress :: Cave -> Cave
+compress cave = M.mapWithKey (compressRoom cave) cave
+
+compressRoom :: Cave -> RoomID -> Room -> Room
+compressRoom cave here room = room & tunnels .~ t'
+  where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
+
+reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
+reachableFrom _ [] _ routes = routes
+reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
+  | here `S.member` found = reachableFrom cave boundary found routes
+  | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
+  where exits = (cave ! here) ^. tunnels
+        exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
+        legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
+        routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
+                  then routes
+                  else S.insert tunnel routes
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser (S.Set Tunnel)
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+    -- where roomify v ts = Room {flowRate = v, tunnels = ts }
+tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input = 
+  case parseOnly caveP input of
+    Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+    Right cave -> cave
\ No newline at end of file
diff --git a/advent16/MainSubsets.hs b/advent16/MainSubsets.hs
new file mode 100644 (file)
index 0000000..b927188
--- /dev/null
@@ -0,0 +1,309 @@
+-- Writeup at https://work.njae.me.uk/2022/12/17/advent-of-code-2022-day-16/
+
+import Debug.Trace
+
+import AoC
+import Data.Text (Text)
+import qualified Data.Text.IO as TIO
+import Data.Attoparsec.Text hiding (take, D)
+import Control.Applicative
+import qualified Data.PQueue.Prio.Max as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+import qualified Data.Map.Strict as M
+import Data.Map.Strict ((!))
+-- import Data.Sequence ((|>), Seq((:|>)), ViewR ((:>))) 
+import Data.Sequence ( (|>), Seq((:|>)) )
+import Data.List
+-- import Data.List.Split (chunksOf)
+import Data.Ord
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+
+type RoomID = String
+
+data Tunnel = Tunnel { _tunnelTo :: RoomID, _tunnelLength :: Int}
+    deriving (Eq, Show, Ord)
+makeLenses ''Tunnel
+
+data Room = Room 
+    { _flowRate :: Int
+    , _tunnels :: S.Set Tunnel
+    } deriving (Eq, Show, Ord)
+makeLenses ''Room
+
+type Cave = M.Map RoomID Room
+data TimedCave = TimedCave { getCave :: Cave, getTimeLimit :: Int , getSortedRooms :: [RoomID]}
+
+type CaveContext = Reader TimedCave
+
+data SearchState = SearchState
+    { _currentRoom :: RoomID
+    , _currentTime :: Int
+    , _openValves :: S.Set RoomID
+    } deriving (Eq, Show, Ord)
+makeLenses ''SearchState
+
+data Agendum = 
+    Agendum { _current :: SearchState
+            , _trail :: Q.Seq SearchState
+            , _trailBenefit :: Int
+            , _benefit :: Int
+            } deriving (Show, Eq, Ord)
+makeLenses ''Agendum   
+
+type Agenda = P.MaxPQueue Int Agendum
+
+-- state, total flowed so far
+type ExploredStates = S.Set (SearchState, Int)
+
+type PartSolutions = M.Map (S.Set RoomID) Int
+
+
+main :: IO ()
+main = 
+  do  dataFileName <- getDataFileName
+      text <- TIO.readFile dataFileName
+      let expandedCave = successfulParse text
+      -- print cave
+      -- print $ reachableFrom cave [Tunnel "AA" 0] S.empty []
+      -- print $ compress cave
+      let cave = compress expandedCave
+      print $ part1 cave
+      print $ part2 cave
+
+part1, part2 :: Cave -> Int
+-- part1 :: Cave -> Int
+part1 cave = runSearch 30 cave
+part2 cave = maximum combinations 
+  where rawSolutions = runSearchAll 26 cave
+        solutionList = M.toList rawSolutions
+        combinations = [ (f1 + f2) 
+                       | (p, f1) <- solutionList
+                       , (e, f2) <- solutionList
+                       , p < e
+                       , S.disjoint p e
+                       ]
+
+includeAgendum :: PartSolutions -> Agendum -> CaveContext PartSolutions
+includeAgendum results agendum =
+  do cf <- currentFlow (agendum ^. current)
+     timeLimit <- asks getTimeLimit
+     let timeLeft = timeLimit - timeOf (agendum ^. current)
+     let remainingFlow = cf * timeLeft
+     let totalFlow = remainingFlow + agendum ^. trailBenefit
+     let visitedSet = agendum ^. current . openValves
+     let currentBest = M.findWithDefault 0 visitedSet results
+     if totalFlow > currentBest
+     then return (M.insert visitedSet totalFlow results)
+     else return results
+
+runSearch :: Int -> Cave -> Int
+runSearch timeLimit cave = maybe 0 _benefit result
+    where result = runReader (searchCave "AA") (TimedCave cave timeLimit sortedRooms)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+
+runSearchAll :: Int -> Cave -> PartSolutions
+runSearchAll timeLimit cave = result
+    where result = runReader (searchCaveAll "AA") (TimedCave cave timeLimit sortedRooms)
+          sortedRooms = sortOn (\r -> Down $ (cave ! r) ^. flowRate ) $ M.keys $ M.filter (\r -> r ^. flowRate > 0) cave
+
+
+searchCave :: String -> CaveContext (Maybe Agendum)
+searchCave startRoom = 
+    do agenda <- initAgenda startRoom
+       aStar agenda S.empty
+
+searchCaveAll :: String -> CaveContext PartSolutions
+searchCaveAll startRoom = 
+    do agenda <- initAgenda startRoom
+       allSolutions agenda S.empty M.empty
+
+initAgenda :: String -> CaveContext Agenda
+initAgenda startID = 
+    do let startState = emptySearchState startID
+       b <- estimateBenefit startState 
+       return $ P.singleton b Agendum { _current = startState, _trail = Q.empty, _trailBenefit = 0, _benefit = b}
+
+aStar ::  Agenda -> ExploredStates -> CaveContext (Maybe Agendum)
+aStar agenda closed 
+    -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda)) False = undefined
+    -- | trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ++ " : foundFlow " ++ (show $ _trailBenefit $ snd $ P.findMax agenda) ++ " : trail " ++ (show $ _trail $ snd $ P.findMax agenda) ++ " : closed " ++ (show closed)) False = undefined
+    -- | trace ("Peeping " ++ (show $ P.findMax agenda)) False = undefined
+    | P.null agenda = return Nothing
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            reachedGoal <- isGoal currentAgendum
+            let cl = (reached, currentAgendum ^. trailBenefit)
+            if reachedGoal
+            then return (Just currentAgendum)
+            else if (cl `S.member` closed)
+                 then aStar (P.deleteMax agenda) closed
+                 else aStar newAgenda (S.insert cl closed)
+
+allSolutions ::  Agenda -> ExploredStates -> PartSolutions -> CaveContext PartSolutions
+allSolutions agenda closed foundSolutions
+    | P.null agenda = return foundSolutions
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum closed
+            let newAgenda = foldl' (\q a -> P.insert (_benefit a) a q) (P.deleteMax agenda) nexts
+            reachedGoal <- isGoal currentAgendum
+            let cl = (reached, currentAgendum ^. trailBenefit)
+            newFoundSolutions <- includeAgendum foundSolutions currentAgendum
+            if reachedGoal
+            then allSolutions (P.deleteMax agenda) closed newFoundSolutions 
+            else if (cl `S.member` closed)
+                 then allSolutions (P.deleteMax agenda) closed foundSolutions 
+                 else allSolutions newAgenda (S.insert cl closed) newFoundSolutions 
+
+
+candidates :: Agendum -> ExploredStates -> CaveContext (Q.Seq Agendum)
+candidates agendum closed = 
+    do  let candidate = agendum ^. current
+        let previous = agendum ^. trail
+        let prevBenefit = agendum ^. trailBenefit
+        succs <- successors candidate
+        succAgs <- mapM (makeAgendum previous prevBenefit) succs
+        let nonloops = Q.filter (\s -> (s ^. current, s ^. trailBenefit) `S.notMember` closed) succAgs
+        return nonloops
+
+emptySearchState :: RoomID -> SearchState
+emptySearchState startID = SearchState 
+    { _currentRoom = startID
+    , _currentTime = 0
+    , _openValves = S.empty 
+    }
+
+currentFlow :: SearchState -> CaveContext Int
+currentFlow state =
+    do cave <- asks getCave
+       let valves = state ^. openValves
+       let presentRooms = cave `M.restrictKeys` valves
+       return $ sumOf (folded . flowRate) presentRooms
+
+timeOf :: SearchState -> Int
+timeOf state = state ^. currentTime
+
+successors :: SearchState -> CaveContext (Q.Seq SearchState)
+successors state = 
+    do isFF <- isFullFlow state
+       cave <- asks getCave 
+       timeLimit <- asks getTimeLimit
+       let here = state ^. currentRoom
+       let opened = state ^. openValves
+       let now = state ^. currentTime
+       let remaining = S.toList $ S.filter (\t -> (t ^. tunnelTo) `S.notMember` opened) ((cave ! here) ^. tunnels)
+       let moves = [ SearchState 
+                         { _currentRoom = (t ^. tunnelTo)
+                         , _currentTime = now + (t ^. tunnelLength)
+                         , _openValves = opened
+                         }
+                   | t <- remaining
+                   , now + (t ^. tunnelLength) <= timeLimit
+                   ]
+       let opens = if here `S.notMember` opened && (cave ! here) ^. flowRate > 0 && now < timeLimit
+                   then [SearchState { _currentRoom = here, _currentTime = now + 1, _openValves = S.insert here opened }]
+                   else []
+       let nexts = if null opens then moves else opens
+       let nexts' =  if null nexts
+                     then [ SearchState 
+                             { _currentRoom = here
+                             , _currentTime = timeLimit
+                             , _openValves = opened
+                             } ]
+                     else nexts
+       let succs = Q.fromList nexts'
+       if isFF || (Q.null succs)
+          then return $ Q.singleton (state & currentTime .~ timeLimit)
+          else return succs
+
+estimateBenefit :: SearchState -> CaveContext Int
+estimateBenefit here =  
+    do cave <- asks getCave
+       timeLimit <- asks getTimeLimit
+       let timeRemaining = timeLimit - (timeOf here)
+       cf <- currentFlow here
+       sortedValves <- asks getSortedRooms
+       let opened = here ^. openValves
+       let sortedClosedValves = [(cave ! v) ^. flowRate | v <- sortedValves, v `S.notMember` opened]
+       let otherValveFlows = sum $ zipWith (*) [timeRemaining, (timeRemaining - 2) .. 0] sortedClosedValves
+       return $ (cf * timeRemaining) + otherValveFlows
+
+makeAgendum :: Q.Seq SearchState -> Int -> SearchState -> CaveContext Agendum
+makeAgendum previous prevBenefit newState = 
+    do predicted <- estimateBenefit newState --  (Q.length previous)
+       -- cf <- currentFlow newState
+       oldFlow <- lastFlow previous (timeOf newState)
+       let newTrail = previous |> newState
+       let incurred = prevBenefit + oldFlow
+       return Agendum { _current = newState
+                      , _trail = newTrail
+                      , _trailBenefit = incurred
+                      , _benefit = incurred + predicted
+                      }
+
+lastFlow :: Q.Seq SearchState -> Int -> CaveContext Int
+lastFlow Q.Empty _ = return 0
+lastFlow (_ :|> previous) newTime = 
+  do cf <- currentFlow previous
+     let dt = newTime - (timeOf previous)
+     return (cf * dt)
+
+isGoal :: Agendum -> CaveContext Bool
+isGoal agendum = 
+  do timeLimit <- asks getTimeLimit
+     let s = agendum ^. current
+     return $ (timeOf s) == timeLimit
+
+isFullFlow :: SearchState -> CaveContext Bool
+isFullFlow state = 
+  do cave <- asks getCave
+     cf <- currentFlow state
+     let ff = sumOf (folded . flowRate) cave
+     return (cf == ff)
+
+compress :: Cave -> Cave
+compress cave = M.mapWithKey (compressRoom cave) cave
+
+compressRoom :: Cave -> RoomID -> Room -> Room
+compressRoom cave here room = room & tunnels .~ t'
+  where t' = reachableFrom cave [Tunnel here 0] S.empty S.empty
+
+reachableFrom :: Cave -> [Tunnel] -> S.Set RoomID -> S.Set Tunnel -> S.Set Tunnel
+reachableFrom _ [] _ routes = routes
+reachableFrom cave (tunnel@(Tunnel here len):boundary) found routes
+  | here `S.member` found = reachableFrom cave boundary found routes
+  | otherwise = reachableFrom cave (boundary ++ (S.toList legs)) (S.insert here found) routes'
+  where exits = (cave ! here) ^. tunnels
+        exits' = S.filter (\t -> (t ^. tunnelTo) `S.notMember` found) exits
+        legs = S.map (\t -> t & tunnelLength .~ (len + 1)) exits'
+        routes' = if (len == 0) || ((cave ! here) ^. flowRate) == 0
+                  then routes
+                  else S.insert tunnel routes
+
+-- Parse the input file
+
+caveP :: Parser Cave
+valveP :: Parser (RoomID, Room)
+roomP :: Parser Room
+tunnelsP :: Parser (S.Set Tunnel)
+tunnelTextP :: Parser Text
+
+caveP = M.fromList <$> valveP `sepBy` endOfLine
+valveP = (,) <$> ("Valve " *> (many1 letter)) <*> roomP
+roomP = Room <$> (" has flow rate=" *> decimal) <*> (tunnelTextP *> tunnelsP)
+    -- where roomify v ts = Room {flowRate = v, tunnels = ts }
+tunnelsP = (S.fromList . (fmap (flip Tunnel 1))) <$> (many1 letter) `sepBy` ", "
+tunnelTextP = "; tunnels lead to valves " <|> "; tunnel leads to valve "
+
+successfulParse :: Text -> Cave
+successfulParse input = 
+  case parseOnly caveP input of
+    Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
+    Right cave -> cave
\ No newline at end of file
diff --git a/advent16/a16-solution.dot.png b/advent16/a16-solution.dot.png
new file mode 100644 (file)
index 0000000..a27adfa
Binary files /dev/null and b/advent16/a16-solution.dot.png differ
diff --git a/advent16/a16.dot b/advent16/a16.dot
new file mode 100644 (file)
index 0000000..295f779
--- /dev/null
@@ -0,0 +1,94 @@
+graph G {
+AA -- DZ;
+AA -- EI;
+AA -- RO;
+AA -- VJ;
+AA -- VQ;
+AD -- BK;
+AD -- RC;
+AJ -- JW;
+AJ -- MK;
+AJ -- QX;
+AJ -- TR;
+AJ [style = filled, color = grey, label="AJ: 12"];
+AV -- AX;
+AV -- PI;
+AX -- HP;
+AX -- TG;
+AX [style = filled, color = grey, label="AX: 5"];
+AZ -- GJ;
+AZ -- ZR;
+BK -- PI;
+CG -- FF;
+CG -- QX;
+CG -- RV;
+CG -- SU;
+CG -- TI;
+CG [style = filled, color = grey, label="CG: 10"];
+DZ -- VO;
+EI -- RV;
+EQ -- RC;
+EQ -- YJ;
+ER -- QO;
+ER -- ZK;
+ET -- HP;
+ET -- ZR;
+EU -- GJ;
+EU -- PI;
+FF -- ZL;
+FR -- TF;
+FR -- ZK;
+FR [style = filled, color = grey, label="FR: 22"];
+FV -- KV;
+FV -- TX;
+FV [style = filled, color = grey, label="FV: 23"];
+GJ -- TG;
+GJ -- YJ;
+GJ -- ZJ;
+GJ [style = filled, color = grey, label="GJ: 21"];
+GQ -- MF;
+GQ -- VD;
+HF -- JI;
+HF -- LM;
+JI -- VD;
+JW -- YI;
+KU -- TC;
+KU -- TF;
+KU -- VY;
+KU -- XL;
+KU -- YW;
+KU [style = filled, color = grey, label="KU: 9"];
+KV -- OF;
+KX -- PI;
+KX -- ZR;
+LM -- SU;
+LM -- UJ;
+LM -- VY;
+LM -- YI;
+LM [style = filled, color = grey, label="LM: 3"];
+LN -- TI;
+LN -- ZR;
+MF -- QO;
+MK -- YW;
+OF [style = filled, color = grey, label="OF: 19"];
+PI -- VQ;
+PI [style = filled, color = grey, label="PI: 4"];
+QO [style = filled, color = grey, label="QO: 24"];
+RC -- WR;
+RC [style = filled, color = grey, label="RC: 18"];
+RO -- TC;
+TR -- VD;
+TX -- WR;
+UJ -- VJ;
+VD -- VO;
+VD -- VS;
+VD [style = filled, color = grey, label="VD: 17"];
+VS -- XL;
+WI -- XO;
+WI -- ZJ;
+WI -- ZL;
+WI [style = filled, color = grey, label="WI: 13"];
+XO -- ZR;
+ZR [style = filled, color = grey, label="ZR: 11"];
+
+}