module HGraph.Undirected.Solvers.Treedepth
        ( optimalDecomposition
        , treedepthAtMost
        , isDecomposition
        , Decomposition(..)
        )
where

import HGraph.Undirected
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Maybe
import Data.List
import Control.Monad

data Decomposition a = 
  Decomposition
  { forall a. Decomposition a -> Map a a
ancestor :: M.Map a a
  , forall a. Decomposition a -> Map a (Set a)
children :: M.Map a (S.Set a)
  , forall a. Decomposition a -> Int
depth    :: Int
  , forall a. Decomposition a -> [a]
roots    :: [a]
  }
  deriving (Decomposition a -> Decomposition a -> Bool
(Decomposition a -> Decomposition a -> Bool)
-> (Decomposition a -> Decomposition a -> Bool)
-> Eq (Decomposition a)
forall a. Eq a => Decomposition a -> Decomposition a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Decomposition a -> Decomposition a -> Bool
== :: Decomposition a -> Decomposition a -> Bool
$c/= :: forall a. Eq a => Decomposition a -> Decomposition a -> Bool
/= :: Decomposition a -> Decomposition a -> Bool
Eq)



optimalDecomposition :: t a -> Decomposition a
optimalDecomposition t a
g = Maybe (Decomposition a) -> Decomposition a
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Decomposition a) -> Decomposition a)
-> Maybe (Decomposition a) -> Decomposition a
forall a b. (a -> b) -> a -> b
$ (Maybe (Decomposition a)
 -> Maybe (Decomposition a) -> Maybe (Decomposition a))
-> Maybe (Decomposition a)
-> [Maybe (Decomposition a)]
-> Maybe (Decomposition a)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Maybe (Decomposition a)
-> Maybe (Decomposition a) -> Maybe (Decomposition a)
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
mplus Maybe (Decomposition a)
forall a. Maybe a
Nothing ([Maybe (Decomposition a)] -> Maybe (Decomposition a))
-> [Maybe (Decomposition a)] -> Maybe (Decomposition a)
forall a b. (a -> b) -> a -> b
$ (Integer -> Maybe (Decomposition a))
-> [Integer] -> [Maybe (Decomposition a)]
forall a b. (a -> b) -> [a] -> [b]
map (t a -> Integer -> Maybe (Decomposition a)
forall {t :: * -> *} {a} {a}.
(Adjacency t, Num a, Mutable t, Ord a, Eq a) =>
t a -> a -> Maybe (Decomposition a)
treedepthAtMost t a
g) [Integer
1..]

treedepthAtMost :: t a -> a -> Maybe (Decomposition a)
treedepthAtMost t a
_ a
0 = Maybe (Decomposition a)
forall a. Maybe a
Nothing
treedepthAtMost t a
g a
k
  | (Maybe (Decomposition a) -> Bool)
-> [Maybe (Decomposition a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Maybe (Decomposition a) -> Bool
forall a. Maybe a -> Bool
isNothing [Maybe (Decomposition a)]
ts = Maybe (Decomposition a)
forall a. Maybe a
Nothing
  | Bool
otherwise = Decomposition a -> Maybe (Decomposition a)
forall a. a -> Maybe a
Just (Decomposition a -> Maybe (Decomposition a))
-> Decomposition a -> Maybe (Decomposition a)
forall a b. (a -> b) -> a -> b
$ (Decomposition a -> Decomposition a -> Decomposition a)
-> Decomposition a -> [Decomposition a] -> Decomposition a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Decomposition a
t0 Decomposition a
t1 -> 
                  Decomposition{ ancestor :: Map a a
ancestor = Map a a -> Map a a -> Map a a
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union (Decomposition a -> Map a a
forall a. Decomposition a -> Map a a
ancestor Decomposition a
t0) (Decomposition a -> Map a a
forall a. Decomposition a -> Map a a
ancestor Decomposition a
t1)
                               , children :: Map a (Set a)
children = Map a (Set a) -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union (Decomposition a -> Map a (Set a)
forall a. Decomposition a -> Map a (Set a)
children Decomposition a
t0) (Decomposition a -> Map a (Set a)
forall a. Decomposition a -> Map a (Set a)
children Decomposition a
t1)
                               , depth :: Int
depth = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Decomposition a -> Int
forall a. Decomposition a -> Int
depth Decomposition a
t0) (Decomposition a -> Int
forall a. Decomposition a -> Int
depth Decomposition a
t1)
                               , roots :: [a]
roots = (Decomposition a -> [a]
forall a. Decomposition a -> [a]
roots Decomposition a
t1) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (Decomposition a -> [a]
forall a. Decomposition a -> [a]
roots Decomposition a
t0)
                  }) Decomposition a
