-- https://gist.github.com/abhin4v/8172534

module AStar where

import qualified Data.PQueue.Prio.Min as PQ
import qualified Data.HashSet as Set
import qualified Data.HashMap.Strict as Map
import Data.Hashable (Hashable)
import Data.List (foldl')
import Data.Maybe (fromJust)

{- cabal:
build-depends: base, pqueue, unordered-containers, hashable
-}

-- A* search: Finds the shortest path from a start node to a goal node using a heuristic function.
astarSearch :: (Eq a, Hashable a) =>
  a                     -- startNode: the node to start the search from
  -> (a -> Bool)        -- isGoalNode: a function to test if a node is the goal node
  -> (a -> [(a, Int)])  -- nextNodeFn: a function which calculates the next nodes for a current node
                        -- along with the costs of moving from the current node to the next nodes
  -> (a -> Int)         -- heuristic: a function which calculates the (approximate) cost of moving
                        -- from a node to the nearest goal node
  -> Maybe (Int, [a])   -- result: Nothing is no path is found else
                        -- Just (path cost, path as a list of nodes)
astarSearch startNode isGoalNode nextNodeFn heuristic =
  astar (PQ.singleton (heuristic startNode) (startNode, 0))
         Set.empty (Map.singleton startNode 0) Map.empty
  where
    -- pq: open set, seen: closed set, tracks: tracks of states
    astar pq seen gscore tracks
      -- If open set is empty then search has failed. Return Nothing
      | PQ.null pq = Nothing
      -- If goal node reached then construct the path from the tracks and node
      | isGoalNode node = Just (gcost, findPath tracks node)
      -- If node has already been seen then discard it and continue
      | Set.member node seen = astar pq' seen gscore tracks
      -- Else expand the node and continue
      | otherwise = astar pq'' seen' gscore' tracks'
      where
        -- Find the node with min f-cost
        (node, gcost) = snd . PQ.findMin $ pq

        -- Delete the node from open set
        pq' = PQ.deleteMin pq

        -- Add the node to the closed set
        seen' =  Set.insert node seen

        -- Find the successors (with their g and h costs) of the node
        -- which have not been seen yet
        successors =
          filter (\(s, g, _) ->
                    not (Set.member s seen') &&
                      (not (s `Map.member` gscore)
                        || g < (fromJust . Map.lookup s $ gscore)))
          $ successorsAndCosts node gcost

        -- Insert the successors in the open set
        pq'' = foldl' (\q (s, g, h) -> PQ.insert (g + h) (s, g) q) pq' successors

        gscore' = foldl' (\m (s, g, _) -> Map.insert s g m) gscore successors

        -- Insert the tracks of the successors
        tracks' = foldl' (\m (s, _, _) -> Map.insert s node m) tracks successors

    -- Finds the successors of a given node and their costs
    successorsAndCosts node gcost =
      map (\(s, g) -> (s, gcost + g, heuristic s)) . nextNodeFn $ node


    -- Constructs the path from the tracks and last node
    findPath tracks node =
      if Map.member node tracks
      then findPath tracks (fromJust . Map.lookup node $ tracks) ++ [node]
      else [node]