Initial attempt at optimising day 23
[advent-of-code-23.git] / advent23 / Main.hs
index 82982614816689c43b009398f45278bf39ef30d2..7a2488445514dca16baaecf91edf30a59250ffb9 100644 (file)
@@ -6,8 +6,14 @@ import AoC
 import Linear -- (V2(..), (^+^))
 import qualified Data.Set as S
 import qualified Data.Map.Strict as M
+import qualified Data.Sequence as Q
+import Data.Sequence (Seq( (:|>), (:<|) ) ) 
 import Control.Lens
 import Data.List (foldl')
+import Control.Monad.Reader
+import qualified Data.PQueue.Prio.Max as P
+import Data.Foldable
+import Data.Maybe 
 
 data Slide = SlideLeft | SlideRight | SlideUp | SlideDown
   deriving (Show, Eq)
@@ -26,6 +32,26 @@ data CompressedPath = CPath { _nextPos :: Position, _pathLen :: Int }
 makeLenses ''CompressedPath
 
 type CompressedMap = M.Map Position [CompressedPath]
+data Mountain = Mountain
+  { _paths :: CompressedMap
+  , _start :: Position
+  , _goal :: Position
+  } deriving (Eq, Show)
+makeLenses ''Mountain
+
+type MountainContext = Reader Mountain
+
+data Agendum = 
+    Agendum { _current :: Position
+            , _trail :: Q.Seq Position
+            , _trailCost :: Int
+            , _cost :: Int
+            } deriving (Show, Eq)
+makeLenses ''Agendum  
+
+type Agenda = P.MaxPQueue Int Agendum
+
+type ExploredStates = M.Map Position Int
 
 
 main :: IO ()
@@ -33,16 +59,18 @@ main =
   do  dataFileName <- getDataFileName
       text <- readFile dataFileName
       let (forest, slides, start, end) = mkGrid text
+      -- print $ compress slides forest start end
       print $ part1 slides forest start end
       print $ part2 slides forest start end
 
 part1, part2 :: Slides -> Grid -> Position -> Position -> Int
-part1 slides forest start end = maximum $ fmap (pathLength cMap) paths
+part1 slides forest start end = searchCompressed $ Mountain cMap start end
   where cMap = compress slides forest start end
-        paths = searchCompressed cMap end [] [[start]]
-part2 _ forest start end = maximum $ fmap (pathLength cMap) paths
+part2 slides forest start end = searchCompressed $ Mountain cMap start end
   where cMap = compress M.empty forest start end
-        paths = searchCompressed cMap end [] [[start]]
+-- part2 _ forest start end = maximum $ fmap (pathLength cMap) paths
+--   where cMap = compress M.empty forest start end
+--         paths = searchCompressed cMap start end
 
 adjacents :: Position -> Slides -> Grid -> [Position]
 adjacents here slides forest = filter (`S.notMember` forest) $ fmap (here ^+^) deltas
@@ -92,23 +120,112 @@ compress slides forest start end = foldl' go compressed0 iPoints
         go com here = search slides forest iPoints com $ fmap (: [here]) $ adjacents here slides forest
 
 
-searchCompressed :: CompressedMap -> Position -> [[Position]] -> [[Position]] -> [[Position]]
--- searchCompressed _ _ _ (c:_) _ | DT.trace (show c) False = undefined
-searchCompressed _ _ found [] = found
-searchCompressed map goal found (current:agenda) 
-  | head current == goal = searchCompressed map goal (current:found) agenda
-  | otherwise = searchCompressed map goal found (nextPositions ++ agenda)
-  where neighbours0 = map M.! (head current)
-        neighbours = neighbours0 ^.. folded . filtered ((`notElem` current) . _nextPos)
-        nextPositions = fmap ((: current) . _nextPos) neighbours
+-- searchCompressed :: CompressedMap -> Position -> [[Position]] -> [[Position]] -> [[Position]]
+-- -- searchCompressed _ _ _ (c:_) _ | DT.trace (show c) False = undefined
+-- searchCompressed _ _ found [] = found
+-- searchCompressed map goal found (current:agenda) 
+--   | head current == goal = searchCompressed map goal (current:found) agenda
+--   | otherwise = searchCompressed map goal found (nextPositions ++ agenda)
+--   where neighbours0 = map M.! (head current)
+--         neighbours = neighbours0 ^.. folded . filtered ((`notElem` current) . _nextPos)
+--         nextPositions = fmap ((: current) . _nextPos) neighbours
+
+searchCompressed :: Mountain -> Int
+searchCompressed mountain = maybe 0 _trailCost result
+  where result = runReader searchMountain mountain
+
+searchMountain :: MountainContext (Maybe Agendum)
+searchMountain = 
+  do agenda <- initAgenda
+     aStar agenda Nothing
+
+initAgenda :: MountainContext Agenda
+initAgenda = 
+   do s <- asks _start
+      c <- estimateCost Q.Empty s
+      let agendum = Agendum { _current = s, _trail = Q.empty, _trailCost = 0, _cost = c}
+      let agenda = P.singleton c agendum
+      return agenda
+
+aStar ::  Agenda -> (Maybe Agendum) -> MountainContext (Maybe Agendum)
+aStar agenda best 
+    -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ _trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
+    -- | DT.trace ("Peeping " ++ (show $ _current $ snd $ P.findMax agenda) ) False = undefined
+    -- | DT.trace ("Peeping " ++ (show $ snd $ P.findMax agenda) ) False = undefined
+    -- | DT.trace ("Peeping " ++ (show agenda) ) False = undefined
+    | P.null agenda = return best
+    | (fst $ P.findMax agenda) < maybe 0 _trailCost best = return best
+    | otherwise = 
+        do  let (_, currentAgendum) = P.findMax agenda
+            let reached = currentAgendum ^. current
+            nexts <- candidates currentAgendum 
+            let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMax agenda) nexts
+            reachedGoal <- isGoal reached
+            let best' = updateBest reachedGoal best currentAgendum
+            -- let closed' = M.insert reached (currentAgendum ^. trailCost) closed
+            if reachedGoal -- || (reached `S.member` closed)
+                then aStar (P.deleteMax agenda) best' -- closed'
+                else aStar newAgenda best' -- closed'
+
+updateBest :: Bool -> Maybe Agendum -> Agendum -> Maybe Agendum
+updateBest False current _ = current
+updateBest True Nothing best 
+ -- | DT.trace ("Nothing " ++ show best) False = undefined 
+ | otherwise = Just best
+updateBest True (Just current) best 
+  -- | DT.trace (show current ++ " " ++ show best) False = undefined 
+  | (current ^. trailCost) > (best ^. trailCost) = Just current
+  | otherwise = Just best
+
+
+
+
+candidates :: Agendum -> MountainContext (Q.Seq Agendum)
+candidates agendum = 
+  do  let here = agendum ^. current
+      let previous = agendum ^. trail
+      let prevCost = agendum ^. trailCost
+      ts <- asks _paths
+      let succs = Q.fromList $ ts M.! here
+      -- succs <- successors candidate
+      let nonloops = Q.filter (\s -> (s ^. nextPos) `notElem` previous) succs
+      mapM (makeAgendum previous prevCost here) nonloops
+
+
+makeAgendum :: (Q.Seq Position) -> Int -> Position -> CompressedPath -> MountainContext Agendum
+makeAgendum previous prevCost here step = 
+   do let newTrail = previous :|> here
+      predicted <- estimateCost newTrail $ step ^. nextPos
+      -- ts <- asks _trails
+      let incurred = prevCost + step ^. pathLen
+      return Agendum { _current = step ^. nextPos
+                     , _trail = newTrail
+                     , _trailCost = incurred
+                     , _cost = incurred + predicted
+                     }
+
+
+isGoal :: Position -> MountainContext Bool
+isGoal here = 
+  do goal <- asks _goal
+     return $ here == goal
+
+estimateCost :: Q.Seq Position -> Position -> MountainContext Int
+estimateCost r e = 
+  do ts <- asks _paths
+     let endCost = fromMaybe 0 $ maximumOf (folded . filtered ((`notElem` r) . _nextPos) . pathLen) $ ts M.! e
+     let res = S.fromList $ toList (r :|> e)
+     let otherPaths = concat $ M.elems $ ts `M.withoutKeys` res
+     let restCost = sumOf (folded . filtered ((`notElem` r) . _nextPos) . pathLen) otherPaths
+     return $ (restCost `div` 2) + endCost
 
-pathLength :: CompressedMap -> [Position] -> Int
-pathLength map ps = sum $ zipWith (stepLength map) ps $ tail ps
+-- pathLength :: CompressedMap -> [Position] -> Int
+-- pathLength map ps = sum $ zipWith (stepLength map) ps $ tail ps
 
-stepLength :: CompressedMap -> Position -> Position -> Int
-stepLength map here there = 
-  -- head $ (map M.! there) ^.. folded . filtered ((== here) . _nextPos) . pathLen
-  head $ (map M.! there) ^.. folded . filteredBy (nextPos . only here) . pathLen
+-- stepLength :: CompressedMap -> Position -> Position -> Int
+-- stepLength map here there = 
+--   -- head $ (map M.! there) ^.. folded . filtered ((== here) . _nextPos) . pathLen
+--   head $ (map M.! there) ^.. folded . filteredBy (nextPos . only here) . pathLen
 
 -- reading the map