forall {a}. Decomposition a
emptyDecomposition
                  ([Decomposition a] -> Decomposition a)
-> [Decomposition a] -> Decomposition a
forall a b. (a -> b) -> a -> b
$ (Maybe (Decomposition a) -> Decomposition a)
-> [Maybe (Decomposition a)] -> [Decomposition a]
forall a b. (a -> b) -> [a] -> [b]
map Maybe (Decomposition a) -> Decomposition a
forall a. HasCallStack => Maybe a -> a
fromJust [Maybe (Decomposition a)]
ts
  where
    gs :: [t a]
gs = ([a] -> t a) -> [[a]] -> [t a]
forall a b. (a -> b) -> [a] -> [b]
map (t a -> [a] -> t a
forall a. t a -> [a] -> t a
forall (t :: * -> *) a. Adjacency t => t a -> [a] -> t a
inducedSubgraph t a
g) ([[a]] -> [t a]) -> [[a]] -> [t a]
forall a b. (a -> b) -> a -> b
$ t a -> [[a]]
forall a. Ord a => t a -> [[a]]
forall (t :: * -> *) a. (Adjacency t, Ord a) => t a -> [[a]]
connectedComponents t a
g
    ts :: [Maybe (Decomposition a)]
ts = (t a -> Maybe (Decomposition a))
-> [t a] -> [Maybe (Decomposition a)]
forall a b. (a -> b) -> [a] -> [b]
map (\t a
g -> t a -> a -> Maybe (Decomposition a)
treedepthAtMost' t a
g a
k) [t a]
gs

treedepthAtMost' :: t a -> a -> Maybe (Decomposition a)
treedepthAtMost' t a
g a
0 = Maybe (Decomposition a)
forall a. Maybe a
Nothing
treedepthAtMost' t a
g a
1
  | t a -> Integer
forall b a. Integral b => t a -> b
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numVertices t a
g Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 = Decomposition a -> Maybe (Decomposition a)
forall a. a -> Maybe a
Just (Decomposition a -> Maybe (Decomposition a))
-> Decomposition a -> Maybe (Decomposition a)
forall a b. (a -> b) -> a -> b
$ Decomposition a
forall {a}. Decomposition a
emptyDecomposition { depth = 1, roots = vertices g }
  | Bool
otherwise = Maybe (Decomposition a)
forall a. Maybe a
Nothing
treedepthAtMost' t a
g a
k = (Maybe (Decomposition a)
 -> Maybe (Decomposition a) -> Maybe (Decomposition a))
-> Maybe (Decomposition a)
-> [Maybe (Decomposition a)]
-> Maybe (Decomposition a)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Maybe (Decomposition a)
-> Maybe (Decomposition a) -> Maybe (Decomposition a)
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
mplus Maybe (Decomposition a)
forall a. Maybe a
Nothing ([Maybe (Decomposition a)] -> Maybe (Decomposition a))
-> [Maybe (Decomposition a)] -> Maybe (Decomposition a)
forall a b. (a -> b) -> a -> b
$ (a -> Maybe (Decomposition a)) -> [a] -> [Maybe (Decomposition a)]
forall a b. (a -> b) -> [a] -> [b]
map a -> Maybe (Decomposition a)
guess ([a] -> [Maybe (Decomposition a)])
-> [a] -> [Maybe (Decomposition a)]
forall a b. (a -> b) -> a -> b
$ t a -> [a]
forall a. t a -> [a]
forall (t :: * -> *) a. UndirectedGraph t => t a -> [a]
vertices t a
g
  where
    guess :: a -> Maybe (Decomposition a)
guess a
v = (Decomposition a -> Decomposition a)
-> Maybe (Decomposition a) -> Maybe (Decomposition a)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> Decomposition a -> Decomposition a
forall {a}. Ord a => a -> Decomposition a -> Decomposition a
addRoot a
v) Maybe (Decomposition a)
td
      where
        td :: Maybe (Decomposition a)
td = t a -> a -> Maybe (Decomposition a)
treedepthAtMost (a -> t a -> t a
forall a. a -> t a -> t a
forall (t :: * -> *) a. Mutable t => a -> t a -> t a
removeVertex a
v t a
g) (a
k a -> a -> a
forall a. Num a => a -> a -> a
- a
1)

isDecomposition :: t k -> Decomposition k -> Bool
isDecomposition t k
g Decomposition k
td =
  ((k, k) -> Bool) -> [(k, k)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(k
v,k
u) -> k
v k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` (Map k (Set k)
ancestors Map k (Set k) -> k -> Set k
forall k a. Ord k => Map k a -> k -> a
M.! k
u) Bool -> Bool -> Bool
|| k
u k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` (Map k (Set k)
ancestors Map k (Set k) -> k -> Set k
forall k a. Ord k => Map k a -> k -> a
M.! k
v)) ([(k, k)] -> Bool) -> [(k, k)] -> Bool
forall a b. (a -> b) -> a -> b
$
            t k -> [(k, k)]
