module Reanimate.Math.FloydWarshall where

import           Control.Monad (join)
import           Data.List
import           Data.Maybe
import qualified Data.Map      as Map
import           Data.Map (Map)

data Shortest b a = Shortest { distance :: a, path :: [b] }
  deriving Show

instance (Ord a, Eq b) => Semigroup (Shortest b a) where
  a <> b = case distance a `compare` distance b of
    GT -> b
    LT -> a
    EQ -> a { path = path a `union` path b }

floydWarshall :: (Semigroup b, Ord b, Ord p) =>
     [p]
     -> Map (p, p) (Shortest p b)
     -> Map (p, p) (Maybe (Shortest p b))
floydWarshall v dist = foldr innerCycle (Just <$> dist) v
  where
    innerCycle k dist = (newDist <$> v <*> v) `setTo` dist
      where
        newDist i j =
          ((i,j), do a <- join $ Map.lookup (i, k) dist
                     b <- join $ Map.lookup (k, j) dist
                     return $ Shortest (distance a <> distance b) (path a))
        setTo = Map.unionWith (<>) . Map.fromList

buildPaths :: Ord p => Map (p,p) (Shortest p b) -> Map (p,p) (Shortest [p] b)
buildPaths d = Map.mapWithKey (\pair s -> s { path = buildPath pair}) d
  where
    buildPath (i,j)
      | i == j = [[j]]
      | otherwise = do
        k <- path $ fromJust $ Map.lookup (i,j) d
        p <- buildPath (k, j)
        [i : p]

findMinDistances :: (Monoid b, Ord p, Ord b) => [p] -> Map (p,p) b -> Map (p,p) (Shortest [p] b)
findMinDistances v g =
  let weights = Map.mapWithKey (\(_,j) w -> Shortest w [j]) g
      trivial = Map.fromList [ ((i,i), Shortest mempty []) | i <- v ]
      clean d = fromJust <$> Map.filter isJust (d Map.\\ trivial)
  in buildPaths $ clean $ floydWarshall v (weights <> trivial)