module NaiveMinCostFlow (minCostFlow) where
import IMinCostFlow
import BellmanFord
import MonadStuff
import Data.Array.IArray
import Data.Array.ST
import Control.Monad
import Control.Monad.ST
import Data.Graph.Inductive.Graph
import Data.Graph.Inductive.Internal.RootPath
import Data.List

data MCFEdge i f c = MCFEdge {
	edgeIdx   :: i,
	edgeCap   :: f,
	edgeCost  :: c,
	edgeIsRev :: Bool
}
data MCFState s gr a b i f c = MCFState {
	mcfGraph  :: gr a (MCFEdge i f c),
	mcfSource :: Node,
	mcfSink   :: Node,
	mcfFlow   :: STArray s i f
}

edgeCapLeft :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
	MCFEdge i f c -> ST s f
edgeCapLeft state (MCFEdge i cap _ isRev) = do
	fwdFlow <- readArray (mcfFlow state) i
	return (if isRev then fwdFlow else cap - fwdFlow)

edgePush :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
	MCFEdge i f c -> f -> ST s ()
edgePush state (MCFEdge i _ _ isRev) nf = do
	oldFwdFlow <- readArray (mcfFlow state) i
	let newFwdFlow = if isRev then oldFwdFlow - nf else oldFwdFlow + nf
	writeArray (mcfFlow state) i newFwdFlow

pathCapLeft :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
	(MCFEdge i f c, BFPath (MCFEdge i f c) c) -> ST s f
pathCapLeft state (lastEdge, BFPath _ _ mFrom) = do
	lastCL <- edgeCapLeft state lastEdge
	case mFrom of
		Nothing -> return lastCL
		Just from -> do
			fromCL <- pathCapLeft state from
			return (min lastCL fromCL)

augment :: (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
	f -> BFPath (MCFEdge i f c) c -> ST s ()
augment state augAmt (BFPath _ _ mFrom) = case mFrom of
	Nothing -> nop
	Just (lastEdge, path1) -> do
		edgePush state lastEdge augAmt
		augment state augAmt path1

doFlow :: forall s gr a b i f c. (Graph gr, Ix i, Real f, Real c) => MCFState s gr a b i f c ->
	ST s ()
doFlow state = do
	filteredEdges <- filterM (\(_, _, l) -> do
			ecl <- edgeCapLeft state l
			return (ecl /= 0)
		) (labEdges (mcfGraph state))
	let filteredGraph = mkGraph (labNodes (mcfGraph state)) filteredEdges :: gr a (MCFEdge i f c)
	-- Why won't we get a negative cycle?  The original graph from the
	-- proposal matcher is acyclic (so no negative cycle), and if we
	-- created a negative cycle, that would contradict the fact that we
	-- always augment along the lowest-cost path.
	let mAugPath = bellmanFord edgeCost (mcfSource state) filteredGraph
		! (mcfSink state)
	case mAugPath of
		Nothing -> nop -- Done.
		-- source /= sink, so augPasth is not empty.
		Just augPath@(BFPath _ _ (Just from)) -> do
			augAmt <- pathCapLeft state from
			augment state augAmt augPath
			doFlow state

-- We need to put the type parameters in scope for the mkGraph call.
minCostFlow :: forall gr a b i f c. (Graph gr, Ix i, Real f, Real c) => MinCostFlowImpl1 gr a b i f c
minCostFlow idxBounds edgeIdx edgeCap edgeCost theGraph (source, sink) = runSTArray (do
		let ourFlipF isRev l =
			MCFEdge (edgeIdx l) (edgeCap l)
				(if isRev then -(edgeCost l) else edgeCost l)
				isRev
		let graph2 = mkGraph (labNodes theGraph) (concatMap
			(\(n1, n2, l) -> [ -- Capacity of reverse edge is never used.
				(n1, n2, MCFEdge (edgeIdx l) (edgeCap l) (  edgeCost l ) False),
				(n2, n1, MCFEdge (edgeIdx l)  undefined  (-(edgeCost l)) True )
			]) $ labEdges theGraph) :: gr a (MCFEdge i f c)
		-- Initialize only the slots of the flow array corresponding to
		-- existing edges to zero to catch buggy callers that query
		-- nonexistent edges.
		flow <- newArray idxBounds undefined
		sequence $ map (\(_, _, l) -> writeArray flow (edgeIdx l) 0)
			(labEdges theGraph)
		let state = MCFState graph2 source sink flow
		doFlow state
		return (mcfFlow state)
	)