forall a. t a -> [(a, a)]
forall (t :: * -> *) a. UndirectedGraph t => t a -> [(a, a)]
edges t k
g
  where
    ancestors :: Map k (Set k)
ancestors = [(k, Set k)] -> Map k (Set k)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [ (k
v, [k] -> Set k
forall a. Ord a => [a] -> Set a
S.fromList ([k] -> Set k) -> [k] -> Set k
forall a b. (a -> b) -> a -> b
$ k -> [k]
ancestry k
v) | k
v <- t k -> [k]
forall a. t a -> [a]
forall (t :: * -> *) a. UndirectedGraph t => t a -> [a]
vertices t k
g]
    ancestry :: k -> [k]
ancestry k
v
      | Maybe k -> Bool
forall a. Maybe a -> Bool
isNothing Maybe k
mu = []
      | Bool
otherwise = k
u k -> [k] -> [k]
forall a. a -> [a] -> [a]
: k -> [k]
ancestry k
u
      where
        mu :: Maybe k
mu = k
v k -> Map k k -> Maybe k
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` (Decomposition k -> Map k k
forall a. Decomposition a -> Map a a
ancestor Decomposition k
td) 
        Just k
u = Maybe k
mu

emptyDecomposition :: Decomposition a
emptyDecomposition = Decomposition { ancestor :: Map a a
ancestor = Map a a
forall k a. Map k a
M.empty, children :: Map a (Set a)
children = Map a (Set a)
forall k a. Map k a
M.empty, roots :: [a]
roots = [], depth :: Int
depth = Int
0 }

addRoot :: a -> Decomposition a -> Decomposition a
addRoot a
r Decomposition a
td = Decomposition{ ancestor :: Map a a
ancestor = Map a a
a' Map a a -> Map a a -> Map a a
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Decomposition a -> Map a a
forall a. Decomposition a -> Map a a
ancestor Decomposition a
td
                            , children :: Map a (Set a)
children = Map a (Set a)
c' Map a (Set a) -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Decomposition a -> Map a (Set a)
forall a. Decomposition a -> Map a (Set a)
children Decomposition a
td
                            , depth :: Int
depth = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Decomposition a -> Int
forall a. Decomposition a -> Int
depth Decomposition a
td
                            , roots :: [a]
roots = [a
r]
                            }
  where
    a' :: Map a a
a' = [(a, a)] -> Map a a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(a, a)] -> Map a a) -> [(a, a)] -> Map a a
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> [(a, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Decomposition a -> [a]
forall a. Decomposition a -> [a]
roots Decomposition a
td) (a -> [a]
forall a. a -> [a]
repeat a
r)
    c' :: Map a (Set a)
c' = a -> Set a -> Map a (Set a)
forall k a. k -> a -> Map k a
M.singleton a
r  ([a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList ([a] -> Set a) -> [a] -> Set a
forall a b. (a -> b) -> a -> b
$ Decomposition a -> [a]
forall a. Decomposition a -> [a]
roots Decomposition a
td)

showTd :: Decomposition a -> [Char]
showTd Decomposition a
td = (a -> [Char]) -> [a] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Char] -> a -> [Char]
showTd' [Char]
"") (Decomposition a -> [a]
forall a. Decomposition a -> [a]
roots Decomposition a
td)
  where
    showTd' :: [Char] -> a -> [Char]
showTd' [Char]
indent a
v = [Char]
indent [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ a -> [Char]
forall a. Show a => a -> [Char]
show a
v [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
rs
      where
        mcs :: Maybe (Set a)
mcs = a -> Map a (Set a) -> Maybe (Set a)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
v (Decomposition a -> Map a (Set a)
forall a. Decomposition a -> Map a (Set a)
children Decomposition a
td)
        Just Set a
cs = Maybe (Set a)
mcs
        rs :: [Char]
rs
          | Maybe (Set a) -> Bool
forall a. Maybe a -> Bool
isNothing Maybe (Set a)
mcs = [Char]
""
          | Bool
otherwise = (a -> [Char]) -> [a] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Char] -> a -> [Char]
showTd' (Char
'-'Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
:[Char]
indent)) (Set a -> [a]
forall a. Set a -> [a]
S.toList Set a
cs)

instance (Ord a, Show a) => Show (Decomposition a) where
  show :: Decomposition a -> [Char]
show = Decomposition a -> [Char]
forall a. (Ord a, Show a) => Decomposition a -> [Char]
showTd