{-# LANGUAGE ViewPatterns #-}
module Satyros.BellmanFord.Propagation where

import           Control.Lens                (_1, _2, at, each, uses, (.=), (^..))
import           Control.Monad               (forM_, when)
import qualified Data.Map                    as Map
import           Data.Maybe                  (fromJust)
import qualified Data.Set                    as Set
import           Satyros.BellmanFord.Effect  (BellmanFord, propagationCheck,
                                              propagationEnd,
                                              propagationFindShorter,
                                              propagationNth)
import           Satyros.BellmanFord.Storage (IDLGraph)

propagation :: IDLGraph -> BellmanFord ()
propagation :: IDLGraph -> BellmanFord ()
propagation (IDLGraph -> [((IDLGraphVertex, IDLGraphVertex), Int)]
forall k a. Map k a -> [(k, a)]
Map.toList -> [((IDLGraphVertex, IDLGraphVertex), Int)]
graph) = do
  [Int] -> (Int -> BellmanFord ()) -> BellmanFord ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
1..Set IDLGraphVertex -> Int
forall a. Set a -> Int
Set.size ([IDLGraphVertex] -> Set IDLGraphVertex
forall a. Ord a => [a] -> Set a
Set.fromList ([((IDLGraphVertex, IDLGraphVertex), Int)]
graph [((IDLGraphVertex, IDLGraphVertex), Int)]
-> Getting
     (Endo [IDLGraphVertex])
     [((IDLGraphVertex, IDLGraphVertex), Int)]
     IDLGraphVertex
-> [IDLGraphVertex]
forall s a. s -> Getting (Endo [a]) s a -> [a]
^.. (((IDLGraphVertex, IDLGraphVertex), Int)
 -> Const
      (Endo [IDLGraphVertex]) ((IDLGraphVertex, IDLGraphVertex), Int))
-> [((IDLGraphVertex, IDLGraphVertex), Int)]
-> Const
     (Endo [IDLGraphVertex]) [((IDLGraphVertex, IDLGraphVertex), Int)]
forall s t a b. Each s t a b => Traversal s t a b
each ((((IDLGraphVertex, IDLGraphVertex), Int)
  -> Const
       (Endo [IDLGraphVertex]) ((IDLGraphVertex, IDLGraphVertex), Int))
 -> [((IDLGraphVertex, IDLGraphVertex), Int)]
 -> Const
      (Endo [IDLGraphVertex]) [((IDLGraphVertex, IDLGraphVertex), Int)])
-> ((IDLGraphVertex
     -> Const (Endo [IDLGraphVertex]) IDLGraphVertex)
    -> ((IDLGraphVertex, IDLGraphVertex), Int)
    -> Const
         (Endo [IDLGraphVertex]) ((IDLGraphVertex, IDLGraphVertex), Int))
-> Getting
     (Endo [IDLGraphVertex])
     [((IDLGraphVertex, IDLGraphVertex), Int)]
     IDLGraphVertex
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((IDLGraphVertex, IDLGraphVertex)
 -> Const (Endo [IDLGraphVertex]) (IDLGraphVertex, IDLGraphVertex))
-> ((IDLGraphVertex, IDLGraphVertex), Int)
-> Const
     (Endo [IDLGraphVertex]) ((IDLGraphVertex, IDLGraphVertex), Int)
forall s t a b. Field1 s t a b => Lens s t a b
_1 (((IDLGraphVertex, IDLGraphVertex)
  -> Const (Endo [IDLGraphVertex]) (IDLGraphVertex, IDLGraphVertex))
 -> ((IDLGraphVertex, IDLGraphVertex), Int)
 -> Const
      (Endo [IDLGraphVertex]) ((IDLGraphVertex, IDLGraphVertex), Int))
-> ((IDLGraphVertex
     -> Const (Endo [IDLGraphVertex]) IDLGraphVertex)
    -> (IDLGraphVertex, IDLGraphVertex)
    -> Const (Endo [IDLGraphVertex]) (IDLGraphVertex, IDLGraphVertex))
-> (IDLGraphVertex -> Const (Endo [IDLGraphVertex]) IDLGraphVertex)
-> ((IDLGraphVertex, IDLGraphVertex), Int)
-> Const
     (Endo [IDLGraphVertex]) ((IDLGraphVertex, IDLGraphVertex), Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IDLGraphVertex -> Const (Endo [IDLGraphVertex]) IDLGraphVertex)
-> (IDLGraphVertex, IDLGraphVertex)
-> Const (Endo [IDLGraphVertex]) (IDLGraphVertex, IDLGraphVertex)
forall s t a b. Field2 s t a b => Lens s t a b
_2)) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2] ((Int -> BellmanFord ()) -> BellmanFord ())
-> (Int -> BellmanFord ()) -> BellmanFord ()
forall a b. (a -> b) -> a -> b
$ \Int
n -> do
    Int -> BellmanFord ()
