```module Data.Graph.Libgraph.Cycles
( CycleTree(..)
, getCycles
, CycleNest
, getCycleNest
) where
import Data.List
import Data.Graph.Libgraph.Core
import Data.Graph.Libgraph.DepthFirst
import Data.Graph.Libgraph.UnionFind(UF)
import Data.Graph.Libgraph.Dot
import qualified  Data.Graph.Libgraph.UnionFind as UF
import Data.Array

data CycleTree vertex = CycleTree vertex [CycleTree vertex]
| Reducible vertex [CycleTree vertex]
| Irreducible      [CycleTree vertex]
deriving Show

getCycles :: Ord vertex => CycleNest vertex -> CycleTree vertex
getCycles nest = {-# SCC "getCycles" #-} cycleTree nest (children nest) 1

cycleTree :: CycleNest vertex -> Array Int [Int] -> Int -> CycleTree vertex
cycleTree nest cs x
= case (vertexType nest) ! x of
IrredHead -> Irreducible ((CycleTree v []) : simplify cts)

where cts  = map ct (cs ! x)
ct c = cycleTree nest cs c
v    = getVertex nest x

-- Flattens irreducible loops.
simplify :: [CycleTree vertex] -> [CycleTree vertex]
simplify vs = foldl simplify' [] vs
where simplify' acc v = case v of
(Irreducible ws) -> ws ++ acc
_                -> v : acc

children :: CycleNest vertex -> Array Int [Int]
children nest
where cs0 = listArray (1,n nest) (cycle [[]])
add cs (p,c) = if p == c then cs else cs // [(c,p : cs ! c)]

-- | Entry vertices of reducible cycles.
getRedHeaders :: CycleNest vertex -> [vertex]
= map i2v (filter (isRedHead . snd) vtyps)
where vtyps = assocs . vertexType \$ nest
i2v   = (getVertex nest) . fst

-- Implementation of Havlaks algorithm.

type S vertex a = State (CycleNest vertex) a

deriving Show

data CycleNest vertex = CycleNest
{ graph        :: Graph Int ()
, getVertex    :: Int -> vertex
, n            :: Int
, dfs          :: Dfs Int ()
, backPreds    :: [[Int]]
, nonBackPreds :: Array Int [Int]
, vertexType   :: Array Int VertexType
, header       :: Array Int Int
, body         :: [Int]              -- P in Havlak's algorithm
, worklist     :: [Int]
, uf           :: UF
}

-- Part a and b of Havlak's algorithm
state0 :: (Show vertex, Ord vertex) => Graph vertex arc -> CycleNest vertex
state0 g = s0
where ps   = map (\w -> partition (isAncestor' (dfs s0) w) (preds (graph s0) w)) [1..n s0]
bps  = map fst ps
nbps = map snd ps
dfsg = dfsGraph g
s0   = CycleNest
{ graph        = fst dfsg
, getVertex    = snd dfsg
, n            = length (vertices g)
, dfs          = getDfs (graph s0)
, backPreds    = bps
, nonBackPreds = listArray (1,n s0) nbps
, vertexType   = listArray (1,n s0) \$ cycle [NonHead]
, header       = listArray (1,n s0) \$ cycle [root . graph \$ s0]
, body         = []
, worklist     = []
, uf           = UF.fromList [1..n s0]
}

isAncestor' d v w = case isAncestor d v w of
Just b  -> b
Nothing -> error "Havlak's algorithm only works for connected graphs!"

getCycleNest :: (Show vertex, Ord vertex) => Graph vertex arc -> CycleNest vertex
getCycleNest g = execState (analyse . reverse \$ [1..n s0]) s0
where s0 = state0 g

-- Part c of Havlak's algorithm
analyse :: Eq vertex => [Int] -> S vertex ()
analyse ws = mapM_ analyse' ws
where analyse' w = do modify \$ \s -> s { body = [] }
analyseBackPreds w
modify \$ \s -> s { worklist = body s }
labelReducible w
work w
merge w

labelReducible :: Eq vertex => Int -> S vertex ()
labelReducible w = do p <- gets \$ body
case p of [] -> return ()
work :: Int -> S vertex ()
work w = do
wl <- gets worklist
case wl of
[]      -> return ()
(x:wl') -> do modify \$ \s -> s { worklist = wl' }
chase w x
work w

merge :: Int -> S vertex ()
merge w = do
p <- gets body
mapM_ (merge' w) p

merge' w x = do
modify \$ \s -> s { header = (header s) // [(x,w)] }
uf_union x w

-- Part d of Havlak's algorithm
analyseBackPreds :: Eq vertex => Int -> S vertex ()
analyseBackPreds w = do bps <- gets backPreds
mapM_ f (bps !!! w)
where f v = if v /= w then do x <- uf_find v

(!!!) :: [a] -> Int -> a
xs !!! i = xs !! (i-1)

-- Part e of Havlak's algorithm

chase :: Int -> Int -> S vertex ()
chase w x = do
nbps <- gets nonBackPreds
mapM_ (chase' w) (nbps ! x)

chase' :: Int -> Int -> S vertex ()
chase' w y = do
y' <- uf_find y
d  <- gets dfs
p  <- gets body
if not \$ isAncestor' d w y' then do
else if not (y' `elem` p) && y' /= w then do
else
return ()

-- Some helper functions

dfsGraph :: (Show vertex, Ord vertex) => Graph vertex arc -> (Graph Int (), Int -> vertex)
dfsGraph g = (mapGraph v2i g', i2v)
where preorder = getPreorder (getDfs g)
i2v i = lookup' i (zip [1..] preorder)
\$ "Libraph.dfsGraph: unreachable from root " ++ show i
v2i v = lookup' v (zip preorder [1..])
\$ "Libraph.dfsGraph: unreachable from root " ++ show v
g' = unitGraph g

modifyVertexType :: (Int,VertexType) -> S vertex ()
modifyVertexType vtyp = do
modify \$ \s -> s { vertexType = (vertexType s) // [vtyp]}

addToNonBackPredsOf :: Int -> Int -> S vertex ()
modify \$ \s -> s { nonBackPreds = (nonBackPreds s) // [(w,y : (nonBackPreds s) ! w)] }

addToBody :: Int -> S vertex ()
addToBody v = modify \$ \s -> s { body = v : body s }

addWork :: Int -> S vertex ()
addWork v = modify \$ \s -> s { worklist = v : worklist s }

uf_find :: Int -> S vertex Int
uf_find v = do
uf' <- gets uf
let r =  UF.find uf' v
-- MF TODO update uf?
return r

uf_union :: Int -> Int -> S vertex ()
uf_union v w = modify \$ \s -> s { uf = UF.union (uf s) v w }

-- Show

instance Show vertex => Show (CycleNest vertex) where
show cycleNest
=  "diGraph G {\n"
++ "rankdir=BT\n"
++ foldl (\s -> (s++) . showVType i2v)  "" (assocs . vertexType \$ cycleNest)
++ foldl (\s -> (s++) . showHeader)     "" (assocs . header     \$ cycleNest)
++ "}\n"
where i2v = getVertex cycleNest

showVType :: Show vertex => (Int -> vertex) -> (Int,VertexType) -> String
showVType i2v (i,vtyp)
= "v" ++ show i ++ " [label=\"" ++ s (i2v i) ++ " (" ++ show i ++ " | " ++ show vtyp ++ ")\"]\n"
where s = escape . show