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
  { ancestor :: M.Map a a
  , children :: M.Map a (S.Set a)
  , depth    :: Int
  , roots    :: [a]
  }
  deriving (Eq)



optimalDecomposition g = fromJust $ foldr mplus Nothing $ map (treedepthAtMost g) [1..]

treedepthAtMost _ 0 = Nothing
treedepthAtMost g k
  | any isNothing ts = Nothing
  | otherwise = Just $ foldl' (\t0 t1 -> 
                  Decomposition{ ancestor = M.union (ancestor t0) (ancestor t1)
                               , children = M.union (children t0) (children t1)
                               , depth = max (depth t0) (depth t1)
                               , roots = (roots t1) ++ (roots t0)
                  }) emptyDecomposition
                  $ map fromJust ts
  where
    gs = map (inducedSubgraph g) $ connectedComponents g
    ts = map (\g -> treedepthAtMost' g k) gs

treedepthAtMost' g 0 = Nothing
treedepthAtMost' g 1
  | numVertices g == 1 = Just $ emptyDecomposition { depth = 1, roots = vertices g }
  | otherwise = Nothing
treedepthAtMost' g k = foldr mplus Nothing $ map guess $ vertices g
  where
    guess v = fmap (addRoot v) td
      where
        td = treedepthAtMost (removeVertex g v) (k - 1)

isDecomposition g td =
  all (\(v,u) -> v `S.member` (ancestors M.! u) || u `S.member` (ancestors M.! v)) $
            edges g
  where
    ancestors = M.fromList [ (v, S.fromList $ ancestry v) | v <- vertices g]
    ancestry v
      | isNothing mu = []
      | otherwise = u : ancestry u
      where
        mu = v `M.lookup` (ancestor td) 
        Just u = mu

emptyDecomposition = Decomposition { ancestor = M.empty, children = M.empty, roots = [], depth = 0 }

addRoot r td = Decomposition{ ancestor = a' `M.union` ancestor td
                            , children = c' `M.union` children td
                            , depth = 1 + depth td
                            , roots = [r]
                            }
  where
    a' = M.fromList $ zip (roots td) (repeat r)
    c' = M.singleton r  (S.fromList $ roots td)

showTd td = concatMap (showTd' "") (roots td)
  where
    showTd' indent v = indent ++ show v ++ "\n" ++ rs
      where
        mcs = M.lookup v (children td)
        Just cs = mcs
        rs
          | isNothing mcs = ""
          | otherwise = concatMap (showTd' ('-':indent)) (S.toList cs)

instance (Ord a, Show a) => Show (Decomposition a) where
  show = showTd