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

import           Control.Lens                 (_1, 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.IDLGraph (IDLGraph,
                                               PositiveInfiniteInt (Finite),
                                               addPositiveInfiniteInt)

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, IDLGraphVertex) -> Int
forall a. Set a -> Int
Set.size ([(IDLGraphVertex, IDLGraphVertex)]
-> Set (IDLGraphVertex, IDLGraphVertex)
forall a. Ord a => [a] -> Set a
Set.fromList ([((IDLGraphVertex, IDLGraphVertex), Int)]
graph [((IDLGraphVertex, IDLGraphVertex), Int)]
-> Getting
     (Endo [(IDLGraphVertex, IDLGraphVertex)])
     [((IDLGraphVertex, IDLGraphVertex), Int)]
     (IDLGraphVertex, IDLGraphVertex)
-> [(IDLGraphVertex, IDLGraphVertex)]
forall s a. s -> Getting (Endo [a]) s a -> [a]
^.. (((IDLGraphVertex, IDLGraphVertex), Int)
 -> Const
      (Endo [(IDLGraphVertex, IDLGraphVertex)])
      ((IDLGraphVertex, IDLGraphVertex), Int))
-> [((IDLGraphVertex, IDLGraphVertex), Int)]
-> Const
     (Endo [(IDLGraphVertex, 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, IDLGraphVertex), Int))
 -> [((IDLGraphVertex, IDLGraphVertex), Int)]
 -> Const
      (Endo [(IDLGraphVertex, IDLGraphVertex)])
      [((IDLGraphVertex, IDLGraphVertex), Int)])
-> (((IDLGraphVertex, IDLGraphVertex)
     -> Const
          (Endo [(IDLGraphVertex, IDLGraphVertex)])
          (IDLGraphVertex, IDLGraphVertex))
    -> ((IDLGraphVertex, IDLGraphVertex), Int)
    -> Const
         (Endo [(IDLGraphVertex, IDLGraphVertex)])
         ((IDLGraphVertex, IDLGraphVertex), Int))
-> Getting
     (Endo [(IDLGraphVertex, IDLGraphVertex)])
     [((IDLGraphVertex, IDLGraphVertex), Int)]
     (IDLGraphVertex, IDLGraphVertex)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((IDLGraphVertex, IDLGraphVertex)
 -> Const
      (Endo [(IDLGraphVertex, IDLGraphVertex)])
      (IDLGraphVertex, IDLGraphVertex))
-> ((IDLGraphVertex, IDLGraphVertex), Int)
-> Const
     (Endo [(IDLGraphVertex, IDLGraphVertex)])
     ((IDLGraphVertex, IDLGraphVertex), Int)
forall s t a b. Field1 s t a b => Lens s t a b
_1))] ((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
_, PositiveInfiniteInt
df) <- LensLike'
  (Const (IDLGraphVertex, PositiveInfiniteInt))
  BellmanFordStore
  (Maybe (IDLGraphVertex, PositiveInfiniteInt))
-> (Maybe (IDLGraphVertex, PositiveInfiniteInt)
    -> (IDLGraphVertex, PositiveInfiniteInt))
-> BellmanFord (IDLGraphVertex, PositiveInfiniteInt)
forall s (m :: * -> *) r a.
MonadState s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
uses (Index BellmanFordStore
-> Lens' BellmanFordStore (Maybe (IxValue BellmanFordStore))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at IDLGraphVertex
Index BellmanFordStore
f) Maybe (IDLGraphVertex, PositiveInfiniteInt)
-> (IDLGraphVertex, PositiveInfiniteInt)
forall a. HasCallStack => Maybe a -> a
fromJust
      (IDLGraphVertex
pt, PositiveInfiniteInt
dt) <- LensLike'
  (Const (IDLGraphVertex, PositiveInfiniteInt))
  BellmanFordStore
  (Maybe (IDLGraphVertex, PositiveInfiniteInt))
-> (Maybe (IDLGraphVertex, PositiveInfiniteInt)
    -> (IDLGraphVertex, PositiveInfiniteInt))
-> BellmanFord (IDLGraphVertex, PositiveInfiniteInt)
forall s (m :: * -> *) r a.
MonadState s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
uses (Index BellmanFordStore
-> Lens' BellmanFordStore (Maybe (IxValue BellmanFordStore))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at IDLGraphVertex
Index BellmanFordStore
t) Maybe (IDLGraphVertex, PositiveInfiniteInt)
-> (IDLGraphVertex, PositiveInfiniteInt)
forall a. HasCallStack => Maybe a -> a
fromJust
      Bool -> BellmanFord () -> BellmanFord ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
addPositiveInfiniteInt PositiveInfiniteInt
df (Int -> PositiveInfiniteInt
Finite Int
w) PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
forall a. Ord a => a -> a -> Bool
< PositiveInfiniteInt
dt) (BellmanFord () -> BellmanFord ())
-> BellmanFord () -> BellmanFord ()
forall a b. (a -> b) -> a -> b
$ do
        Index BellmanFordStore
-> Lens' BellmanFordStore (Maybe (IxValue BellmanFordStore))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at IDLGraphVertex
Index BellmanFordStore
t ((Maybe (IDLGraphVertex, PositiveInfiniteInt)
  -> Identity (Maybe (IDLGraphVertex, PositiveInfiniteInt)))
 -> BellmanFordStore -> Identity BellmanFordStore)
-> Maybe (IDLGraphVertex, PositiveInfiniteInt) -> BellmanFord ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= (IDLGraphVertex, PositiveInfiniteInt)
-> Maybe (IDLGraphVertex, PositiveInfiniteInt)
forall a. a -> Maybe a
Just (IDLGraphVertex
f, PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
addPositiveInfiniteInt PositiveInfiniteInt
df (Int -> PositiveInfiniteInt
Finite Int
w))
        IDLGraphVertex
-> (IDLGraphVertex, PositiveInfiniteInt) -> BellmanFord ()
propagationFindShorter IDLGraphVertex
t (IDLGraphVertex
pt, PositiveInfiniteInt
dt)
  BellmanFord ()
propagationEnd