{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE TupleSections #-}

-- | Linear expressions of variables.
module Math.Programming.LinExpr where

import Data.List (foldl', sortOn)

-- | A linear expression.
--
-- Linear expressions contain symbolic variables of type @b@ and
-- numeric coefficients of type @a@. Often @a@ will be 'Double', and
-- @b@ will be whatever variable type your linear program uses.
data LinExpr a b
  = LinExpr ![(a, b)] !a
  deriving (LinExpr a b -> LinExpr a b -> Bool
(LinExpr a b -> LinExpr a b -> Bool)
-> (LinExpr a b -> LinExpr a b -> Bool) -> Eq (LinExpr a b)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall a b. (Eq a, Eq b) => LinExpr a b -> LinExpr a b -> Bool
/= :: LinExpr a b -> LinExpr a b -> Bool
$c/= :: forall a b. (Eq a, Eq b) => LinExpr a b -> LinExpr a b -> Bool
== :: LinExpr a b -> LinExpr a b -> Bool
$c== :: forall a b. (Eq a, Eq b) => LinExpr a b -> LinExpr a b -> Bool
Eq, ReadPrec [LinExpr a b]
ReadPrec (LinExpr a b)
Int -> ReadS (LinExpr a b)
ReadS [LinExpr a b]
(Int -> ReadS (LinExpr a b))
-> ReadS [LinExpr a b]
-> ReadPrec (LinExpr a b)
-> ReadPrec [LinExpr a b]
-> Read (LinExpr a b)
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
forall a b. (Read a, Read b) => ReadPrec [LinExpr a b]
forall a b. (Read a, Read b) => ReadPrec (LinExpr a b)
forall a b. (Read a, Read b) => Int -> ReadS (LinExpr a b)
forall a b. (Read a, Read b) => ReadS [LinExpr a b]
readListPrec :: ReadPrec [LinExpr a b]
$creadListPrec :: forall a b. (Read a, Read b) => ReadPrec [LinExpr a b]
readPrec :: ReadPrec (LinExpr a b)
$creadPrec :: forall a b. (Read a, Read b) => ReadPrec (LinExpr a b)
readList :: ReadS [LinExpr a b]
$creadList :: forall a b. (Read a, Read b) => ReadS [LinExpr a b]
readsPrec :: Int -> ReadS (LinExpr a b)
$creadsPrec :: forall a b. (Read a, Read b) => Int -> ReadS (LinExpr a b)
Read, Int -> LinExpr a b -> ShowS
[LinExpr a b] -> ShowS
LinExpr a b -> String
(Int -> LinExpr a b -> ShowS)
-> (LinExpr a b -> String)
-> ([LinExpr a b] -> ShowS)
-> Show (LinExpr a b)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. (Show a, Show b) => Int -> LinExpr a b -> ShowS
forall a b. (Show a, Show b) => [LinExpr a b] -> ShowS
forall a b. (Show a, Show b) => LinExpr a b -> String
showList :: [LinExpr a b] -> ShowS
$cshowList :: forall a b. (Show a, Show b) => [LinExpr a b] -> ShowS
show :: LinExpr a b -> String
$cshow :: forall a b. (Show a, Show b) => LinExpr a b -> String
showsPrec :: Int -> LinExpr a b -> ShowS
$cshowsPrec :: forall a b. (Show a, Show b) => Int -> LinExpr a b -> ShowS
Show, (forall a b. (a -> b) -> LinExpr a a -> LinExpr a b)
-> (forall a b. a -> LinExpr a b -> LinExpr a a)
-> Functor (LinExpr a)
forall a b. a -> LinExpr a b -> LinExpr a a
forall a b. (a -> b) -> LinExpr a a -> LinExpr a b
forall a a b. a -> LinExpr a b -> LinExpr a a
forall a a b. (a -> b) -> LinExpr a a -> LinExpr a b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> LinExpr a b -> LinExpr a a
$c<$ :: forall a a b. a -> LinExpr a b -> LinExpr a a
fmap :: forall a b. (a -> b) -> LinExpr a a -> LinExpr a b
$cfmap :: forall a a b. (a -> b) -> LinExpr a a -> LinExpr a b
Functor, (forall m. Monoid m => LinExpr a m -> m)
-> (forall m a. Monoid m => (a -> m) -> LinExpr a a -> m)
-> (forall m a. Monoid m => (a -> m) -> LinExpr a a -> m)
-> (forall a b. (a -> b -> b) -> b -> LinExpr a a -> b)
-> (forall a b. (a -> b -> b) -> b -> LinExpr a a -> b)
-> (forall b a. (b -> a -> b) -> b -> LinExpr a a -> b)
-> (forall b a. (b -> a -> b) -> b -> LinExpr a a -> b)
-> (forall a. (a -> a -> a) -> LinExpr a a -> a)
-> (forall a. (a -> a -> a) -> LinExpr a a -> a)
-> (forall a. LinExpr a a -> [a])
-> (forall a. LinExpr a a -> Bool)
-> (forall a. LinExpr a a -> Int)
-> (forall a. Eq a => a -> LinExpr a a -> Bool)
-> (forall a. Ord a => LinExpr a a -> a)
-> (forall a. Ord a => LinExpr a a -> a)
-> (forall a. Num a => LinExpr a a -> a)
-> (forall a. Num a => LinExpr a a -> a)
-> Foldable (LinExpr a)
forall a. Eq a => a -> LinExpr a a -> Bool
forall a. Num a => LinExpr a a -> a
forall a. Ord a => LinExpr a a -> a
forall m. Monoid m => LinExpr a m -> m
forall a. LinExpr a a -> Bool
forall a. LinExpr a a -> Int
forall a. LinExpr a a -> [a]
forall a. (a -> a -> a) -> LinExpr a a -> a
forall a a. Eq a => a -> LinExpr a a -> Bool
forall a a. Num a => LinExpr a a -> a
forall a a. Ord a => LinExpr a a -> a
forall m a. Monoid m => (a -> m) -> LinExpr a a -> m
forall a m. Monoid m => LinExpr a m -> m
forall a a. LinExpr a a -> Bool
forall a a. LinExpr a a -> Int
forall a a. LinExpr a a -> [a]
forall b a. (b -> a -> b) -> b -> LinExpr a a -> b
forall a b. (a -> b -> b) -> b -> LinExpr a a -> b
forall a a. (a -> a -> a) -> LinExpr a a -> a
forall a m a. Monoid m => (a -> m) -> LinExpr a a -> m
forall a b a. (b -> a -> b) -> b -> LinExpr a a -> b
forall a a b. (a -> b -> b) -> b -> LinExpr a a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => LinExpr a a -> a
$cproduct :: forall a a. Num a => LinExpr a a -> a
sum :: forall a. Num a => LinExpr a a -> a
$csum :: forall a a. Num a => LinExpr a a -> a
minimum :: forall a. Ord a => LinExpr a a -> a
$cminimum :: forall a a. Ord a => LinExpr a a -> a
maximum :: forall a. Ord a => LinExpr a a -> a
$cmaximum :: forall a a. Ord a => LinExpr a a -> a
elem :: forall a. Eq a => a -> LinExpr a a -> Bool
$celem :: forall a a. Eq a => a -> LinExpr a a -> Bool
length :: forall a. LinExpr a a -> Int
$clength :: forall a a. LinExpr a a -> Int
null :: forall a. LinExpr a a -> Bool
$cnull :: forall a a. LinExpr a a -> Bool
toList :: forall a. LinExpr a a -> [a]
$ctoList :: forall a a. LinExpr a a -> [a]
foldl1 :: forall a. (a -> a -> a) -> LinExpr a a -> a
$cfoldl1 :: forall a a. (a -> a -> a) -> LinExpr a a -> a
foldr1 :: forall a. (a -> a -> a) -> LinExpr a a -> a
$cfoldr1 :: forall a a. (a -> a -> a) -> LinExpr a a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> LinExpr a a -> b
$cfoldl' :: forall a b a. (b -> a -> b) -> b -> LinExpr a a -> b
foldl :: forall b a. (b -> a -> b) -> b -> LinExpr a a -> b
$cfoldl :: forall a b a. (b -> a -> b) -> b -> LinExpr a a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> LinExpr a a -> b
$cfoldr' :: forall a a b. (a -> b -> b) -> b -> LinExpr a a -> b
foldr :: forall a b. (a -> b -> b) -> b -> LinExpr a a -> b
$cfoldr :: forall a a b. (a -> b -> b) -> b -> LinExpr a a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> LinExpr a a -> m
$cfoldMap' :: forall a m a. Monoid m => (a -> m) -> LinExpr a a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> LinExpr a a -> m
$cfoldMap :: forall a m a. Monoid m => (a -> m) -> LinExpr a a -> m
fold :: forall m. Monoid m => LinExpr a m -> m
$cfold :: forall a m. Monoid m => LinExpr a m -> m
Foldable, Functor (LinExpr a)
Foldable (LinExpr a)
Functor (LinExpr a)
-> Foldable (LinExpr a)
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> LinExpr a a -> f (LinExpr a b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    LinExpr a (f a) -> f (LinExpr a a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> LinExpr a a -> m (LinExpr a b))
-> (forall (m :: * -> *) a.
    Monad m =>
    LinExpr a (m a) -> m (LinExpr a a))
-> Traversable (LinExpr a)
forall a. Functor (LinExpr a)
forall a. Foldable (LinExpr a)
forall a (m :: * -> *) a.
Monad m =>
LinExpr a (m a) -> m (LinExpr a a)
forall a (f :: * -> *) a.
Applicative f =>
LinExpr a (f a) -> f (LinExpr a a)
forall a (m :: * -> *) a b.
Monad m =>
(a -> m b) -> LinExpr a a -> m (LinExpr a b)
forall a (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> LinExpr a a -> f (LinExpr a b)
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a.
Monad m =>
LinExpr a (m a) -> m (LinExpr a a)
forall (f :: * -> *) a.
Applicative f =>
LinExpr a (f a) -> f (LinExpr a a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> LinExpr a a -> m (LinExpr a b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> LinExpr a a -> f (LinExpr a b)
sequence :: forall (m :: * -> *) a.
Monad m =>
LinExpr a (m a) -> m (LinExpr a a)
$csequence :: forall a (m :: * -> *) a.
Monad m =>
LinExpr a (m a) -> m (LinExpr a a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> LinExpr a a -> m (LinExpr a b)
$cmapM :: forall a (m :: * -> *) a b.
Monad m =>
(a -> m b) -> LinExpr a a -> m (LinExpr a b)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
LinExpr a (f a) -> f (LinExpr a a)
$csequenceA :: forall a (f :: * -> *) a.
Applicative f =>
LinExpr a (f a) -> f (LinExpr a a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> LinExpr a a -> f (LinExpr a b)
$ctraverse :: forall a (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> LinExpr a a -> f (LinExpr a b)
Traversable)

instance Num a => Semigroup (LinExpr a b) where
  <> :: LinExpr a b -> LinExpr a b -> LinExpr a b
(<>) = LinExpr a b -> LinExpr a b -> LinExpr a b
forall a b. Num a => LinExpr a b -> LinExpr a b -> LinExpr a b
(.+.)

instance Num a => Monoid (LinExpr a b) where
  mempty :: LinExpr a b
mempty = a -> LinExpr a b
forall a b. a -> LinExpr a b
con a
0

-- | Construct a term in a linear expression by multiplying a constant
-- by a variable.
(*.) :: Num a => a -> b -> LinExpr a b
*. :: forall a b. Num a => a -> b -> LinExpr a b
(*.) a
x b
y = [(a, b)] -> a -> LinExpr a b
forall a b. [(a, b)] -> a -> LinExpr a b
LinExpr [(a
x, b
y)] a
0

infixl 7 .*

-- | Construct a term in a linear expression by multiplying a variable
-- by a constant.
(.*) :: Num a => b -> a -> LinExpr a b
.* :: forall a b. Num a => b -> a -> LinExpr a b
(.*) = (a -> b -> LinExpr a b) -> b -> a -> LinExpr a b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> LinExpr a b
forall a b. Num a => a -> b -> LinExpr a b
(*.)

infixl 7 *.

-- | Construct a term in a linear expression by dividing a variable by
-- a constant.
(./) :: Fractional a => b -> a -> LinExpr a b
./ :: forall a b. Fractional a => b -> a -> LinExpr a b
(./) b
x a
y = b
x b -> a -> LinExpr a b
forall a b. Num a => b -> a -> LinExpr a b
.* (a
1 a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
y)

infixl 7 ./

-- | Multiplication of linear expressions by a constant.
scale :: Num a => a -> LinExpr a b -> LinExpr a b
scale :: forall a b. Num a => a -> LinExpr a b -> LinExpr a b
scale a
coef (LinExpr [(a, b)]
terms a
constant) = [(a, b)] -> a -> LinExpr a b
forall a b. [(a, b)] -> a -> LinExpr a b
LinExpr [(a, b)]
terms' a
constant'
  where
    terms' :: [(a, b)]
terms' = [(a
c a -> a -> a
forall a. Num a => a -> a -> a
* a
coef, b
x) | (a
c, b
x) <- [(a, b)]
terms]
    constant' :: a
constant' = a
constant a -> a -> a
forall a. Num a => a -> a -> a
* a
coef

-- | Addition of linear expressions.
(.+.) :: Num a => LinExpr a b -> LinExpr a b -> LinExpr a b
.+. :: forall a b. Num a => LinExpr a b -> LinExpr a b -> LinExpr a b
(.+.) (LinExpr [(a, b)]
terms a
constant) (LinExpr [(a, b)]
terms' a
constant') =
  [(a, b)] -> a -> LinExpr a b
forall a b. [(a, b)] -> a -> LinExpr a b
LinExpr ([(a, b)]
terms [(a, b)] -> [(a, b)] -> [(a, b)]
forall a. Semigroup a => a -> a -> a
<> [(a, b)]
terms') (a
constant a -> a -> a
forall a. Num a => a -> a -> a
+ a
constant')

infixl 6 .+.

-- | The difference of linear expressions.
(.-.) :: Num a => LinExpr a b -> LinExpr a b -> LinExpr a b
.-. :: forall a b. Num a => LinExpr a b -> LinExpr a b -> LinExpr a b
(.-.) LinExpr a b
x LinExpr a b
y = LinExpr a b
x LinExpr a b -> LinExpr a b -> LinExpr a b
forall a b. Num a => LinExpr a b -> LinExpr a b -> LinExpr a b
.+. a -> LinExpr a b -> LinExpr a b
forall a b. Num a => a -> LinExpr a b -> LinExpr a b
scale (-a
1) LinExpr a b
y

infixl 6 .-.

-- | A linear expression with a single variable term.
var :: Num a => b -> LinExpr a b
var :: forall a b. Num a => b -> LinExpr a b
var b
x = [(a, b)] -> a -> LinExpr a b
forall a b. [(a, b)] -> a -> LinExpr a b
LinExpr [(a
1, b
x)] a
0

-- | A linear expression with only a constant term.
con :: a -> LinExpr a b
con :: forall a b. a -> LinExpr a b
con = [(a, b)] -> a -> LinExpr a b
forall a b. [(a, b)] -> a -> LinExpr a b
LinExpr []

-- | The sum of variable terms with coefficients of unity.
vsum :: Num a => [b] -> LinExpr a b
vsum :: forall a b. Num a => [b] -> LinExpr a b
vsum = ([(a, b)] -> a -> LinExpr a b) -> a -> [(a, b)] -> LinExpr a b
forall a b c. (a -> b -> c) -> b -> a -> c
flip [(a, b)] -> a -> LinExpr a b
forall a b. [(a, b)] -> a -> LinExpr a b
LinExpr a
0 ([(a, b)] -> LinExpr a b)
-> ([b] -> [(a, b)]) -> [b] -> LinExpr a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> (a, b)) -> [b] -> [(a, b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a
1,)

-- | The sum of linear expressions.
esum :: Num a => Foldable t => t (LinExpr a b) -> LinExpr a b
esum :: forall a (t :: * -> *) b.
(Num a, Foldable t) =>
t (LinExpr a b) -> LinExpr a b
esum = (LinExpr a b -> LinExpr a b -> LinExpr a b)
-> LinExpr a b -> t (LinExpr a b) -> LinExpr a b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' LinExpr a b -> LinExpr a b -> LinExpr a b
forall a b. Num a => LinExpr a b -> LinExpr a b -> LinExpr a b
(.+.) LinExpr a b
forall a. Monoid a => a
mempty

-- | Reduce an expression to its value.
eval :: Num a => LinExpr a a -> a
eval :: forall a. Num a => LinExpr a a -> a
eval (LinExpr [(a, a)]
terms a
constant) = a
constant a -> a -> a
forall a. Num a => a -> a -> a
+ [a] -> a
sum' (((a, a) -> a) -> [(a, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Num a => a -> a -> a
(*)) [(a, a)]
terms)
  where
    sum' :: [a] -> a
sum' = (a -> a -> a) -> a -> [a] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0

-- | Simplify an expression by grouping like terms.
simplify :: (Num a, Ord b) => LinExpr a b -> LinExpr a b
simplify :: forall a b. (Num a, Ord b) => LinExpr a b -> LinExpr a b
simplify (LinExpr [(a, b)]
terms a
constant) =
  [(a, b)] -> a -> LinExpr a b
forall a b. [(a, b)] -> a -> LinExpr a b
LinExpr (((a, b) -> [(a, b)] -> [(a, b)])
-> [(a, b)] -> [(a, b)] -> [(a, b)]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a, b) -> [(a, b)] -> [(a, b)]
forall {b} {a}. (Eq b, Num a) => (a, b) -> [(a, b)] -> [(a, b)]
f [] (((a, b) -> b) -> [(a, b)] -> [(a, b)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (a, b) -> b
forall a b. (a, b) -> b
snd [(a, b)]
terms)) a
constant
  where
    f :: (a, b) -> [(a, b)] -> [(a, b)]
f (a
c, b
x) [] = [(a
c, b
x)]
    f (a
c, b
x) ((a
c', b
x') : [(a, b)]
xs) =
      if b
x b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
x'
        then (a
c a -> a -> a
forall a. Num a => a -> a -> a
+ a
c', b
x) (a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
: [(a, b)]
xs
        else (a
c, b
x) (a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
: (a
c', b
x') (a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
: [(a, b)]
xs