propagationNth Int
n
    [((IDLGraphVertex, IDLGraphVertex), Int)]
-> (((IDLGraphVertex, IDLGraphVertex), Int) -> BellmanFord ())
-> BellmanFord ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [((IDLGraphVertex, IDLGraphVertex), Int)]
graph ((((IDLGraphVertex, IDLGraphVertex), Int) -> BellmanFord ())
 -> BellmanFord ())
-> (((IDLGraphVertex, IDLGraphVertex), Int) -> BellmanFord ())
-> BellmanFord ()
forall a b. (a -> b) -> a -> b
$ \((IDLGraphVertex
f, IDLGraphVertex
t), Int
w) -> do
      (IDLGraphVertex, IDLGraphVertex) -> BellmanFord ()
propagationCheck (IDLGraphVertex
f, IDLGraphVertex
t)
      (IDLGraphVertex
_, Int
df) <- LensLike'
  (Const (IDLGraphVertex, Int)) Storage (Maybe (IDLGraphVertex, Int))
-> (Maybe (IDLGraphVertex, Int) -> (IDLGraphVertex, Int))
-> BellmanFord (IDLGraphVertex, Int)
forall s (m :: * -> *) r a.
MonadState s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
uses (Index Storage -> Lens' Storage (Maybe (IxValue Storage))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at IDLGraphVertex
Index Storage
f) Maybe (IDLGraphVertex, Int) -> (IDLGraphVertex, Int)
forall a. HasCallStack => Maybe a -> a
fromJust
      (IDLGraphVertex
pt, Int
dt) <- LensLike'
  (Const (IDLGraphVertex, Int)) Storage (Maybe (IDLGraphVertex, Int))
-> (Maybe (IDLGraphVertex, Int) -> (IDLGraphVertex, Int))
-> BellmanFord (IDLGraphVertex, Int)
forall s (m :: * -> *) r a.
MonadState s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
uses (Index Storage -> Lens' Storage (Maybe (IxValue Storage))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at IDLGraphVertex
Index Storage
t) Maybe (IDLGraphVertex, Int) -> (IDLGraphVertex, Int)
forall a. HasCallStack => Maybe a -> a
fromJust
      Bool -> BellmanFord () -> BellmanFord ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
df Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
dt) (BellmanFord () -> BellmanFord ())
-> BellmanFord () -> BellmanFord ()
forall a b. (a -> b) -> a -> b
$ do
        Index Storage -> Lens' Storage (Maybe (IxValue Storage))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at IDLGraphVertex
Index Storage
t ((Maybe (IDLGraphVertex, Int)
  -> Identity (Maybe (IDLGraphVertex, Int)))
 -> Storage -> Identity Storage)
-> Maybe (IDLGraphVertex, Int) -> BellmanFord ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= (IDLGraphVertex, Int) -> Maybe (IDLGraphVertex, Int)
forall a. a -> Maybe a
Just (IDLGraphVertex
f, Int
df Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w)
        IDLGraphVertex -> (IDLGraphVertex, Int) -> BellmanFord ()
propagationFindShorter IDLGraphVertex
t (IDLGraphVertex
pt, Int
dt)
  BellmanFord ()
propagationEnd