-

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-- | A simple graph library
module Graph
  ( -- * Creating graphs
    runGraph
  , newGraph
  , node
  , Graph(..)
  , Vertex(..)
  , Node(..)

    -- * Topological sort
  , topologicalSort

    -- * Drawing graphs
  , viewGraph
  , graphviz

    -- * Internal types
  , MGraph(..)
  , Vertex_(..)
  ) where

import Control.Monad
import Control.Monad.ST
import Data.List
import Data.STRef
import Data.Vector (Vector)
import Data.Vector.Mutable (MVector)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import System.IO.Temp
import System.Process

-- * Graph API

-- | a 'Node' is just an integer identifier.
newtype Node = Node
  { nodeId :: Int
  } deriving (Eq, Show, Ord, Num)

-- | An "incomplete" version of 'Vertex' where we do not
--   yet have the output nodes handy, nor the labels. Used
--   when building the graph.
data Vertex_ a = Vertex_
  { vlabel  :: !a
  , vid     :: Node
  , vinputs :: [Node]
  } deriving (Functor, Eq, Show)

-- | A "mutable graph".
data MGraph s a = MGraph
  { graphData :: !(STRef s (MVector s (Vertex_ a)))
  , nextNode  :: {-# UNPACK #-} !(STRef s Node)
  }

-- | A vertex in our 'Graph's. Contains the 'label',
--   the node 'ident'ifier, the 'inputs' nodes and 'outputs'
--   nodes of this vertex.
data Vertex a = Vertex
  { label :: !a
  , ident :: Node
  , inputs :: [(Node, a)]
  , outputs :: [(Node, a)]
  } deriving (Functor, Eq, Show)

-- | A 'Graph' as collection of vertices.
newtype Graph a = Graph
  { graphArray :: Vector (Vertex a)
  } deriving (Functor, Eq, Show)

-- | Initialize a new mutable graph.
newGraph :: ST s (MGraph s a)
newGraph = do
  mv0 <- MV.new 8
  v0 <- newSTRef mv0
  n0 <- newSTRef 0
  return (MGraph v0 n0)

-- | @node g a xs@ adds a vertex with label @a@
--   and inputs @xs@ to the graph @g@, returning
--   the identifier of the newly added node.
node :: MGraph s a -> a -> [Node] -> ST s Node
node mg a xs = do
  mv <- readSTRef (graphData mg)
  Node nodeid <- readSTRef (nextNode mg)
  let arrSize = MV.length mv
      available = arrSize - nodeid
      getTarget
        | available < 1 = do mv' <- MV.grow mv (2 * arrSize)
                             writeSTRef (graphData mg) mv'
                             return mv'
        | otherwise     = return mv
  targetVec <- getTarget
  MV.write targetVec nodeid (Vertex_ a (Node nodeid) xs)
  modifySTRef' (nextNode mg) (+1)
  return (Node nodeid)

buildOutputs :: Vector (Vertex_ a) -> Vector [(Node, a)]
buildOutputs v = runST $ do
  mv <- MV.new (V.length v)
  forM_ v $ \(Vertex_ a n@(Node i) is) -> do
    MV.write mv i []
    forM_ is $ \(Node j) ->
      MV.modify mv (++[(n, a)]) j
  V.freeze mv

buildGraph :: Vector (Vertex_ a) -> Graph a
buildGraph v = Graph (V.map f v)

  where f (Vertex_ a nid@(Node i) is) = Vertex a nid (map lkp is) (outputs V.! i)
        lkp n@(Node i) = (n, vlabel (v V.! i))
        outputs = buildOutputs v

freezeGraph :: MGraph s a -> ST s (Graph a)
freezeGraph mg = do
  numNodes <- (\(Node n) -> n) <$> readSTRef (nextNode mg)
  mv <- MV.take numNodes <$> readSTRef (graphData mg)
  v <- V.freeze mv
  return (buildGraph v)

-- | Turn a computation returning a mutable graph into
--   an immutable 'Graph'.
--
-- @
-- ex = runGraph $ do
--   g <- newGraph
--   a <- node g \"A\" []
--   b <- node g \"B\" [a]
--   c <- node g \"C\" [a, b]
--   return g
-- @
runGraph :: (forall s. ST s (MGraph s a)) -> Graph a
runGraph f = runST (freezeGraph =<< f)

buildEdges :: Vector (Vertex a) -> [(Node, Node)]
buildEdges v = [ (ni, n)
               | Vertex _ n is _ <- V.toList v
               , (ni, _) <- is
               ]

-- | Returns an error if the graph has cycles, a list of
--   all the vertices of the graph sorted in topological order
--   otherwise.
topologicalSort :: Graph a -> [Vertex a]
topologicalSort (Graph v) =
  let (nodes, result) = runST $ do
        nodes0 <- newSTRef startNodes
        edges0 <- newSTRef edgeList
        res0 <- newSTRef []
        go nodes0 edges0 res0
  in
  if not (null nodes)
  then error "topologicalSort: the graph has cycles"
  else result

  where startNodes = V.toList (V.filter (null . inputs) v)
        edgeList = buildEdges v

        go nodes edges res = do
          ns <- readSTRef nodes
          case ns of
            [] -> (,) <$> readSTRef edges <*> readSTRef res
            (vtx:vs) -> do
              modifySTRef' nodes tail
              modifySTRef' res (++[vtx])
              forM_ (outputs $ v V.! nodeId (ident vtx)) $ \(i, a) -> do
                es <- readSTRef edges
                writeSTRef edges (es \\ [(ident vtx, i)])
                let noOtherIncomingEdge = (==1) . length $ filter ((==i) . snd) es
                when noOtherIncomingEdge $
                  modifySTRef' nodes (\xs -> xs ++ [v V.! nodeId i])
              go nodes edges res

-- * Graphviz

-- | Generates the graphviz description of the graph
graphviz :: Graph String -> String
graphviz g = unlines . wrapGraph . foldMap collect $ graphArray g

  where collect (Vertex a _ is os) = ("  \"" ++ a ++ "\";") : map (ppEdgeToFrom a . snd) is
        ppEdgeToFrom lbl1 lbl2 = "  \"" ++ lbl2 ++ "\" -> \"" ++ lbl1 ++ "\";"
        wrapGraph xs = [ "digraph g {" ] ++ xs ++ [ "}", "" ]

invokeDot :: Graph String -> IO FilePath
invokeDot g = do
  fpin <- writeSystemTempFile "graph.dot" (graphviz g)
  fpout <- emptySystemTempFile "graph.png"
  callCommand $ "dot -Tpng " ++ fpin ++ " -o " ++ fpout
  return fpout

viewDot :: FilePath -> IO ()
viewDot fp = callCommand $ "xdg-open " ++ fp

-- | Visualise the graph by calling out to the @dot@ command
--   and then using @xdg-open@ to open an image viewer.
viewGraph :: Graph String -> IO ()
viewGraph = invokeDot >=> viewDot