module GraphWrapper (
                  Graph,
                  makeGraph, addNodes, addEdge, addEdgesWithNodes,
                  graphNodes, graphEdges, graphNeighbors,
                  findCycles, tSort, tSortInt,
                  findReachables, findReachablesIO,
                  hasPath,
                  findPath,
                  topSort,
                  toDotFormat
                  ) where

-- ==================================================
-- GraphWrapper
--
-- This package extends GHC's Data.Graph package to be more accessible.
--
-- GHC's package uses a contiguous range of Int as nodes, for fast
-- indexing into arrays.  This util package provides functions for
-- converting back and forth to this representation.
--
-- This package also defines a cycle detection function based on
-- the available graph functions, since one is not available.
--
-- ==================================================

import Data.Array
import Data.Array.IO
import System.IO.Unsafe
import Data.Maybe( mapMaybe)
import qualified Data.List as L
import qualified Data.Graph as G
import qualified Data.Map as M
import qualified Data.Set as S

import PPrint
import Eval
import Util(map_insertMany, fromJustOrErr)

-- ===============
-- Graph data type

-- A graph is:
--  * the GHC graph (a table mapping a key to the keys of neighbor nodes)
--     note that the GHC graph is mutable so addEdge can be efficient
--     we use unsafe hackery to get the immutable graph the GHC libraries expect
--  * mapping of nodes to their key
--  * mapping of keys to node
data Graph nodeT = Graph (IOArray G.Vertex [G.Vertex]) (M.Map nodeT G.Vertex) (Array G.Vertex nodeT) Int

instance (PPrint a) => PPrint (Graph a) where
    pPrint d p g =
        let ns = graphNodes g
            es = graphEdges g
        in  text "Graph" <+> text "{" $+$
            nest 2 (text "nodes" <+> text "=" <+>
                    pPrint d p ns <+> text "," $+$
                    text "edges" <+> text "=" <+>
                    pPrint d p es) $+$
            text "}"

instance (Hyper a) => Hyper (Graph a) where
    hyper g y = hyper2 (graphNodes g) (graphEdges g) y

-- ===============
-- Graph construction

vToInt :: (Ord nodeT) => Graph nodeT -> nodeT -> G.Vertex
vToInt (Graph _ ordmap _ _) x = fromJustOrErr "vToInt" (M.lookup x ordmap)

intToV :: Graph nodeT -> G.Vertex -> nodeT
intToV (Graph _ _ lookuptable _) x = lookuptable!x

makeGraph :: (Ord nodeT) => [nodeT] -> [(nodeT,nodeT)] -> IO (Graph nodeT)
makeGraph ns es = do
    let size = length ns
        pairs = zip ns [0..]
        ordmap = M.fromList pairs
        vToInt' x = fromJustOrErr "makeGraph"  (M.lookup x ordmap)
        int_es = mapPair vToInt' es
        bnds = (0, size - 1)
        graph0 = G.buildG bnds int_es
        lookuptable = listArray bnds ns
    graph <- thaw graph0
    return (Graph graph ordmap lookuptable size)

-- ===============
-- Selection functions

graphNodes :: Graph nodeT -> [nodeT]
graphNodes (Graph _ _ lookuptable _) = elems lookuptable

graphEdges :: Graph nodeT -> [(nodeT,nodeT)]
graphEdges g@(Graph graph _ lookuptable _) =
    let g_immut = unsafePerformIO (freeze graph) in
    mapPair (intToV g) (G.edges g_immut)

graphNeighbors :: (Ord nodeT) => Graph nodeT -> nodeT -> [nodeT]
graphNeighbors g@(Graph graph _ _ _) n =
  let neighborInts = unsafePerformIO (readArray graph (vToInt g n))
  in map (intToV g) neighborInts

mapPair :: (a -> b) -> [(a, a)] -> [(b, b)]
mapPair f ps = [ (f p1, f p2) | (p1, p2) <- ps ]

