+-- Writeup at https://work.njae.me.uk/2022/12/14/advent-of-code-2022-day-12/
+
+import AoC
+
+import qualified Data.PQueue.Prio.Min as P
+import qualified Data.Set as S
+import qualified Data.Sequence as Q
+-- import Data.Sequence ((<|), (|>), (><))
+import Data.Sequence ((|>))
+import Data.Foldable (foldl')
+import Data.Char
+import Control.Monad.Reader
+import Control.Lens hiding ((<|), (|>), (:>), (:<), indices)
+import Linear (V2(..), (^+^), (^-^))
+import Data.Array.IArray
+
+-- 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 Position = V2 Int -- r, c
+type Grid = Array Position Int
+
+data Mountain = Mountain
+ { _grid :: Grid
+ , _start :: Position
+ , _goal :: Position
+ } deriving (Eq, Ord, 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.MinPQueue Int Agendum
+
+type ExploredStates = S.Set Position
+
+main :: IO ()
+main =
+ do dataFileName <- getDataFileName
+ text <- readFile dataFileName
+ let mountain = mkMountain text
+ -- print mountain
+ print $ part1 mountain
+ print $ part2 mountain
+
+part1, part2 :: Mountain -> Int
+part1 mountain = maybe 0 _cost result
+ where s = mountain ^. start
+ result = runReader (searchMountain s) mountain
+
+part2 mountain = minimum results
+ where starts = possibleStarts mountain
+ results = fmap (runSearch mountain) starts
+
+runSearch :: Mountain -> Position -> Int
+runSearch mountain s = maybe maxCost _cost result
+ where result = runReader (searchMountain s) mountain
+ maxCost = length $ indices $ mountain ^. grid
+
+possibleStarts :: Mountain -> [Position]
+possibleStarts mountain = map fst $ filter ((== 0) . snd)
+ $ assocs $ mountain ^. grid
+
+mkMountain :: String -> Mountain
+mkMountain text = Mountain { _grid = grid, _start = s, _goal = g }
+ where rows = lines text
+ r = length rows - 1
+ c = (length $ head rows) - 1
+ grid0 = listArray ((V2 0 0), (V2 r c)) $ map mkCell $ concat rows
+ mkCell e = ord e - ord 'a'
+ s = head [i | i <- range ((V2 0 0), (V2 r c)), grid0 ! i == (mkCell 'S')]
+ g = head [i | i <- range ((V2 0 0), (V2 r c)), grid0 ! i == (mkCell 'E')]
+ grid = grid0 // [(s, mkCell 'a'), (g, mkCell 'z')]
+
+searchMountain :: Position -> MountainContext (Maybe Agendum)
+searchMountain startPos =
+ do agenda <- initAgenda startPos
+ aStar agenda S.empty
+
+initAgenda :: Position -> MountainContext Agenda
+initAgenda pos =
+ do c <- estimateCost pos
+ return $ P.singleton c Agendum { _current = pos, _trail = Q.empty, _trailCost = 0, _cost = c}
+
+aStar :: Agenda -> ExploredStates -> MountainContext (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.findMin agenda) ) False = undefined
+ | P.null agenda = return Nothing
+ | otherwise =
+ do let (_, currentAgendum) = P.findMin agenda
+ let reached = currentAgendum ^. current
+ nexts <- candidates currentAgendum closed
+ let newAgenda = foldl' (\q a -> P.insert (_cost a) a q) (P.deleteMin agenda) nexts
+ reachedGoal <- isGoal reached
+ if reachedGoal
+ then return (Just currentAgendum)
+ else if reached `S.member` closed
+ then aStar (P.deleteMin agenda) closed
+ else aStar newAgenda (S.insert reached closed)
+
+candidates :: Agendum -> ExploredStates -> MountainContext (Q.Seq Agendum)
+candidates agendum closed =
+ do let candidate = agendum ^. current
+ let previous = agendum ^. trail
+ let prevCost = agendum ^. trailCost
+ succs <- successors candidate
+ let nonloops = Q.filter (\s -> s `S.notMember` closed) succs
+ mapM (makeAgendum previous prevCost) nonloops
+
+
+makeAgendum :: Q.Seq Position -> Int -> Position -> MountainContext Agendum
+makeAgendum previous prevCost newPosition =
+ do predicted <- estimateCost newPosition
+ grid <- asks _grid
+ let newTrail = previous |> newPosition
+ let incurred = prevCost + 1
+ return Agendum { _current = newPosition
+ , _trail = newTrail
+ , _trailCost = incurred
+ , _cost = incurred + predicted
+ }
+
+isGoal :: Position -> MountainContext Bool
+isGoal here =
+ do goal <- asks _goal
+ return $ here == goal
+
+successors :: Position -> MountainContext (Q.Seq Position)
+successors here =
+ do grid <- asks _grid
+ let heightHere = grid ! here
+ let neighbours =
+ filter (\p -> (grid ! p) - heightHere <= 1)
+ $
+ filter (inRange (bounds grid))
+ [ here ^+^ delta
+ | delta <- [V2 -1 0, V2 1 0, V2 0 -1, V2 0 1]
+ ]
+ let succs = Q.fromList neighbours
+ return succs
+
+estimateCost :: Position -> MountainContext Int
+estimateCost here =
+ do goal <- asks _goal
+ let (V2 dr dc) = here ^-^ goal
+ return $ (abs dr) + (abs dc)