module Data.BellmanFord where

import qualified Data.Map as M
import Data.Maybe
import Control.Monad
import Data.List.Extras.Argmax

-- | Edge relaxation.
relaxEdge nodeWeights edgeWeights x (k, weight) = argmin snd ((k, weight) : map (\y -> (y, maybe (1/0) (\z -> snd (fromJust (M.lookup y nodeWeights)) + z) (M.lookup (y, x) edgeWeights))) (M.keys nodeWeights))

-- | The Bellman-Ford shortest path algorithm.
bellmanFord :: (Ord a) => M.Map (a, a) Double -> a -> M.Map a (a, Double)
bellmanFord gr x = foldl (\weights _ -> M.mapWithKey (relaxEdge weights gr) weights) weights [1..M.size gr+1] where
	keys = map fst (M.keys gr) ++ map snd (M.keys gr)
	weights = M.insert x (x, 0) $ M.fromList $ map (\x -> (x, (x, 1/0))) keys

retrievePath :: (Ord a) => M.Map a (a, Double) -> a -> a -> [a]
retrievePath mp a a2 | a == a2 = [a]
retrievePath mp a a2 = retrievePath mp a (fst (fromJust (M.lookup a2 mp))) ++ [a2]

-- | Cycle finding
cycles :: (Ord a) => M.Map (a, a) () -> a -> Maybe a
cycles gr x = do
	(y, z) <- M.lookup x weights
	guard (round z < 0)
	return y where
	gr' = fmap (const (-1)) gr
	weights = bellmanFord gr' x