module LOAG.Graphs where

import Control.Monad (forM, forM_, when)
import Control.Monad.ST
import Control.Monad.State
import CommonTypes
import Data.STRef
import Data.Maybe (catMaybes, isNothing, fromJust)
import Data.Tuple (swap)
import qualified Data.Array as A
import           Data.Array.IArray (amap)
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Array.MArray (mapArray)
import Data.Array.ST

type Schedule   = (A.Array Vertex (Maybe Int), A.Array Int [Vertex])
type Vertex     = Int
type Cycle      = IS.IntSet
type Vertices   = IS.IntSet
type Edge       = (Vertex, Vertex)
type Edges      = S.Set Edge
-- Maps that are suitable for Graphs (from 1 node to a set of nodes)
type Graph s    = (DirGraphRef s, DirGraphRef s)
-- | Frozen version of a graph
type FrGraph    = (DirGraph, DirGraph)
type DirGraph   = A.Array Vertex Vertices
type DirGraphRef s = STArray s Vertex Vertices

-- |----------------------------------------------------------------------
-- | Functions for changing the state within AOAG
-- |  possibly catching errors from creating cycles

addEDs :: Graph s -> [Edge] -> (ST s) (Maybe (Edge, Cycle))
addEDs :: forall s. Graph s -> [Edge] -> ST s (Maybe (Edge, Cycle))
addEDs Graph s
_ [] = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
addEDs Graph s
edp (Edge
e:[Edge]
es) = do
    Either Cycle [Edge]
res <- Edge
e forall s. Edge -> Graph s -> ST s (Either Cycle [Edge])
`inserT` Graph s
edp
    case Either Cycle [Edge]
res of
        Right [Edge]
_ -> forall s. Graph s -> [Edge] -> ST s (Maybe (Edge, Cycle))
addEDs Graph s
edp [Edge]
es
        Left Cycle
c  -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Edge
e,Cycle
c)
       
-- | Draws an edge from one node to another, by adding the latter to the
--    node set of the first
insErt :: Edge -> Graph s -> (ST s) ()
insErt :: forall s. Edge -> Graph s -> ST s ()
insErt (Vertex
f, Vertex
t) g :: Graph s
g@(DirGraphRef s
ft,DirGraphRef s
tf) = do 
    Cycle
ts <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
ft Vertex
f
    Cycle
fs <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
tf Vertex
t
    forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray DirGraphRef s
ft Vertex
f (Vertex
t Vertex -> Cycle -> Cycle
`IS.insert` Cycle
ts)
    forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray DirGraphRef s
tf Vertex
t (Vertex
f Vertex -> Cycle -> Cycle
`IS.insert` Cycle
fs)

removE :: Edge -> Graph s -> (ST s) ()
removE :: forall s. Edge -> Graph s -> ST s ()
removE e :: Edge
e@(Vertex
f,Vertex
t) g :: Graph s
g@(DirGraphRef s
ft,DirGraphRef s
tf) = do 
    Cycle
ts <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
ft Vertex
f
    Cycle
fs <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
tf Vertex
t
    forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray DirGraphRef s
ft Vertex
f (Vertex
t Vertex -> Cycle -> Cycle
`IS.delete` Cycle
ts)
    forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray DirGraphRef s
tf Vertex
t (Vertex
f Vertex -> Cycle -> Cycle
`IS.delete` Cycle
fs)


-- | Revert an edge in the graph
revErt :: Edge -> Graph s -> (ST s) ()
revErt :: forall s. Edge -> Graph s -> ST s ()
revErt Edge
e Graph s
g = do
    Bool
present <- forall s. Edge -> Graph s -> ST s Bool
member Edge
e Graph s
g
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
present forall a b. (a -> b) -> a -> b
$ forall s. Edge -> Graph s -> ST s ()
removE Edge
e Graph s
g forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. Edge -> Graph s -> ST s ()
insErt (forall a b. (a, b) -> (b, a)
swap Edge
e) Graph s
g

-- | Assuming the given graph is already transitively closed, and
-- |    not cyclic, insert an 
-- |    edge such that the graph maintains transitively closed.
-- |    returns the cycle if this results in a cycle or returns a pair
-- |    (graph, edges) if not. Where graph is the new Graph and 
-- |    edges represent the edges that were required for transitively
-- |    closing the graph.
inserT :: Edge -> Graph s -> (ST s) (Either Cycle [Edge])
inserT :: forall s. Edge -> Graph s -> ST s (Either Cycle [Edge])
inserT e :: Edge
e@(Vertex
f, Vertex
t) g :: Graph s
g@(DirGraphRef s
gft,DirGraphRef s
gtf)
    | Vertex
f forall a. Eq a => a -> a -> Bool
== Vertex
t     = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ Vertex -> Cycle
IS.singleton Vertex
f
    | Bool
otherwise  = do
        Bool
present <- forall s. Edge -> Graph s -> ST s Bool
member Edge
e Graph s
g
        if Bool
present 
         then (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right [])
         else do
          Cycle
rs <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gtf Vertex
f
          Cycle
us <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gft Vertex
t
          Cycle
pointsToF <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gtf Vertex
f
          Cycle
pointsToT <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gtf Vertex
t
          Cycle
tPointsTo <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gft Vertex
t
          let new2t :: Cycle
new2t = Cycle
pointsToF Cycle -> Cycle -> Cycle
IS.\\ Cycle
pointsToT
          -- extras from f connects all new nodes pointing to f with t
          let extraF :: [Edge]
