ai-examples.hs

module AI.Search.Examples.Graph where

import           Control.Monad
import           Control.Monad.ST
import           Control.Applicative
import           Data.STRef
import           Data.Map (Map, (!))
import qualified Data.Map                                 as Map
import           Data.List (nub)
import           Data.Graph.Inductive (LNode, LEdge, Gr)
import qualified Data.Graph.Inductive                     as G

import           AI.Search.Uninformed



type EdgeList   v e = [(v,e)]
type VertexList v e = [(v, EdgeList v e)]

posInf :: Double
posInf = 1 / 0

mkGraphProblem :: (Ord v) => v -> v -> VertexList v Double -> Problem v v
mkGraphProblem initial final vertices = Problem
  { probInitialState = initial
  , probSuccessor    = successor
  , probGoalTest     = goaltest
  , probCost         = cost
  }
 where

  graph = mkGraphMap vertices

  successor v = [ (x,x) | x <- neighbors graph v ]

  goaltest s  = s == final

  cost s a s' = case edge graph s s' of
    Nothing -> posInf
    Just c  -> c

-- | An abstraction around the graphs from fgl, allowing nodes
-- to be referred to via their labels, as opposed to by an abstract
-- node reference. This requires that labels are unique (a constraint
-- not imposed in fgl).
data Graph v e = Graph
  { graphRep  :: G.Gr v e
  , label2key :: Map v Int
  , key2label :: Map Int v
  }
  deriving (Show)

-- | Return the label of an edge between two nodes of a graph, or
-- return 'Nothing' if no such edge exists.
edge :: (Ord v) => Graph v e -> v -> v -> Maybe e
edge (Graph gr label2key key2label) v1 v2 = result
 where
  k1 = label2key ! v1
  k2 = label2key ! v2
  (Just ctxt, _) = G.match k1 gr
  (adj, _, _, _) = ctxt
  result = lookup k2 [ (k,e) | (e,k) <- adj ]

-- | Return the neighbors of a given vertex in a graph.
neighbors :: (Ord v) => Graph v e -> v -> [v]
neighbors (Graph gr label2key key2label) node =
  case Map.lookup node label2key of
    Nothing -> error "AI.Search.Example.Graph.neighbors"
    Just k  -> nub [ key2label ! v | v <- G.neighbors gr k ]

-- | Make an undirected graph from a vertex list.
mkGraphMap :: (Ord v) => VertexList v e -> Graph v e
mkGraphMap vertices = runST $ do
  labelToKey <- newSTRef Map.empty
  keyToLabel <- newSTRef Map.empty
  nextKey    <- newSTRef 0
  edgeList   <- newSTRef []
  mkGraphMap' labelToKey keyToLabel nextKey edgeList
 where
  mkGraphMap' labelToKeyRef keyToLabelRef nextKeyRef edgeListRef = do
    forM_ vertices $ \(u, conxns) -> do
      registerVertex u
      forM_ conxns $ \(v, e) -> do
        registerVertex v
        keys <- readSTRef labelToKeyRef
        let uKey = keys ! u
            vKey = keys ! v
        modifySTRef edgeListRef (\l -> (uKey, vKey, e) : (vKey, uKey, e) : l)
    label2key <- readSTRef labelToKeyRef
    key2label <- readSTRef keyToLabelRef
    edges     <- readSTRef edgeListRef
    let graph = G.mkGraph (Map.toList key2label) edges
    return $ Graph graph label2key key2label
   where
    registerVertex label = do
      labels <- readSTRef labelToKeyRef
      when (Map.notMember label labels) $ do
        key <- readSTRef nextKeyRef
        modifySTRef labelToKeyRef (Map.insert label key)
        modifySTRef keyToLabelRef (Map.insert key label)
        modifySTRef nextKeyRef (+1)


{-------------------------------------
  Australia Example
--------------------------------------}

--australia :: Gr String Double
--australia = mkGraphMap

aus :: VertexList String Double
aus = 
    [ (t,   [])
    , (sa,  [(wa,1), (nt,1), (q,1), (nsw,1), (v,1)])
    , (nt,  [(wa,1), (q,1)])
    , (nsw, [(q, 1), (v,1)]) ]
 where
  sa  = "South Australia"
  wa  = "Western Australia"
  nt  = "Northern Territory"
  q   = "Queensland"
  nsw = "New South Wales"
  v   = "Victoria"
  t   = "Tasmania"