module Satyros.BellmanFord.Storage
  ( Storage
  , IDLGraph
  , IDLGraphVertex
  , IDLWeightMap
  , rootIDLGraphVertex
  , initializeStorage
  , storageToValues

  , PositiveInfiniteInt(Finite, PositiveInfinity)
  , addPositiveInfiniteInt
  ) where

import           Data.Map      (Map)
import qualified Data.Map      as Map
import           Data.Maybe    (mapMaybe)
import qualified Data.Set      as Set
import           GHC.Generics  (Generic)
import qualified Satyros.QFIDL as QFIDL

type Storage = IDLWeightMap

type IDLGraph = Map (IDLGraphVertex, IDLGraphVertex) Int
type IDLGraphVertex = Maybe QFIDL.Variable
type IDLWeightMap = Map IDLGraphVertex (IDLGraphVertex, PositiveInfiniteInt)

rootIDLGraphVertex :: IDLGraphVertex
rootIDLGraphVertex :: IDLGraphVertex
rootIDLGraphVertex = IDLGraphVertex
forall a. Maybe a
Nothing

initializeStorage :: [QFIDL.Expressed] -> (IDLGraph, Storage)
initializeStorage :: [Expressed] -> (IDLGraph, Storage)
initializeStorage [Expressed]
es = ([((IDLGraphVertex, IDLGraphVertex), Int)] -> IDLGraph
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([((IDLGraphVertex, IDLGraphVertex), Int)] -> IDLGraph)
-> [((IDLGraphVertex, IDLGraphVertex), Int)] -> IDLGraph
forall a b. (a -> b) -> a -> b
$ [((IDLGraphVertex
rootIDLGraphVertex, IDLGraphVertex
rootIDLGraphVertex), Int
0)] [((IDLGraphVertex, IDLGraphVertex), Int)]
-> [((IDLGraphVertex, IDLGraphVertex), Int)]
-> [((IDLGraphVertex, IDLGraphVertex), Int)]
forall a. Semigroup a => a -> a -> a
<> ((, Int
0) ((IDLGraphVertex, IDLGraphVertex)
 -> ((IDLGraphVertex, IDLGraphVertex), Int))
-> (Variable -> (IDLGraphVertex, IDLGraphVertex))
-> Variable
-> ((IDLGraphVertex, IDLGraphVertex), Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IDLGraphVertex
rootIDLGraphVertex, ) (IDLGraphVertex -> (IDLGraphVertex, IDLGraphVertex))
-> (Variable -> IDLGraphVertex)
-> Variable
-> (IDLGraphVertex, IDLGraphVertex)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Variable -> IDLGraphVertex
forall a. a -> Maybe a
Just (Variable -> ((IDLGraphVertex, IDLGraphVertex), Int))
-> [Variable] -> [((IDLGraphVertex, IDLGraphVertex), Int)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Variable]
vars) [((IDLGraphVertex, IDLGraphVertex), Int)]
-> [((IDLGraphVertex, IDLGraphVertex), Int)]
-> [((IDLGraphVertex, IDLGraphVertex), Int)]
forall a. Semigroup a => a -> a -> a
<> [((IDLGraphVertex, IDLGraphVertex), Int)]
edges, [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
-> Storage
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
 -> Storage)
-> ([(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
    -> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))])
-> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
-> Storage
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((IDLGraphVertex
rootIDLGraphVertex, (IDLGraphVertex
rootIDLGraphVertex, Int -> PositiveInfiniteInt
Finite Int
0)) (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
-> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
-> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
forall a. a -> [a] -> [a]
:) ([(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
 -> Storage)
-> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
-> Storage
forall a b. (a -> b) -> a -> b
$ ((,) (IDLGraphVertex
 -> (IDLGraphVertex, PositiveInfiniteInt)
 -> (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt)))
-> (IDLGraphVertex -> (IDLGraphVertex, PositiveInfiniteInt))
-> IDLGraphVertex
-> (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (, PositiveInfiniteInt
PositiveInfinity)) (IDLGraphVertex
 -> (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt)))
-> (Variable -> IDLGraphVertex)
-> Variable
-> (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Variable -> IDLGraphVertex
forall a. a -> Maybe a
Just (Variable
 -> (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt)))
-> [Variable]
-> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Variable]
vars)
  where
    vars :: [Variable]
vars = Set Variable -> [Variable]
forall a. Set a -> [a]
Set.toList (Set Variable -> [Variable])
-> ([Set Variable] -> Set Variable) -> [Set Variable] -> [Variable]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Set Variable] -> Set Variable
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set Variable] -> [Variable]) -> [Set Variable] -> [Variable]
forall a b. (a -> b) -> a -> b
$ (Expressed -> Set Variable) -> [Expressed] -> [Set Variable]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Expressed -> Set Variable
QFIDL.variablesInExpressed [Expressed]
es

    edges :: [((IDLGraphVertex, IDLGraphVertex), Int)]
edges = (\(QFIDL.LessThanEqualTo Variable
x1 Variable
x2 Int
v) -> ((Variable -> IDLGraphVertex
forall a. a -> Maybe a
Just Variable
x1, Variable -> IDLGraphVertex
forall a. a -> Maybe a
Just Variable
x2), Int
v)) (Expressed -> ((IDLGraphVertex, IDLGraphVertex), Int))
-> [Expressed] -> [((IDLGraphVertex, IDLGraphVertex), Int)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Expressed]
es

storageToValues :: IDLWeightMap -> [Int]
storageToValues :: Storage -> [Int]
storageToValues Storage
m
  | Variable -> IDLGraphVertex