extraF = forall a. (a -> Vertex -> a) -> a -> Cycle -> a
IS.foldl' (\[Edge]
acc Vertex
tf -> (Vertex
tf,Vertex
t) forall a. a -> [a] -> [a]
: [Edge]
acc) [] Cycle
new2t
          -- extras of t connects all nodes that will be pointing to t
          -- in the new graph, with all the nodes t points to in the
          -- current graph
          STRef s [Edge]
all2tPointsTo <- forall a s. a -> ST s (STRef s a)
newSTRef []
          forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Cycle -> [Vertex]
IS.toList Cycle
tPointsTo) forall a b. (a -> b) -> a -> b
$ \Vertex
ft -> do
            [Edge]
current  <- forall s a. STRef s a -> ST s a
readSTRef STRef s [Edge]
all2tPointsTo
            Cycle
existing <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
gtf Vertex
ft
            let new4ft :: [Edge]
new4ft = forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) Vertex
ft) forall a b. (a -> b) -> a -> b
$ Cycle -> [Vertex]
IS.toList forall a b. (a -> b) -> a -> b
$ 
                            -- removing existing here matters a lot
                            (Vertex
f Vertex -> Cycle -> Cycle
`IS.insert` Cycle
pointsToF) Cycle -> Cycle -> Cycle
IS.\\ Cycle
existing
            forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [Edge]
all2tPointsTo forall a b. (a -> b) -> a -> b
$ [Edge]
current forall a. [a] -> [a] -> [a]
++ [Edge]
new4ft
                  
          [Edge]
extraT <- forall s a. STRef s a -> ST s a
readSTRef STRef s [Edge]
all2tPointsTo
        -- the extras consists of extras from f and extras from t
        -- both these extra sets dont contain edges if they are already 
        -- present in the old graph
          let extra :: [Edge]
extra  = [Edge]
extraF forall a. [a] -> [a] -> [a]
++ [Edge]
extraT
          forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall s. Edge -> Graph s -> ST s ()
`insErt` Graph s
g) (Edge
e forall a. a -> [a] -> [a]
: [Edge]
extra) 
        -- the new graph contains a cycle if there is a self-edge
        -- this cycle will contain both f and t
          Bool
cyclic <- forall s. Edge -> Graph s -> ST s Bool
member (Vertex
f,Vertex
f) Graph s
g
          if Bool
cyclic
           then do
            Cycle
cycle <- forall s. STArray s Vertex Cycle -> ST s Cycle
getCycle DirGraphRef s
gft
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left Cycle
cycle
           else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right [Edge]
extra
        
       where
        -- given that there is a cycle,all elements of this cycle are being
        -- pointed at by f. However, not all elements that f points to are 
        -- part of the cycle. Only those that point back to f.
        getCycle :: STArray s Vertex Vertices -> (ST s) Cycle
        getCycle :: forall s. STArray s Vertex Cycle -> ST s Cycle
getCycle STArray s Vertex Cycle
gft = do
            Cycle
ts <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Vertex Cycle
gft Vertex
f
            [Maybe Vertex]
mnodes <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Cycle -> [Vertex]
IS.toList Cycle
ts) forall a b. (a -> b) -> a -> b
$ \Vertex
t' -> do
                Cycle
fs' <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Vertex Cycle
gft Vertex
t'
                if Vertex
f Vertex -> Cycle -> Bool
`IS.member` Cycle
fs'
                 then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just Vertex
t'
                 else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Maybe a
Nothing
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Vertex] -> Cycle
IS.fromList forall a b. (a -> b) -> a -> b
$ forall a. [Maybe a] -> [a]
catMaybes [Maybe Vertex]
mnodes

-- | Check if a certain edge is part of a graph which means that,
-- |  the receiving node must be in the node set of the sending
member :: Edge -> Graph s -> (ST s) Bool
member :: forall s. Edge -> Graph s -> ST s Bool
member (Vertex
f, Vertex
t) (DirGraphRef s
ft, DirGraphRef s
tf) = do
    Cycle
ts <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray DirGraphRef s
ft Vertex
f
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Vertex -> Cycle -> Bool
IS.member Vertex
t Cycle
ts

-- | Check whether an edge is part of a frozen graph
fr_member :: FrGraph -> Edge -> Bool
fr_member :: FrGraph -> Edge -> Bool
fr_member (DirGraph
ft, DirGraph
tf) (Vertex
f, Vertex
t) = Vertex -> Cycle -> Bool
IS.member Vertex
t (DirGraph
ft forall i e. Ix i => Array i e -> i -> e
A.! Vertex
f)

-- | Flatten a graph, meaning that we transform this graph to 
-- |  a set of Edges by combining a sending node with all the
-- |  receiving nodes in its node set
flatten :: Graph s -> (ST s) Edges 
flatten :: forall s. Graph s -> ST s Edges
flatten (DirGraphRef s
gft, DirGraphRef s
_) = do
    [(Vertex, Cycle)]
list <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m [(i, e)]
getAssocs DirGraphRef s
gft
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap 
                (\(Vertex
f, Cycle
ts) -> forall a b. (a -> b) -> [a] -> [b]
map ((,) Vertex
f) forall a b. (a -> b) -> a -> b
$ Cycle -> [Vertex]
IS.toList Cycle
ts) [(Vertex, Cycle)]
list

freeze_graph :: Graph s -> (ST s) FrGraph
freeze_graph :: forall s. Graph s -> ST s FrGraph
freeze_graph (DirGraphRef s
mf, DirGraphRef s
mt) = do
    DirGraph
fr_f <- forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
freeze DirGraphRef s
mf
    DirGraph
fr_t <- forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
freeze DirGraphRef s
mt
    forall (m :: * -> *) a. Monad m => a -> m a
return (DirGraph
fr_f, DirGraph
fr_t)