module HGraph.Directed.Connectivity.IntegralLinkage
       ( extendLinkage
       , linkage
       , linkageI
       , LinkageInstance(..)
       )
where

import HGraph.Directed.Connectivity.Flow
import HGraph.Utils
import HGraph.Directed
import qualified Data.Map as M
import Data.Maybe

data LinkageInstance a = 
  LinkageInstance
  { liTerminalPairs :: M.Map Int (a,a)
  , liLinkage       :: M.Map a Int
  , liPath          :: M.Map Int [a]
  }

extendLinkage d inst = 
  case extendLinkage' $ M.keys $ liTerminalPairs inst of
    Nothing -> Nothing
    Just [] -> Just inst
    Just ext ->
      let link' = M.union (foldr (\(v,i) -> 
                                   M.insert v i)
                                 M.empty ext)
                          (liLinkage inst)
          st' = M.union (M.fromList $ [ (i, (v, t))
                                    | (v,i) <- ext
                                    , let (s,t) = (liTerminalPairs inst) M.! i
                                    , v `elem` (outneighbors d s)
                                    ] ++
                                    [ (i, (s, v))
                                    | (v,i) <- ext
                                    , let (s,t) = (liTerminalPairs inst) M.! i
                                    , v `elem` (inneighbors d t)
                                    ]
                        )
                        (liTerminalPairs inst)
      in extendLinkage d inst{liTerminalPairs = st', liLinkage = link'}
  where
    extendLinkage' [] = Just []
    extendLinkage' (i:is)
      | s == t  = extendLinkage' is
      | null cut = Nothing
      | not $ null $ drop 1 cut = extendLinkage' is
      | maybe True (i/=) (cv `M.lookup` (liLinkage inst)) = Just [(cv,i)]
      where
        (s,t) = (liTerminalPairs inst) M.! i
        d' = foldr removeVertex d
                   [ v
                   | v <- vertices d
                   , maybe False (i ==) (v `M.lookup` (liLinkage inst))
                   ]
        cut = minCutI d' s t
        cv = head cut

-- | Finds an integral linkaged connecting the given terminal pairs, if one exists.
linkage :: (DirectedGraph t, Adjacency t, Mutable t, Eq a) => t a -> [(a,a)] -> Maybe [((a,a), [a])]
linkage d st = fmap (map convertResult) $ linkageI di sti
  where
    (di, itova) = linearizeVertices d
    sti = [ (si, ti)
          | (s,t) <- st
          , let si = fst $ head $ filter (\(_, v) -> v == s) itova
          , let ti = fst $ head $ filter (\(_, v) -> v == t) itova
          ]
    iToV = M.fromList itova
    convertResult ((v,u), ps) = ((iToV M.! v, iToV M.! u), map (iToV M.!) ps )

-- | Special case of `linkage` where vertices are of type `Int`.
-- | Faster than calling `linkage` if vertices of the digraph are already of type `Int`.
linkageI :: (DirectedGraph t, Adjacency t, Mutable t, Integral a, Ord a, Eq a) => t a -> [(a,a)] -> Maybe [((a,a), [a])]
linkageI d st = linkage' inst0
  where
    sti = zip [0..] st
    terminalPairs0 = M.fromList sti
    inst0 = LinkageInstance
            { liLinkage = M.fromList $ 
                concatMap (\(i, (s,t)) ->  [(s, i), (t, i)]) sti
            , liTerminalPairs = terminalPairs0
            , liPath = M.empty
            }
    -- linkage' :: (Eq a, Ord a, Num a) => LinkageInstance a -> Maybe [((a,a), [a])]
    linkage' inst
      | M.null $ liTerminalPairs inst = Just [ (terminalPairs0 M.! t, reverse ps) | (t, ps) <- M.assocs $ liPath inst]
      | otherwise = 
        let (i, (s,t)) = head $ M.assocs $ liTerminalPairs inst
            tries = do
              v <- filter (\u -> not $ isJust $ M.lookup u $ liLinkage inst) $ inneighbors d t
              let inst' = extendLinkage d $ inst
                          { liTerminalPairs = M.insert i (s,v) (liTerminalPairs inst)
                          , liPath = M.insertWith (++) i [v] $ liPath inst
                          , liLinkage = M.insert v i $ liLinkage inst
                          }
              case fmap linkage' inst' of
                Just (Just r) -> return r
                Nothing -> []
        in mhead tries