module BellmanFord (bellmanFord, BFPath(BFPath)) where
import Data.Graph.Inductive.Graph
import Data.Graph.Inductive.Internal.Queue
import Data.Array.IArray
import Data.Array.ST
import Data.STRef
import Control.Monad.ST
import ArrayQueue
import MonadStuff
import Data.List

-- Path to a node
data Real w => BFPath b w = BFPath {
	pLen  :: w,                    -- Total distance at the end of the path
	pDest :: Node,                 -- Destination of the path
	pFrom :: Maybe (b, BFPath b w) -- Last edge and remaining path (Nothing for source)
} deriving Show

data (Graph gr, Real w) => BFState s gr a b w = BFState {
	bfsGraph :: gr a b,
	bfsEdgeWt :: b -> w,
	bfsMPaths :: STArray s Node (Maybe (BFPath b w)),
	bfsChanged :: ArrayQueue s,
	-- Sentinel in the queue that ends a pass.
	bfsPassEnder :: Node,
	-- The length of the paths we are currently considering.
	-- If this reaches the number of nodes in the graph, we must have a negative cycle.
	bfsPass :: STRef s Int
}

{--
negativeCycleCheck :: (Graph gr, Real w) => BFState s gr a b w ->
	BFPath b w -> ST s ()
negativeCycleCheck state path =
	let getNodes (BFPath _ dst from) = dst : case from of
		Nothing -> []
		Just (_, p1) -> getNodes p1 in
	let nodes = getNodes path in
	if length (nub nodes) < length nodes
	then error ("Negative cycle detected: " ++ show path)
	else nop
--}

offerPath :: (Graph gr, Real w) => BFState s gr a b w ->
	BFPath b w -> ST s ()
offerPath state newPath@(BFPath newLen dest _) = do
	oldMPath <- readArray (bfsMPaths state) dest
	-- Is newPath the first path to dest or shorter than the existing one?
	let adoptPath = case oldMPath of
		Nothing -> True
		Just (BFPath oldLen _ _) -> newLen < oldLen
	if adoptPath
		then do
			-- Save the new path.
			writeArray (bfsMPaths state) dest (Just newPath)
			-- Mark dest as changed.
			aqEnqueue (bfsChanged state) dest
			nop
		else nop -- Don't update anything.

processEdge :: (Graph gr, Real w) => BFState s gr a b w ->
	BFPath b w -> LEdge b -> ST s ()
processEdge state path1@(BFPath len1 _ _) (_, dest, label) = do
	let edgeLen = (bfsEdgeWt state) label
	let newPath = BFPath (len1 + edgeLen) dest (Just (label, path1))
	offerPath state newPath

endPass :: (Graph gr, Real w) => BFState s gr a b w -> ST s ()
endPass state = do
	qEmpty <- aqIsEmpty (bfsChanged state)
	if qEmpty
		then nop -- No nodes to visit on the next pass.  We're done.
		else do
			-- Increment the pass number.
			modifySTRef (bfsPass state) (+ 1)
			p <- readSTRef (bfsPass state)
			if p == noNodes (bfsGraph state)
				then error "BellmanFord: Negative cycle detected!"
				else do
					-- Re-enqueue the pass ender.
					aqEnqueue (bfsChanged state) (bfsPassEnder state)
					search state -- Continue with the next pass.

search :: (Graph gr, Real w) => BFState s gr a b w -> ST s ()
search state = do
	Just node <- aqDequeue (bfsChanged state)
	if node == bfsPassEnder state
		then endPass state
		else do
			mPath1 <- readArray (bfsMPaths state) node
			let Just path1 = mPath1
			sequence $ map (processEdge state path1) $
				out (bfsGraph state) node
			search state -- Keep going.

bellmanFord :: (Graph gr, Real w) =>
	(b -> w) ->                     -- Edge label -> weight
	Node ->                         -- Source node
	gr a b ->                       -- Graph
	Array Node (Maybe (BFPath b w)) -- ! node -> maybe a path
bellmanFord edgeWt source theGraph = runSTArray (do
		mPaths <- newArray (nodeRange theGraph) Nothing
		let (nlo, nhi) = nodeRange theGraph
		let passEnder = nlo - 1
		changed <- newArrayQueue (passEnder, nhi) -- Queue can contain the pass ender.
		pass <- newSTRef 0
		let state = BFState theGraph edgeWt mPaths changed passEnder pass
		-- Pass is 0, and the queue contains a node that was offered a path with 0 edges.
		offerPath state (BFPath 0 source Nothing)
		aqEnqueue (bfsChanged state) (bfsPassEnder state)
		search state
		return (bfsMPaths state)
	)