module BellmanFord {-(spTree)-} where
import Data.Graph.Inductive.Graph
import Data.Graph.Inductive.Internal.Queue
import Data.Graph.Inductive.Internal.RootPath
import Data.Array.Diff

data NodeInfo b = NodeInfo {
	path :: Maybe [LNode b],
	changed :: Bool
}
data Graph gr => BFState gr a b = BFState {
	theGraph :: gr a b,
	nis :: DiffArray Node (NodeInfo b),
	changedQ :: Queue Node
}

nisToLRTree nis = do
	ni <- elems nis
	case path ni of
		Just lnl -> return (LP lnl)
		Nothing -> fail "Node is unreachable"

offerPath :: (Graph gr, Real b) => BFState gr a b -> [LNode b] -> BFState gr a b 
offerPath bfs newPath@((dest, newDist): _) =
	-- Is newPath the first path to dest, or better than the previous one?
	let adoptPath =
		case path (nis bfs ! dest) of
			Nothing -> True
			Just ((_, oldDist) : _) -> newDist < oldDist
		in
	if adoptPath then
		bfs{
			-- Update NodeInfo with the new path.
			nis = nis bfs // [(dest, NodeInfo (Just newPath) True)],
			changedQ = if changed (nis bfs ! dest)
				then changedQ bfs -- Already on the queue; leave as is.
				else queuePut dest (changedQ bfs) -- Add to queue.
			}
		else bfs -- Don't update anything.

processEdge :: (Graph gr, Real b) => [LNode b] -> LEdge b -> BFState gr a b -> BFState gr a b
processEdge srcPath@((_, srcDist) : _) (_, dest, edgeLen) bfs =
	let newPath = (dest, srcDist + edgeLen) : srcPath in
	offerPath bfs newPath

search :: (Graph gr, Real b) => BFState gr a b -> LRTree b
search bfs =
	if queueEmpty (changedQ bfs) then
		-- Finished.
		nisToLRTree (nis bfs)
	else
		-- Process a changed node from the queue.
		let (src, moreQ) = queueGet (changedQ bfs) in
		let srcNI = nis bfs ! src in
		-- Clear src's changed flag.
		let bfs1 = bfs{nis = nis bfs // [(src, srcNI{changed = False})], changedQ = moreQ} in
		let Just srcPath = path srcNI in
		let outEdges = out (theGraph bfs) src in
		let newBFS = foldr (processEdge srcPath) bfs1 outEdges in
		search newBFS

spTree :: (Graph gr, Real b) => Node -> gr a b -> LRTree b
spTree source theGraph = let theNodes = range (nodeRange theGraph) in
	let emptyBFS = BFState theGraph
		(array (nodeRange theGraph) (map (\n -> (n, NodeInfo Nothing False)) theNodes))
		mkQueue in
	-- Start with a zero-length path to the source.
	let initBFS = offerPath emptyBFS [(source, 0)] in
	search initBFS