{-# LANGUAGE ImplicitParams, Safe #-}

module Data.BellmanFord where

import Data.Maybe
import Control.Monad
import Control.CUtils.DataParallel
import Control.CUtils.Conc
import Data.List.Extras.Argmax
import Data.List
import Data.Array

keys :: [(a, b)] -> [a]
keys = map fst

mapWithKey :: (t -> u -> v) -> Array Int (t, u) -> Array Int (t, v)
mapWithKey f ar = let
		?seq = True
		?pool = BoxedThreadPool NoPool in
	arr_concF(\(_, i)->let (a, b) = ar!i in ((,) $! a) $! f a b) ((), snd(bounds ar)+1)

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

-- | The Bellman-Ford shortest path algorithm.
bellmanFord :: (Eq a) => [((a, a), Double)] -> a -> [(a, (a, Double))]
bellmanFord gr x = elems$foldl' (\weights _ -> mapWithKey (relaxEdge weights gr) weights) weights [1..length gr+1] where
	ks = nub (map fst (keys gr) ++ map snd (keys gr))
	weights = newArray$(x, (x, 0)) : map (\x -> (x, (x, 1/0))) (delete x ks)

retrievePath :: (Eq a) => [(a, (a, Double))] -> a -> a -> [a]
retrievePath mp a a2 | a == a2 = [a]
retrievePath mp a a2 = retrievePath mp a (fst (fromJust (lookup a2 mp))) ++ [a2]

-- | Cycle finding
cycles :: (Eq a) => [((a, a), ())] -> a -> Maybe a
cycles gr x = do
	(y, z) <- lookup x weights
	guard (round z < 0)
	return y where
	gr' = mapWithKey (\_ _ -> -1)$newArray gr
	weights = bellmanFord(elems gr') x