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

import           Data.Functor  (($>))
import           Data.Map      (Map)
import qualified Data.Map      as Map
import           Data.Maybe    (mapMaybe)
import qualified Data.Set      as Set
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, Int)

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

initializeStorage :: [QFIDL.Expressed] -> (IDLGraph, Storage)
initializeStorage :: [Expressed] -> (IDLGraph, Storage)
initializeStorage [Expressed]
es = ((Int -> Int -> Int)
-> [((IDLGraphVertex, IDLGraphVertex), Int)] -> IDLGraph
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
min ([((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, Int))] -> Storage
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(IDLGraphVertex, (IDLGraphVertex, Int))] -> Storage)
-> ([(IDLGraphVertex, (IDLGraphVertex, Int))]
    -> [(IDLGraphVertex, (IDLGraphVertex, Int))])
-> [(IDLGraphVertex, (IDLGraphVertex, Int))]
-> Storage
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((IDLGraphVertex
rootIDLGraphVertex, (IDLGraphVertex
rootIDLGraphVertex, Int
0)) (IDLGraphVertex, (IDLGraphVertex, Int))
-> [(IDLGraphVertex, (IDLGraphVertex, Int))]
-> [(IDLGraphVertex, (IDLGraphVertex, Int))]
forall a. a -> [a] -> [a]
:) ([(IDLGraphVertex, (IDLGraphVertex, Int))] -> Storage)
-> [(IDLGraphVertex, (IDLGraphVertex, Int))] -> Storage
forall a b. (a -> b) -> a -> b
$ ((,) (IDLGraphVertex
 -> (IDLGraphVertex, Int)
 -> (IDLGraphVertex, (IDLGraphVertex, Int)))
-> (IDLGraphVertex -> (IDLGraphVertex, Int))
-> IDLGraphVertex
-> (IDLGraphVertex, (IDLGraphVertex, Int))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (, Int
0)) (IDLGraphVertex -> (IDLGraphVertex, (IDLGraphVertex, Int)))
-> (Variable -> IDLGraphVertex)
-> Variable
-> (IDLGraphVertex, (IDLGraphVertex, Int))
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)
  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, Int)) -> Maybe Int)
-> [(IDLGraphVertex, (IDLGraphVertex, Int))] -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(IDLGraphVertex, (IDLGraphVertex, Int))
x -> (IDLGraphVertex, (IDLGraphVertex, Int)) -> IDLGraphVertex
forall a b. (a, b) -> a
fst (IDLGraphVertex, (IDLGraphVertex, Int))
x IDLGraphVertex -> Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Int -> Int
forall a. Num a => a -> a
negate ((IDLGraphVertex, Int) -> Int
forall a b. (a, b) -> b
snd ((IDLGraphVertex, (IDLGraphVertex, Int)) -> (IDLGraphVertex, Int)
forall a b. (a, b) -> b
snd (IDLGraphVertex, (IDLGraphVertex, Int))
x))) ([(IDLGraphVertex, (IDLGraphVertex, Int))] -> [Int])
-> (Storage -> [(IDLGraphVertex, (IDLGraphVertex, Int))])
-> Storage
-> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Storage -> [(IDLGraphVertex, (IDLGraphVertex, Int))]
forall k a. Map k a -> [(k, a)]
Map.toAscList (Storage -> [Int]) -> Storage -> [Int]
forall a b. (a -> b) -> a -> b
$ Storage
m