module HGraph.Directed.Connectivity.Flow
       ( maxFlow
       , maxDisjointPaths
       , minCut
       , minCutI
       )
where

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

maxFlow :: (Ord a, Adjacency t, DirectedGraph t) => t a -> a -> a -> M.Map (a, a) Bool
maxFlow d s t = maxFlow' $ foldr (\a -> M.insert a False) M.empty (arcs d)
  where
    maxFlow' flow 
      | null p = flow
      | otherwise = maxFlow' flow'
      where
        p = shortestPathResidual d s t flow
        flow' = foldr (M.adjust not) flow $ zip p (tail p)

shortestPathResidual d s t flow = path (S.singleton s) M.empty
  where
    path active preds
      | t `M.member` preds = reverse $ makePath preds t
      | S.null active = []
      | otherwise = path (S.fromList $ M.keys newPred) (preds `M.union` newPred)
        where
          newPred = M.fromList $ [ (u,v)
                             | v <- S.toList active
                             , u <- outneighbors d v
                             , (not $ flow M.! (v,u)) && (not $ u `M.member` preds)
                             ]
                             ++
                             [ (u,v)
                             | v <- S.toList active
                             , u <- inneighbors d v
                             , flow M.! (u, v) && (not $ u `M.member` preds)
                             ]
    makePath preds v
      | v == s = [v]
      | otherwise = v : makePath preds (preds M.! v)

maxDisjointPaths :: (Mutable t, DirectedGraph t, Adjacency t, Integral a) => t a -> a -> a -> [[a]]
maxDisjointPaths d s t = [s : makePath v | v <- outneighbors d s, (2*v + 1) `M.member` succs]
  where
    d'  = foldr addVertex (empty d) (concat [[2*v, 2*v+1] | v <- vertices d])
    d'' = foldr addArc d' ([(2*v, 2*v + 1) | v <- vertices d] ++ [(2*v+1, 2*u) | (v,u) <- arcs d])
    succs = M.fromList $ M.keys $ M.filter (id) $ maxFlow d'' (2*s+1) (2*t)
    makePath v
      | v == t = [t]
      | otherwise = v : makePath ((succs M.! (2*v + 1)) `div` 2)

minCut :: (Mutable t, DirectedGraph t, Adjacency t, Eq a) => t a -> a -> a -> [a]
minCut d s t = map (iToV M.!) $ minCutI di si ti
  where
    (di, itova) = linearizeVertices d
    iToV = M.fromList itova
    Just si = fmap fst $ find ((==s) . snd) itova
    Just ti = fmap fst $ find ((==t) . snd) itova

minCutI :: (Mutable t, DirectedGraph t, Adjacency t, Integral a) => t a -> a -> a -> [a]
minCutI d s t = [u `div` 2 | v <- S.toList $ reach, u <- outneighbors d'' v, not $ u `S.member` reach]
  where
    d'  = foldr addVertex (empty d) (concat [[2*v, 2*v+1] | v <- vertices d])
    d'' = foldr addArc d' ([(2*v, 2*v + 1) | v <- vertices d] ++ [(2*v+1, 2*u) | (v,u) <- arcs d])
    flow = M.filter (id) $ maxFlow d'' (2*s+1) (2*t)
    reach = bfs (S.singleton (2*s+1)) (S.singleton (2*s+1))
    bfs active reached
      | S.null active = reached
      | otherwise = bfs new (S.union reached new)
        where
          new = S.fromList $ [ u
                           | v <- S.toList active
                           , u <- outneighbors d'' v
                           , (not $ (v,u) `M.member` flow) && (not $ u `S.member` reached)
                           ]
                           ++
                           [ u
                           | v <- S.toList active
                           , u <- inneighbors d'' v
                           , (u,v) `M.member` flow && (not $ u `S.member` reached)
                           ]