-- (c) 2010 by Daneel S. Yaitskov
-- | additional functions for the graph structure defined in fgl library
module Data.Graph.InductivePlus (module Data.Graph.Inductive, module Data.Graph.InductivePlus) where

import Data.Maybe
import Data.List (nub,union)
import Control.Monad.State (execState, get, put, modify)
import Data.Graph.Inductive
import Debug.Trace


delUEdge e@(v,w) g = delEdge e $ delEdge (w,v) g

unear n g = nub $ neighbors g n

-- | the function updates a label of v node in the graph g
setVLabel :: (Node, a) -- ^ node number paired with old node's label value
             -> a -- ^ new label value
             -> Gr a b  
             -> Gr a b 
setVLabel (v,_) newL g =
 let ne = unear v g
     s = lsuc g v
     p = lpre g v
     g' = insNode (v,newL) $ delNode v g
     newEdges = [ (fst x,v, snd x) | x <- p ]  ++ [ (v,fst x, snd x) | x <- s ]
 in  insEdges newEdges g'

-- | the function updates a label of v node in the graph g
setVLabel' :: Node -> a -> Gr a b -> Gr a b
setVLabel' v newL g = setVLabel (v,newL) newL g

-- | the function extracts a label of v node from the graph g
getVLabel :: Node -> Gr a b -> a
getVLabel v g = fromJust $ lab g v

setELabel :: (Node, Node, a) -> a -> Gr b a -> Gr b a
setELabel (v,w,_) newL g =
 let g' = delEdge (v,w) g
     num = length . filter (w == ) $ suc g v
  in insEdges (replicate 1 {-num-} (v,w,newL)) g'

-- | the function version of setELabel for an undirected graph
setUELabel e@(v,w,_) newL g =
 setELabel e newL $ setELabel (w,v,newL) newL g

setELabel' (v,w) newL g = setELabel (v,w,newL) newL g
setUELabel' (v,w) newL g = setUELabel (v,w,newL) newL g
  
getELabel :: (Node, Node) -> Gr a b -> b
getELabel (v,w) g = fromJust . lookup w $ lsuc g v
getELabel' (v,w) g = lookup w $ lsuc g v

isEdge :: (Node, Node) -> Gr a b -> Bool
isEdge (v,w) g  = w  `elem` suc g v

-- | the function merges two graphs. E.i. if first graph hasn't got vertex v 
--   but second one has got it then the node is inserted into first graph with 
--   same label. Edges are processed too.
mergeTwoGraphs :: Gr a b -> Gr a b -> Gr a b
mergeTwoGraphs g1 g2 =
 let mergeNode n g = if n `gelem` g1
                        then g
                        else insNode (n, getVLabel n g2) g
     g1' = foldr mergeNode g1 $ nodes g2
     mergeEdge e@(v,w) g = if isEdge e g
                              then g
                              else insEdge (v,w, getELabel e g2) g
  in foldr mergeEdge g1' $ edges g2

-- | it finds and returns the path consiting of edges from first node to second one.
-- If the path doesn't exist then the function returns the empty list.
findPaths :: Node -> Node -> Gr a b -> [ Path ]
findPaths v0 w g = execState (f [] v0) []
 where
  f curP v =
   let nei = suc g v 
       subf nv = if nv == w
                    then modify ( curP : ) 
                    else if nv `elem` curP || nv == v0
                            then return ()
                            else f ( curP ++ [ nv ] ) nv 
    in mapM_ subf nei 


instance  (Eq a, Eq b) => Eq (Gr a b) where
 g1 == g2 =
   let nsg1 = nodes g1
       nsg2 = nodes g2
       nsg12 = zip nsg1 nsg2
       cmpTwoNodes (n1,n2) =
        let neOf1 = lsuc g1 n1            
            neOf2 = lsuc g2 n2
            len2 = length neOf2
            len1 = length neOf1            
         in lab g1 n1 == lab g2 n2 &&
            {- trace ("v " ++ show n2 ++ "Len2 = " ++ show len2) -} len2 ==
            {- trace ("v " ++ show n2 ++ "Len1 = " ++ show len1) -} len1  &&
            length (neOf1 `union` neOf2) == len2
    in nsg1 == nsg2 && all cmpTwoNodes nsg12



getSources g = filter (null . lpre g) $ nodes g

getSinks   g = filter (null . lsuc g) $ nodes g


filterVertexes predicate g =
  filter (uncurry predicate)
   . map (\n -> (n , getVLabel n g))
         $ nodes g

findVertex predicate g =
 let matching = filterVertexes predicate g
   in if null matching
         then error $ "findVertex: got empty list\n" ++ show g
         else head matching

fst3 (a,_,_) = a
snd3 (_, a,_) = a
thd3 (_, _, a) = a