hbayes-0.4: Bayes.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{- | Discrete Bayesian Network Library.

It is a very preliminary version. It has only been tested on very simple
examples where it worked. It should be considered as experimental and not used
in any production work.

Look at the "Bayes.Examples" and "Bayes.Examples.Tutorial" in this package 
to see how to use the library.

In "Bayes.Examples.Influence" you'll find additional examples about influence diagrams.

-}
module Bayes(
  -- * Graph
  -- ** Graph classes
    Graph(..)
  , UndirectedGraph(..)
  , DirectedGraph(..)
  , FoldableWithVertex(..)
  , NamedGraph(..)
  -- ** Graph Monad
  , GraphMonad
  , GMState(..)
  , graphNode
  , runGraph
  , execGraph
  , evalGraph
  , emptyAuxiliaryState
  , getNewEmptyVariable
  , isRoot
  , rootNode
  , parentNodes
  , childrenNodes
  -- ** Support functions for Graph constructions
  , Vertex
  , Edge 
  , edge
  , newEdge
  , getVertex
  , edgeEndPoints
  , connectedGraph
  , dag
  , printGraphValues
  -- * SimpleGraph implementation
  -- ** The SimpleGraph type
  , DirectedSG
  , UndirectedSG
  , SBN(..)
  , varMap
  , displaySimpleGraph
  -- ** Bayesian network
  , BayesianNetwork(..)
  -- * Testing
  , testEdgeRemoval_prop
  , testVertexRemoval_prop
) where

import qualified Data.IntMap as IM
import qualified Data.Map as M
import Control.Monad.State.Strict
import Control.Monad.Writer.Strict
import Control.Applicative((<$>))
import Bayes.Factor hiding(isEmpty)
import Bayes.Factor.CPT(CPT(..))
import Bayes.Factor.MaxCPT(MAXCPT(..))
import Data.Maybe
import qualified Data.Map as Map
import qualified Data.Foldable as F
import qualified Data.Traversable as T 
import Control.Applicative 
import qualified Data.Set as Set

