module Util.Search (
  astar
) where

import Data.List
import qualified Data.Map as M (delete, empty, insert, lookup, singleton)

import Util.PQueue

-- astar is the A* pathfinding algorithm, which is similar to Dijkstra's
-- but it uses an estimation function to guide the search.
-- The function returns the cheapest path from start to goal, or an
-- empty list if there is no such path.
-- * "next" is the function that generates the nodes adjacent to a node,
--   together with the costs of moving into those nodes.
-- * "est" is a function that estimates the total cost of moving from a
--   node to the goal.  It must never over-estimate the cost, but it must
--   always be >= 0.  Note: If est is always 0, then this algorithm
--   reduces to Dijkstra's algorithm.
-- * "is_goal" is a function that indicates if a node is a goal node.

-- Information about a node.  The node itself can be found at the
-- head of the path.
-- realcost is strict because it is always checked after node creation.
-- estcost is not strict because many nodes are discarded on basis of
-- realcost before they are ever used.
data Node a = Node { node :: a, realcost :: !Int, estcost :: Int,
                     parent_node :: Node a }

instance (Eq a) => Eq (Node a) where
  x == y = node x == node y

-- Node information is ordered by cost.  The path is not used.
-- The realcost is compared if the estcost are equal, in order to force
-- backtracking to start from the earliest point, rather than randomly
-- (which would cause repeated invalidation of the same nodes, thus
-- making the backtracking take exponential time).
instance Priority (Node a) where
  -- The lowest-cost node has highest priority, so we swap x and y.
  pcompare x y = {-# SCC "Node.pcompare" #-}
    case compare (estcost y) (estcost x) of
      EQ -> compare (realcost y) (realcost x)
      other -> other

node_path :: (Eq a) => Node a -> [a]
node_path n | (node . parent_node) n == node n = [node n]
node_path n = node n : node_path (parent_node n)

-- Make sure astar performs well on simple grids.
{-# SPECIALIZE astar :: ((Int, Int) -> (Int, Int) -> [((Int, Int), Int)]) -> ((Int, Int) -> Int) -> ((Int, Int) -> Bool) -> (Int, Int) -> [(Int, Int)] #-}

astar :: (Ord a) => (a -> a -> [(a, Int)]) -> (a -> Int) -> (a -> Bool) -> a -> [a]
astar next est is_goal start =
  astar_next (unitPQ startN) (M.singleton start startN) M.empty where
  startN = Node { node = start, parent_node = startN,
                  realcost = 0, estcost = est start }
  -- The "open" map contains nodes we have not yet examined.
  -- The "openq" priority queue contains the same nodes as open, but ordered.
  -- The "closed" map contains nodes we have examined and expanded.
  astar_next q _ _ | isEmptyPQ q = [] -- Ran out of nodes! No path to goal.
  astar_next openq open closed = {-# SCC "astar_next" #-}
    let cheapest = {-# SCC "cheapest" #-} peekPQ openq
        newq = {-# SCC "cheapest" #-} popPQ openq
        chnode = node cheapest
        pnode = node (parent_node cheapest)
        open' = {-# SCC "open'" #-} (flip M.delete) open chnode
        closed' = {-# SCC "closed'" #-} M.insert chnode cheapest closed
    in if is_goal chnode
      then reverse (node_path cheapest)
      else astar_process newq open' closed' cheapest (next pnode chnode)
  astar_process openq open closed _ [] = {-# SCC "astar_process" #-}
                                         astar_next openq open closed
  astar_process openq open closed parent ((n, c) : ns) =
    {-# SCC "astar_process" #-}
    let newnode = {-# SCC "newnode" #-} Node {
          node = n,
          parent_node = parent,
          realcost = realcost parent + c,
          estcost = realcost newnode + (est n)
        }
        nopen = M.insert n newnode open
        nopenq = addToPQ openq newnode
        ropenq oldn = addToPQ (delFromPQ openq oldn) newnode
    in case (flip M.lookup) open n of
         Nothing -> case (flip M.lookup) closed n of
           Nothing -> astar_process nopenq nopen closed parent ns -- keep
           Just bn -> if realcost bn <= realcost newnode
             then astar_process openq open closed parent ns  -- discard
             else -- keep new, discard old
                  astar_process nopenq nopen ((flip M.delete) closed n) parent ns
         Just bn -> if realcost bn <= realcost newnode
           then astar_process openq open closed parent ns  -- discard
           else -- keep new, discard old (by shadowing it in the map, with
                -- difficulty in PQ)
                astar_process (ropenq bn) nopen closed parent ns