-- ===============
-- Graph manipulation

-- Note, this is not a fast operation, so use sparingly
addNodes :: Ord nodeT => Graph nodeT -> [nodeT] -> IO (Graph nodeT)
addNodes g@(Graph graph ordmap lookuptable size) ns = do
    let nodesSet = M.keysSet ordmap
    let new_nodes = filter (not . ((flip S.member) nodesSet)) ns
    let new_size = size + (length new_nodes)
    let new_bnds = (0, new_size - 1)
    let new_map = map_insertMany (zip new_nodes [size..]) ordmap
    let new_table = listArray new_bnds (elems lookuptable ++ new_nodes)
    old_elems <- getElems graph
    new_graph <- newListArray new_bnds (old_elems ++ map (const []) new_nodes)
    return (Graph new_graph new_map new_table new_size)

-- Note, this requires that the nodes of all the edges already be in the graph
-- Use addEdgesWithNodes if that is not the case
addEdge :: Ord nodeT => Graph nodeT -> (nodeT, nodeT) -> IO ()
addEdge g@(Graph graph _ _ _) (v1, v2) = do
  let i1 = vToInt g v1
  let i2 = vToInt g v2
  old_list <- readArray graph i1
  writeArray graph i1 (L.insert i2 old_list)

-- this could be more efficient if we coalesced the edges before adding them
addEdgesWithNodes :: Ord nodeT => Graph nodeT -> [(nodeT,nodeT)] -> IO (Graph nodeT)
addEdgesWithNodes g es = do
    let new_nodes = L.nub ((map fst es) ++ (map snd es))
    g' <- addNodes g new_nodes
    mapM_ (addEdge g') es
    return g'

-- ===============
-- Find cycles

-- This function returns sets of nodes which are strongly connected.
-- The set may include more than one loop.
findCycles :: Graph nodeT -> IO [[nodeT]]
findCycles g@(Graph graph ordmap lookuptable _) = do
    g_assocs <- getAssocs graph
    let ns = elems lookuptable
    let (keys, neighbors) = unzip g_assocs
    let graph_as_list = zip3 ns keys neighbors

    let components = G.stronglyConnComp graph_as_list

    let f (G.AcyclicSCC _) = Nothing
        f (G.CyclicSCC vs) = Just vs

    return (mapMaybe f components)

-- ===============
-- Topological sort

tSort :: (Ord nodeT) => [(nodeT, [nodeT])] -> Either [[nodeT]] [nodeT]
tSort g =
   let nodes = map fst g
       pairs = zip nodes ([0..]::[Int])
       ordmap = M.fromList pairs
       nodeToInt x = fromJustOrErr "tSort" (M.lookup x ordmap)
       convEdge (e1, e2s) = (e1, nodeToInt e1, map nodeToInt e2s)
       g' = map convEdge g
       scc = G.stronglyConnComp g'
       acyclic = [ v  | G.AcyclicSCC v <- scc ]
       cycles  = [ vs | G.CyclicSCC vs <- scc ]
   in
       case (cycles) of
         [] -> Right acyclic
         _  -> Left cycles

{-
-- For data types with an Int conversion function, using that
-- conversion preserves any order implicit in that mapping
-- XXX the default ordering appears to be the reverse of the numbering
tSortByInt :: (nodeT -> Int) -> [(nodeT, [nodeT])] -> Either [[nodeT]] [nodeT]
tSortByInt nodeToInt g =
   let convEdge (n1, n2s) = (n1, nodeToInt n1, map nodeToInt n2s)
       g' = map convEdge g
       scc = G.stronglyConnComp g'
       acyclic = [ v  | G.AcyclicSCC v <- scc ]
       cycles  = [ vs | G.CyclicSCC vs <- scc ]
   in
       case (cycles) of
         [] -> Right acyclic
         _  -> Left cycles
-}

-- For Ints, we can sort them directly
-- XXX is it inefficient to define this as "tSortByInt id" ?
tSortInt :: [(Int, [Int])] -> Either [[Int]] [Int]
tSortInt g = case cycles of
             [] -> Right acyclic
             _  -> Left cycles

  where g'  = [(v, v, es) | (v, es) <- g ]
        scc = G.stronglyConnComp g'
        acyclic = [ v  | G.AcyclicSCC v <- scc ]
        cycles  = [ vs | G.CyclicSCC vs <- scc ]


-- ===============
-- Find reachable nodes (with paths)

-- Takes a graph and a list of starting nodes, and returns a list of
-- the nodes reachable from those starting points (paired with the path).
findReachablesIO :: Ord a => Graph a -> [a] -> IO ([[ (a,[a]) ]])
findReachablesIO g@(Graph graph ordmap lookuptable _) vs = do
    let int_vs = map (vToInt g) vs
    let pairToV (a,path) = (intToV g a, map (intToV g) path)
    g_immut <- freeze graph
    return (map (map pairToV . reachable_withPath g_immut) int_vs)

findReachables :: Ord a => Graph a -> [a] -> [[ (a, [a])]]
findReachables g vs = unsafePerformIO $ findReachablesIO g vs

-- ----------

-- This is like "preorderF" from Data.Graph, except we record the path

-- the path is constructed in reverse
preorder_withPath :: [a] -> G.Tree a -> [(a,[a])]
preorder_withPath path (G.Node a ts) = (a,(a:path)) : preorderF_withPath (a:path) ts

preorderF_withPath :: [a] -> G.Forest a -> [(a,[a])]
preorderF_withPath path ts = concat (map (preorder_withPath path) ts)

-- This is like "reachable" from Data.Graph, except with path included

reachable_withPath :: G.Graph -> G.Vertex -> [(G.Vertex,[G.Vertex])]
reachable_withPath g v = preorderF_withPath [] (G.dfs g [v])

-- ----------

-- Check if a path exists, but don't return it
hasPath :: Ord a => Graph a -> a -> a -> Bool
hasPath g@(Graph graph ordmap lookuptable _) in_v out_v =
    let in_int = vToInt g in_v
        out_int = vToInt g out_v
        g_immut = unsafePerformIO (freeze graph)
    in G.path g_immut in_int out_int

-- The path list starts with in_v and ends with out_v
findPath :: Ord a => Graph a -> a -> a -> Maybe [a]
findPath g@(Graph graph ordmap lookuptable _) in_v out_v =
    let
        in_int = vToInt g in_v
        out_int = vToInt g out_v
        g_immut = unsafePerformIO (freeze graph)
    in case (lookup out_int (reachable_withPath g_immut in_int)) of
         Nothing -> Nothing
         Just path -> Just (map (intToV g) (reverse path))

-- ===============
-- Topological sort

topSort :: Graph nodeT -> [nodeT]
topSort g@(Graph graph ordmap lookuptable _) =
     map (intToV g) (G.topSort g_immut)
  where g_immut = unsafePerformIO (freeze graph)

-- ===============
-- Conversion to "dot" (GraphViz) format

toDotFormat :: String -> Graph nodeT -> (nodeT -> String) -> String
toDotFormat name g@(Graph graph _ _ _) toString =
    let g_immut = unsafePerformIO (freeze graph)
        ns = G.vertices g_immut
        es = G.edges g_immut
        label n = "[label=\"" ++ (toString (intToV g n)) ++ "\"]"
    in unlines $ [ "digraph " ++ name ++ " {"
                 , "// nodes"
                 ] ++
                 [ (show n) ++ " " ++ (label n) | n <- ns ] ++
                 [ ""
                 , "//edges"
                 ] ++
                 [ (show u) ++ " -> " ++ (show v) | (u,v) <- es ] ++
                 [ "}" ]

-- ===============