import Test.QuickCheck hiding ((.&.),Testable)
import Test.QuickCheck.Arbitrary
import Data.List(sort,intercalate,nub,foldl')
import Bayes.PrivateTypes hiding(isEmpty)
import GHC.Float(float2Double)

--import Debug.Trace
--debug a = trace (show a) a

-- | An implementation of the BayesianNetwork using the simple graph and no value for the edges
type SBN f = DirectedSG () f

-- | Bayesian network. g must be a directed graph and f a factor
type BayesianNetwork g f = g () f

instance Arbitrary (DirectedSG String String) where
  arbitrary = do 
    let createVertex g i = do 
          name <- arbitrary :: Gen String
          return $ addVertex (Vertex i) name g
        createEdge g (va,vb) = do 
          name <- arbitrary :: Gen String
          return $ addEdge (edge va vb) name g 

    nbVertex <- choose (1,8) :: Gen Int
    
    g <- foldM createVertex emptyGraph [1..nbVertex]

    let allPairs = [(Vertex x,Vertex y) | x <- [1..nbVertex], y <- [1..nbVertex], x /= y]
        anEdge (x,y) = arbitrary :: Gen Bool

    edges <- filterM anEdge allPairs

    foldM createEdge g edges

instance Arbitrary (DirectedSG () String) where
  arbitrary = do 
    let createVertex g i = do 
          name <- arbitrary :: Gen String
          return $ addVertex (Vertex i) name g
        createEdge g (va,vb) = do 
          return $ addEdge (edge va vb) () g 

    nbVertex <- choose (1,8) :: Gen Int
    
    g <- foldM createVertex emptyGraph [1..nbVertex]

    let allPairs = [(Vertex x,Vertex y) | x <- [1..nbVertex], y <- [1..nbVertex], x /= y]
        anEdge (x,y) = arbitrary :: Gen Bool

    edges <- filterM anEdge allPairs

    foldM createEdge g edges   

-- | Warning : the generated graph is not at all a bayesian network
-- The variables in the CPT have no reason to correspond to the edges
-- connected to that CPT.
-- Only the main variable (first variable) is linked to the right vertex
instance Factor f => Arbitrary (DirectedSG () f) where
  arbitrary = do 
    let createVertex g i = do 
          let value = fromJust $ factorWithVariables [DV (Vertex i) 2] [0.1,0.9]
          return $ addVertex (Vertex i) value g
        createEdge g (va,vb) = do 
          return $ addEdge (edge va vb) () g 

    nbVertex <- choose (1,8) :: Gen Int
    
    g <- foldM createVertex emptyGraph [1..nbVertex]

    let allPairs = [(Vertex x,Vertex y) | x <- [1..nbVertex], y <- [1..nbVertex], x /= y]
        anEdge (x,y) = arbitrary :: Gen Bool

    edges <- filterM anEdge allPairs

    foldM createEdge g edges

testEdgeRemoval_prop :: DirectedSG String String -> Property
testEdgeRemoval_prop g = (not . hasNoEdges) g ==> 
  let Just e = someEdge g
      Just (vs,ve) = edgeVertices g e
      Just bi = ingoing g ve
      Just bo = outgoing g vs
      g' = removeEdge e g 
      Just bi' = ingoing g' ve
      Just bo' = outgoing g' vs
  in
  (map (sort . (:) e ) [bi', bo'] == map sort [bi,bo]) &&
   (sort (allEdges g) == sort (e:allEdges g'))

testVertexRemoval_prop :: DirectedSG String String -> Property
testVertexRemoval_prop g = (not . hasNoVertices) g ==>
    let Just v = someVertex g
        Just bi = ingoing g v 
        Just bo = outgoing g v
        g' = removeVertex v g
        srcVertices = mapMaybe (startVertex g') bi
        dstVertices = mapMaybe (endVertex g') bo 
        isNotDstVertex = not . (v `elem`) . mapMaybe (endVertex g') . fromJust . outgoing g'
        isNotStartVertex = not . (v `elem`) . mapMaybe (startVertex g') . fromJust . ingoing g'
    in 
    (sort (allVertices g) == sort (v:allVertices g')) &&
      (all isNotDstVertex srcVertices) && (all isNotStartVertex dstVertices)


-- | Graph class used for graph processing algorithms.
-- A graph processing algorithm does not have to know how the graph is implemented nor if it is
-- directed or undirected
class Graph g where
    -- | Add a new vertex
    addVertex :: Vertex -> b -> g a b -> g a b
    -- | Remove a vertex
    removeVertex :: Vertex -> g a b -> g a b
    -- | Get the vertex value if the vertex is found in the graph
    vertexValue :: g a b -> Vertex -> Maybe b
    -- | Change the vertex value if the vertex is found in the graph
    changeVertexValue :: Vertex -> b -> g a b -> Maybe (g a b)
    -- | Generate a \"random\" vertex
    someVertex :: g a b -> Maybe Vertex

    -- | Check is the graph has no vertrex
    hasNoVertices :: g a b -> Bool

    -- | Generate all vertices
    allVertices :: g a b -> [Vertex]

    -- | Get all the values
    allVertexValues :: g a b -> [b]

    -- | Get all nodes
    allNodes :: g a b -> [(Vertex,b)]

    -- | Check if two vertices are linked by a vertex
    isLinkedWithAnEdge :: g a b -> Vertex -> Vertex -> Bool

    -- | Add an edge
    addEdge :: Edge -> a -> g a b  -> g a b

    -- | Remove an dedge
    removeEdge :: Edge -> g a b -> g a b

    -- | Vertices for an edge
    edgeVertices :: g a b -> Edge -> Maybe (Vertex,Vertex)

    -- | Edge value if the edge is found in the graph
    edgeValue :: g a b -> Edge -> Maybe a

    -- | Return a \"random\" edge
    someEdge :: g a b -> Maybe Edge

    -- | Check if the graph has no edges
    hasNoEdges :: g a b -> Bool

    -- | One extremity of the edge (which is the end only for directed edge)
    endVertex :: g a b -> Edge -> Maybe Vertex
    endVertex g e = do 
      (_,ve) <- edgeVertices g e
      return ve 
        
    -- | One extremity of the edge (which is the start only for directed edge)
    startVertex :: g a b -> Edge -> Maybe Vertex
    startVertex g e = do 
      (vs,_) <- edgeVertices g e
      return vs

    -- | All edges of the graph
    allEdges :: g a b -> [Edge]

    -- | All values of the graph
    allEdgeValues :: g a b -> [a]
   
    -- | Returns an empty graph
    emptyGraph :: g a b

    -- | Check if the graph is empty
    isEmpty :: g a b -> Bool
    isEmpty g = hasNoVertices g && hasNoEdges g

    -- | Check if the graph is oriented
    oriented :: g a b -> Bool

    -- | All the neighbors of a vertex
    neighbors :: g a b -> Vertex -> Maybe [Vertex]

-- | A named graph is a graph where the vertices have a name.
-- This name is not a vertex value. Putting this name in the vertex value
-- would make algorithm less readable.
-- A vertex name is only useful to display the graph.
-- Labeled graph has a different meaning in graph theory.
class Graph g => NamedGraph g where
    -- | Add a vertex with a vertex name in addition to the value
    addLabeledVertex :: String -> Vertex -> b -> g a b -> g a b
    -- | Returns the vertex label
    vertexLabel :: g a b -> Vertex -> Maybe String


-- | Undirected graph
class Graph g => UndirectedGraph g where
    edges :: g a b -> Vertex -> Maybe [Edge]

-- | Directed graph
class Graph g => DirectedGraph g where
    ingoing :: g a b -> Vertex -> Maybe [Edge]
    outgoing :: g a b -> Vertex -> Maybe [Edge]

-- | Return the parents of a node
parentNodes :: DirectedGraph g => g a b -> Vertex -> [Vertex]
parentNodes g v = maybe [] id $ do 
  ie <- ingoing g v
  p <- mapM (startVertex g) ie
  return p

-- | Return the children of a node
childrenNodes :: DirectedGraph g => g a b -> Vertex -> [Vertex]
childrenNodes g v = maybe [] id $ do 
  ie <- outgoing g v
  p <- mapM (endVertex g) ie
  return p

isRoot :: DirectedGraph g => g a b -> Vertex -> Bool
{-# INLINE isRoot #-}
isRoot g v =
  case ingoing g v of 
    Just [] -> True 
    _ -> False

-- | Get the root node for the graph
rootNode :: DirectedGraph g => g a b -> Maybe Vertex
rootNode g = 
  let someRoots = filter (isRoot g) . allVertices $ g
  in 
  case someRoots of 
    (h:l) -> Just h 
    _ -> Nothing

-- | Check if the graph is a directed Acyclic graph
dag :: DirectedGraph g => g a b -> Bool 
dag g = case rootNode g of 
  Nothing -> isEmpty g 
  Just r -> dag (removeVertex r g)

-- | Check if the graph is connected
connectedGraph :: Graph g => g a b -> Bool 
connectedGraph g = 
  let visited = visitVertex g (Set.empty) ([fromJust $ someVertex g])
      vertices = Set.fromList $ allVertices g
      equalSets a b = Set.isSubsetOf a b && Set.isSubsetOf b a
  in 
  equalSets visited vertices
 where 
  visitVertex _ visited [] = visited
  visitVertex theGraph visited (current:n) = 
    if Set.member current visited
      then 
        visitVertex theGraph visited n
      else
        let n' = fromJust $ neighbors theGraph current
        in
        visitVertex theGraph (Set.insert current visited) (n ++ n')



-- | Create an edge description
edge :: Vertex -> Vertex -> Edge 
{-# INLINE edge #-}
edge a b = Edge a b

-- | Endpoints of an edge
edgeEndPoints :: Edge -> (Vertex,Vertex)
edgeEndPoints (Edge va vb) = (va,vb)





-- | Class used to share as much code as possible between
-- directed and undirected graphs without
-- implementing an undirected graph as a graph where
-- we have a directed edge in both directions 
class NeighborhoodStructure n where
  -- | Return an empty neighborhood
  emptyNeighborhood :: n 
  -- | Ingoing edges
  ingoingNeighbors :: n -> [Edge]
  -- | Outgoing edge
  outgoingNeighbors :: n -> [Edge]
  -- | Remove an edge
  removeNeighborsEdge :: Edge -> n -> n
  -- | Add an outgoing edge
  addOutgoingEdge :: Edge -> n -> n
  -- Add in ingoing edge
  addIngoingEdge :: Edge -> n -> n

-- | Directed neighborhood structure for a vertex
instance NeighborhoodStructure DE where
  emptyNeighborhood = DE [] []
  ingoingNeighbors (DE i _) = i
  outgoingNeighbors (DE _ o) = o 
  removeNeighborsEdge e (DE i o) = 
    let i' = filter (/= e) i
        o' = filter (/= e) o 
    in 
    DE i' o'
  addOutgoingEdge e (DE i o) = DE i (e:o)
  addIngoingEdge e (DE i o) = DE (e:i) o

-- | Undirected neighborhood structure for a vertex
instance NeighborhoodStructure UE where
  emptyNeighborhood = UE []
  ingoingNeighbors (UE e) = e
  outgoingNeighbors (UE e) = e
  removeNeighborsEdge e (UE l) = 
    let l' = filter (/= e) l
    in 
    UE l'
  addOutgoingEdge e (UE l) = UE (e:l)
  addIngoingEdge e (UE l) = UE (e:l)



-- | Directed simple graph
type DirectedSG = SimpleGraph DE

-- | Undirected simple graph
type UndirectedSG = SimpleGraph UE

-- | Get the variable name mapping
varMap :: SimpleGraph n e v -> M.Map String Vertex 
varMap (SP _ _ n) = M.fromList . map (\(i,s) -> (s, Vertex i)) . IM.toList $ n

instance (Eq a, Eq b) => Eq (SimpleGraph DE a b) where
    (==) (SP a b _) (SP a' b' _) = a == a' && b == b'

-- | An empty simple graph
emptySimpleGraph = SP M.empty IM.empty IM.empty

-- | Used to prevent adding duplicates to a graph
noRedundancy new old = old

instance FactorContainer (SimpleGraph local edge) where 
   changeFactor = changeFactorInFunctor 

instance Functor (SimpleGraph local edge) where 
  fmap f (SP em vm nm) = SP em (IM.map (\(l,d) -> (l, f d)) vm) nm

instance F.Foldable (SimpleGraph local edge) where
  foldr f c (SP _ vm _) = IM.foldr (\(_,d) s -> f d s) c vm

instance T.Traversable (SimpleGraph local edge) where
  traverse f (SP em vm nm) = 
    let l = IM.toList vm -- [(IM.Key, (DE, String))]
        onTriple f (k,(l,v)) = (\z -> (k,(l,z))) <$> f v
        l' = T.traverse (onTriple f) l -- f [(k,(l,z))]
        result y =  (\x -> SP em (IM.fromList x) nm) <$> y
    in 
    result l'

-- | The foldable class is limited. For a graph g we may need the vertex in addition to the value
class FoldableWithVertex g where
  -- | Fold with vertex 
  foldrWithVertex :: (Vertex -> a -> b -> b) -> b -> g c a -> b 
  foldlWithVertex' :: (b -> Vertex -> a -> b) -> b -> g c a -> b 

instance FoldableWithVertex (SimpleGraph local) where
  foldrWithVertex f s (SP _ vm _) = IM.foldrWithKey (\k (_,v) y -> f (Vertex k) v y) s vm
  foldlWithVertex' f s (SP _ vm _) = IM.foldlWithKey' (\y k (_,v)  -> f y (Vertex k) v) s vm

_addLabeledVertex vertexName vert@(Vertex v) value (SP em vm name) =
  let vm' = IM.insertWith' noRedundancy v (emptyNeighborhood,value) vm
      name' = IM.insert v vertexName name 
  in
  SP em vm' name'

_vertexLabel (SP _ _ name) (Vertex v) = IM.lookup v name

instance NamedGraph DirectedSG where
      addLabeledVertex = _addLabeledVertex
      vertexLabel = _vertexLabel

instance NamedGraph UndirectedSG where
      addLabeledVertex = _addLabeledVertex
      vertexLabel = _vertexLabel

-- | SimpleGraph is an instance of Graph.
instance Graph DirectedSG where
    addVertex = _addVertex
    removeVertex = _removeVertex
    vertexValue = _vertexValue
    changeVertexValue = _changeVertexValue
    someVertex = _someVertex
    hasNoVertices = _hasNoVertices
    allVertices = _allVertices
    allVertexValues = _allVertexValues
    allNodes = _allNodes
    isLinkedWithAnEdge = _isLinkedWithAnEdge
    addEdge = _addEdge
    removeEdge = _removeEdge
    edgeVertices = _edgeVertices
    edgeValue = _edgeValue
    someEdge = _someEdge
    hasNoEdges = _hasNoEdges
    allEdges = _allEdges
    allEdgeValues = _allEdgeValues
    emptyGraph = _emptyGraph
    oriented _ = True
    neighbors g v = nub <$> liftA2 (++) 
             (map (\(Edge _ e) -> e) <$> (outgoing g v)) 
             (map (\(Edge s _) -> s) <$> (ingoing g v))

-- | Reverse edge direction
reverseEdge :: Edge -> Edge 
reverseEdge (Edge va vb) = edge vb va

-- | SimpleGraph is an instance of Graph.
instance Graph UndirectedSG where
    addVertex = _addVertex
    removeVertex = _removeVertex
    vertexValue = _vertexValue
    changeVertexValue = _changeVertexValue
    someVertex = _someVertex
    hasNoVertices = _hasNoVertices
    allVertices = _allVertices
    allVertexValues = _allVertexValues
    allNodes = _allNodes
    isLinkedWithAnEdge = _isLinkedWithAnEdge
    addEdge = _addEdge
    removeEdge e g = _removeEdge (reverseEdge e) (_removeEdge e g)
    edgeVertices = _edgeVertices
    edgeValue g e = case _edgeValue g e of 
                       Nothing -> _edgeValue g (reverseEdge e) 
                       r@(Just _) -> r
    someEdge = _someEdge
    hasNoEdges = _hasNoEdges
    allEdges = _allEdges
    allEdgeValues = _allEdgeValues
    emptyGraph = _emptyGraph
    oriented _ = False
    -- in undirected graphs the edge direction does not count so we need to get both
    -- ends to be sure we don not forget a vertex. In addition to that, an end may be the current vertex which
    -- is not part of the neighbors. So it has to be filtered out. Obviously, a better solution will
    -- have to be designed.
    neighbors g v = filter (/= v) <$> nub <$> liftA2 (++) 
       (map (\(Edge _ e) -> e) <$> (edges g v)) 
       (map (\(Edge s _) -> s) <$> (edges g v))

_emptyGraph = emptySimpleGraph

_hasNoVertices (SP _ vm _) = IM.null vm

_hasNoEdges (SP em _ _) = M.null em

_allVertices (SP _ vm _) = map Vertex . IM.keys $ vm

_allEdges (SP em _ _) = M.keys $ em

_allNodes (SP _ vm _) = map (\(k,(_,v)) -> (Vertex k,v)) . IM.assocs $ vm

_allVertexValues (SP _ vm _) = map snd (IM.elems vm)

_allEdgeValues (SP em _ _) = M.elems em

_isLinkedWithAnEdge :: SimpleGraph n e v -> Vertex -> Vertex -> Bool 
{-# INLINE _isLinkedWithAnEdge #-}
_isLinkedWithAnEdge (SP em _ _) va vb = M.member (edge va vb) em || M.member (edge vb va) em

_someVertex (SP _ vm _) = 
  if IM.null vm 
    then 
      Nothing 
    else 
      Just . Vertex . head . IM.keys $ vm

_someEdge (SP em _ _) = 
  if M.null em 
    then 
      Nothing 
    else 
      Just . head . M.keys $ em

_addVertex vert@(Vertex v) value (SP em vm nm) = SP em (IM.insertWith' noRedundancy v (emptyNeighborhood,value) vm) nm

_removeVertex v@(Vertex vertex) g@(SP _ vm _)  = maybe g removeVertexWithValue $! (IM.lookup vertex vm)
  where
    removeVertexWithValue (n,_) = let g' = foldr _removeEdge g (ingoingNeighbors n)
                                      SP em vm' nm' = foldr _removeEdge g' (outgoingNeighbors n)
                                  in 
                                  SP em (IM.delete vertex vm') nm'
_vertexValue g@(SP _ vm _) (Vertex i) = maybe Nothing (Just . extractValue) $! (IM.lookup i vm)
  where
    extractValue (_,d) = d

_changeVertexValue v@(Vertex vi) newValue g@(SP e vm nm)  = 
  let newVertexMap = do
       (n,_) <- IM.lookup vi vm
       return $ IM.insert vi (n,newValue) vm
  in 
  case newVertexMap of 
    Nothing -> Just g
    Just nvm -> Just $ SP e nvm nm

_removeEdge e@(Edge (Vertex vs) (Vertex ve)) g@(SP em vm nm)  = 
  let r = do 
        _ <- M.lookup e em -- Check e is member of the graph
        (ns,vsdata) <- IM.lookup vs vm
        (ne,vedata) <- IM.lookup ve vm
        return ((vs,(removeNeighborsEdge e ns,vsdata)),(ve,(removeNeighborsEdge e ne,vedata)))
      updateGraph ((vs,vsdata),(ve,vedata)) =
        let vm' = IM.insert ve vedata . IM.insert vs vsdata $ vm
            em' = M.delete e em 
        in 
        SP em' vm' nm
  in 
  maybe g updateGraph r

_edgeVertices (SP em _ _) e@(Edge vs ve) =
     if M.member e em 
      then 
        Just (vs,ve)
      else
        Nothing

_edgeValue :: SimpleGraph n e v -> Edge -> Maybe e 
{-# INLINE _edgeValue #-}
_edgeValue (SP em _ _) e = do
     v <- M.lookup e em
     return v

addEdgeReference :: NeighborhoodStructure local 
                 => Edge 
                 -> IM.IntMap (local, vertexdata) 
                 -> Vertex 
                 -> Vertex 
                 -> IM.IntMap (local, vertexdata)
{-# INLINE addEdgeReference #-}
addEdgeReference newEdge vm (Vertex vsi) (Vertex vei) = id $! IM.adjust addi vei $! (IM.adjust addo vsi vm)
 where
   addi (n,v) = (addIngoingEdge newEdge n,v)  
   addo (n,v) = (addOutgoingEdge newEdge n,v)  

_addEdge :: (NeighborhoodStructure n,Graph (SimpleGraph n)) => Edge -> e -> SimpleGraph n e v -> SimpleGraph n e v 
{-# INLINE _addEdge #-}
_addEdge newEdge@(Edge vs ve) value g@(SP em vm nm)   = 
  if testEdgeExistence g em vs ve 
    then 
      g
    else
      SP (M.insert newEdge value em) (addEdgeReference newEdge vm vs ve) nm
  where
    testEdgeExistence g em va vb = 
      if (oriented g)
        then 
          M.member (Edge va vb) em
        else 
          M.member (Edge va vb) em || M.member (Edge vb va) em 
    

instance UndirectedGraph UndirectedSG where
  edges g@(SP _ vm _) v@(Vertex vi) =
      do 
        (n,_) <- IM.lookup vi vm
        return (ingoingNeighbors n)

instance DirectedGraph DirectedSG where
  ingoing g@(SP _ vm _) v@(Vertex vi) =
      do 
        (n,_) <- IM.lookup vi vm
        return (ingoingNeighbors n)

  outgoing g@(SP _ vm _) v@(Vertex vi) =
      do 
        (n,_) <- IM.lookup vi vm
        return (outgoingNeighbors n) 

{-
 
Following code is used to display a graph in a form adapted to humans.

-}

bracketS :: String -> String 
bracketS [] = []
bracketS s = " [" ++ s ++ "];"

createNodeStyle :: (MonadWriter String m) 
                => (Vertex -> n -> Maybe String)
                -> (Vertex -> n -> Maybe String)
                -> Maybe String 
                -> Vertex 
                -> n 
                -> m ()
createNodeStyle nodeShape nodeColor maybeLabel v n = 
  let apply f = f v n
      label _ _ = case maybeLabel of 
         Nothing -> Nothing 
         Just s -> Just $ "label=\"" ++ s ++ "\""
  in 
  tell $ bracketS . intercalate "," . mapMaybe apply $ [nodeShape,nodeColor, label]
  

createEdgeStyle :: (MonadWriter String m) 
                => (Edge -> e -> Maybe String)
                -> (Edge -> e -> Maybe String)
                -> Edge
                -> e 
                -> m ()
createEdgeStyle edgeShape edgeColor e n = 
  let apply f = f e n
  in 
  tell $ bracketS . intercalate "," . mapMaybe apply $ [edgeShape,edgeColor]



printNode nm (Vertex k,v) = do 
  tell "\n"
  let r = IM.lookup k nm
  when (isJust r) $ do
     tell $ "Node " ++ fromJust r
  tell "\n"
  tell $ show v
  tell "\n"

addVertexToGraphviz nodeShape nodeColor nm (k,(_,v)) = do
  tell $ show k
  let r = IM.lookup k $ nm 
  createNodeStyle nodeShape nodeColor r (Vertex k) v
  tell "\n"

addVertexToUndirectedGraphviz nm (k,(_,v)) = do
  tell $ show k
  tell "\n"

-- | Print the values of the graph vertices
printGraphValues :: (Graph (SimpleGraph n), Show b) => SimpleGraph n e b -> IO () 
printGraphValues g@(SP _ _ nm) = putStrLn . execWriter $ mapM_ (printNode nm) (allNodes g)

displaySimpleGraph :: (Vertex -> n -> Maybe String)
                   -> (Vertex -> n -> Maybe String)
                   -> (Edge -> e -> Maybe String)
                   -> (Edge -> e -> Maybe String)
                   -> SimpleGraph local e n 
                   -> String 
displaySimpleGraph  nodeShape nodeColor edgeShape edgeColor g@(SP em vm nm) = execWriter $ do 
  tell "digraph dot {\n"
  mapM_ (addVertexToGraphviz nodeShape nodeColor  nm) $ IM.toList vm
  tell "\n"
  mapM_ (addEdgeToGraphviz edgeShape edgeColor ) $ M.toList em
  tell "}\n"
   where
     addEdgeToGraphviz es ec (e@(Edge (Vertex vs) (Vertex ve)),l) = do
       tell $ show vs 
       tell " -> "
       tell $ show ve
       createEdgeStyle es ec e l
       tell "\n"  

noNodeStyle _ _ = Nothing 
noEdgeStyle _ _ = Nothing

instance Show (DirectedSG () CPT) where
  show g = displaySimpleGraph noNodeStyle noNodeStyle noEdgeStyle noEdgeStyle g

instance Show (DirectedSG () MAXCPT) where
  show g = displaySimpleGraph noNodeStyle noNodeStyle noEdgeStyle noEdgeStyle g

instance Show (DirectedSG String String) where
  show g = displaySimpleGraph noNodeStyle noNodeStyle noEdgeStyle noEdgeStyle g

instance (Show b, Show e) => Show (UndirectedSG e b)where
  show g@(SP em vm nm) = execWriter $ do
  tell "graph dot {\n"
  mapM_ (addVertexToUndirectedGraphviz nm) $ IM.toList vm
  tell "\n"
  mapM_ (addEdgeToGraphviz) $ M.toList em
  tell "}\n"
   where
     addEdgeToGraphviz (e@(Edge (Vertex vs) (Vertex ve)),l) = do
       tell $ show vs 
       tell " -- "
       tell $ show ve
       tell "\n"


displayFactors :: (NeighborhoodStructure n, Show f, Factor f, Graph (SimpleGraph n)) => SimpleGraph n a f -> String
displayFactors g@(SP _ _ nm) = 
  let nodes = allNodes g
      displayFactor (Vertex i,f) = 
        let s = fromJust . IM.lookup i $ nm
        in
        s ++ "\n" ++ show f
  in
  intercalate "\n" $ map displayFactor nodes

{-

Graph Monad

-}



-- | State used for the construction of the graph in the monad and containing
-- auxiliary informations like vertex name to vertex id and vertex count
type AuxiliaryState = (M.Map String Int, Int)

emptyAuxiliaryState = (M.empty,0)

-- | The state of the graph monad : the graph and auxiliary data
-- useful during the construction
type GMState g e f = (AuxiliaryState,g e f)

-- | Graph monad.
-- The monad used to simplify the description of a new graph
-- g is the graph type. e the edge type. f the node type (generally a 'Factor')
newtype GraphMonad g e f a = GM {runGraphMonad :: State (GMState g e f) a} deriving(Monad, MonadState (GMState g e f))

-- | Get a named vertex from the graph monad
getVertex :: Graph g => String -> GraphMonad g e f (Maybe Vertex)
getVertex a = do
  (namemap,_) <- gets fst
  return $ do
    i <- M.lookup a namemap
    return (Vertex i)

-- | Add a new labeled edge to the graph
newEdge :: Graph g => Vertex -> Vertex -> e -> GraphMonad g e f ()
newEdge va vb e = do
  (aux,g) <- get 
  let g1 = addEdge (edge va vb) e g
  put $! (aux,g1)
  return ()

-- | Add a node in the graph using the graph monad
graphNode :: NamedGraph g => String -> f -> GraphMonad g e f Vertex 
graphNode vertexName initValue = do
  ((namemap,_),_) <- get
  maybe (getNewEmptyVariable (Just vertexName) initValue) returnVertex $! (M.lookup vertexName namemap)
   where
    returnVertex i = return (Vertex i)

-- | Generate a new unique unamed empty variable
getNewEmptyVariable :: NamedGraph g => Maybe String -> f -> GraphMonad g e f Vertex  
getNewEmptyVariable name initValue = do 
  ((namemap,count),g) <- get 
  let vertexName = maybe ("unamed" ++ show count) id name
      g1 = addLabeledVertex vertexName (Vertex count) initValue g
      namemap1 = M.insert vertexName count namemap
  put $! ((namemap1,count+1),g1)
  return (Vertex count)

runGraph :: Graph g => GraphMonad g e f a -> (a,g e f)
runGraph = removeAuxiliaryState . flip runState (emptyAuxiliaryState,emptyGraph) . runGraphMonad 
 where 
  removeAuxiliaryState (r,(_,g)) = (r,g)

evalGraph :: Graph g => GraphMonad g e f a -> a
evalGraph = flip evalState (emptyAuxiliaryState,emptyGraph) . runGraphMonad 

execGraph :: Graph g => GraphMonad g e f a -> g e f
execGraph = snd . flip execState (emptyAuxiliaryState,emptyGraph) . runGraphMonad