module UnitMinCostFlow (umcf, flowDiff) where
import BellmanFord
import Data.Graph.Inductive.Graph
import Data.Graph.Inductive.Internal.RootPath
import Data.List

maybeDelete :: Eq a => a -> [a] -> Maybe [a]
maybeDelete _ [] = Nothing
maybeDelete e (h:t) = if e == h
	then return t
	else do t1 <- maybeDelete e t; return (h:t1)

-- If the edge occurs in the graph, return Just the graph with one occurrence
-- deleted; otherwise return Nothing.  (delLEdge deletes all occurrences.)
maybeDelOneLEdge :: (DynGraph gr, Eq b) => LEdge b -> gr a b -> Maybe (gr a b)
maybeDelOneLEdge (src, dest, lbl) theGraph =
	let (mc, moreGraph) = match src theGraph in do
		(p, v, l, s) <- mc
		s2 <- maybeDelete (lbl, dest) s
		return ((p, v, l, s2) & moreGraph)

flipEdge (src, dest, lbl) = (dest, src, -lbl)

flipEdgeIn :: (DynGraph gr, Real b) => LEdge b -> gr a b -> gr a b
flipEdgeIn edge theGraph =
	let Just graph1 = maybeDelOneLEdge edge theGraph in
	insEdge (flipEdge edge) graph1

augment :: (DynGraph gr, Real b) => [LNode b] -> gr a b -> gr a b
augment augPath@((v1, d1) : t1) theGraph =
	case t1 of
	[] -> theGraph
	(v2, d2) : t2 -> augment (tail augPath)
		(flipEdgeIn (v1, v2, d2 - d1) theGraph)

-- Find a min-cost flow from s to t in theGraph.
-- Each edge of the input graph has unit capacity and cost given by its label.
-- Returns: flow value, residual graph.
umcf :: (DynGraph gr, Real b) => Node -> Node -> gr a b -> (b, gr a b)
umcf s t theGraph =
	-- Use Bellman-Ford to find an augmenting path from s to t, if one exists.
	-- NOTE: getLPath reverses it into s-to-t order!
	let LP augPath = getLPath t (spTree s theGraph) in
	if null augPath then
		-- Finished.
		(0, theGraph)
	else
		-- Augment, and continue flowing in the resulting graph.
		let graph2 = augment augPath theGraph in
		let (fval1, resid) = umcf s t graph2 in (fval1 + 1, resid)

-- Diffs an original graph and a residual graph, producing the flow graph.
flowDiff :: (DynGraph gr, Real b) => gr a b -> gr a b -> gr a b
flowDiff ograph rgraph = case labEdges ograph of
	[] -> mkGraph (labNodes ograph) []
	oedge:_ -> let Just ograph2 = maybeDelOneLEdge oedge ograph in
		case maybeDelOneLEdge oedge rgraph of
			Just rgraph2 -> flowDiff ograph2 rgraph2
			Nothing -> let Just rgraph2 = maybeDelOneLEdge (flipEdge oedge) rgraph in
				insEdge oedge (flowDiff ograph2 rgraph2)