forall a. a -> Maybe a
Just Variable
QFIDL.ZeroVariable IDLGraphVertex -> Storage -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`Map.member` Storage
m = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract (Int -> Int -> Int) -> Int -> Int -> Int
forall a b. (a -> b) -> a -> b
$ [Int] -> Int
forall a. [a] -> a
head [Int]
vs) ([Int] -> [Int]
forall a. [a] -> [a]
tail [Int]
vs)
  | Bool
otherwise = [Int]
vs
  where
    vs :: [Int]
vs = ((IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
 -> Maybe Int)
-> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
-> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
x -> (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
-> IDLGraphVertex
forall a b. (a, b) -> a
fst (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
x IDLGraphVertex -> Maybe Int -> Maybe Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int
forall a. Num a => a -> a
negate (Int -> Int) -> Maybe Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PositiveInfiniteInt -> Maybe Int
toInt ((IDLGraphVertex, PositiveInfiniteInt) -> PositiveInfiniteInt
forall a b. (a, b) -> b
snd ((IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
-> (IDLGraphVertex, PositiveInfiniteInt)
forall a b. (a, b) -> b
snd (IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))
x))) ([(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
 -> [Int])
-> (Storage
    -> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))])
-> Storage
-> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Storage
-> [(IDLGraphVertex, (IDLGraphVertex, PositiveInfiniteInt))]
forall k a. Map k a -> [(k, a)]
Map.toAscList (Storage -> [Int]) -> Storage -> [Int]
forall a b. (a -> b) -> a -> b
$ Storage
m

data PositiveInfiniteInt
  = Finite Int
  | PositiveInfinity
  deriving stock ((forall x. PositiveInfiniteInt -> Rep PositiveInfiniteInt x)
-> (forall x. Rep PositiveInfiniteInt x -> PositiveInfiniteInt)
-> Generic PositiveInfiniteInt
forall x. Rep PositiveInfiniteInt x -> PositiveInfiniteInt
forall x. PositiveInfiniteInt -> Rep PositiveInfiniteInt x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep PositiveInfiniteInt x -> PositiveInfiniteInt
$cfrom :: forall x. PositiveInfiniteInt -> Rep PositiveInfiniteInt x
Generic, PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
(PositiveInfiniteInt -> PositiveInfiniteInt -> Bool)
-> (PositiveInfiniteInt -> PositiveInfiniteInt -> Bool)
-> Eq PositiveInfiniteInt
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
$c/= :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
== :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
$c== :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
Eq, Eq PositiveInfiniteInt
Eq PositiveInfiniteInt
-> (PositiveInfiniteInt -> PositiveInfiniteInt -> Ordering)
-> (PositiveInfiniteInt -> PositiveInfiniteInt -> Bool)
-> (PositiveInfiniteInt -> PositiveInfiniteInt -> Bool)
-> (PositiveInfiniteInt -> PositiveInfiniteInt -> Bool)
-> (PositiveInfiniteInt -> PositiveInfiniteInt -> Bool)
-> (PositiveInfiniteInt
    -> PositiveInfiniteInt -> PositiveInfiniteInt)
-> (PositiveInfiniteInt
    -> PositiveInfiniteInt -> PositiveInfiniteInt)
-> Ord PositiveInfiniteInt
PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
PositiveInfiniteInt -> PositiveInfiniteInt -> Ordering
PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
$cmin :: PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
max :: PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
$cmax :: PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
>= :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
$c>= :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
> :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
$c> :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
<= :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
$c<= :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
< :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
$c< :: PositiveInfiniteInt -> PositiveInfiniteInt -> Bool
compare :: PositiveInfiniteInt -> PositiveInfiniteInt -> Ordering
$ccompare :: PositiveInfiniteInt -> PositiveInfiniteInt -> Ordering
$cp1Ord :: Eq PositiveInfiniteInt
Ord, Int -> PositiveInfiniteInt -> ShowS
[PositiveInfiniteInt] -> ShowS
PositiveInfiniteInt -> String
(Int -> PositiveInfiniteInt -> ShowS)
-> (PositiveInfiniteInt -> String)
-> ([PositiveInfiniteInt] -> ShowS)
-> Show PositiveInfiniteInt
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PositiveInfiniteInt] -> ShowS
$cshowList :: [PositiveInfiniteInt] -> ShowS
show :: PositiveInfiniteInt -> String
$cshow :: PositiveInfiniteInt -> String
showsPrec :: Int -> PositiveInfiniteInt -> ShowS
$cshowsPrec :: Int -> PositiveInfiniteInt -> ShowS
Show)

addPositiveInfiniteInt :: PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
addPositiveInfiniteInt :: PositiveInfiniteInt -> PositiveInfiniteInt -> PositiveInfiniteInt
addPositiveInfiniteInt (Finite Int
n)       (Finite Int
m)       = Int -> PositiveInfiniteInt
Finite (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m)
addPositiveInfiniteInt PositiveInfiniteInt
_                PositiveInfiniteInt
PositiveInfinity = PositiveInfiniteInt
PositiveInfinity
addPositiveInfiniteInt PositiveInfiniteInt
PositiveInfinity PositiveInfiniteInt
_                = PositiveInfiniteInt
PositiveInfinity

toInt :: PositiveInfiniteInt -> Maybe Int
toInt :: PositiveInfiniteInt -> Maybe Int
toInt (Finite Int
n)       = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n
toInt PositiveInfiniteInt
PositiveInfinity = Maybe Int
forall a. Maybe a
Nothing