-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Arith.LPUtil
-- Copyright   :  (c) Masahiro Sakai 2012
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  portable
--
-----------------------------------------------------------------------------
module ToySolver.Arith.LPUtil
  ( toStandardForm
  , toStandardForm'
  ) where

import Control.Exception
import Control.Monad
import Control.Monad.State
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.Maybe
import Data.VectorSpace

import qualified Data.Interval as Interval

import ToySolver.Data.OrdRel
import qualified ToySolver.Data.LA as LA
import ToySolver.Data.IntVar
import qualified ToySolver.Arith.BoundsInference as BI

toStandardForm
  :: (LA.Expr Rational, [OrdRel (LA.Expr Rational)])
  -> ( (LA.Expr Rational, [(LA.Expr Rational, Rational)])
     , Model Rational -> Model Rational
     )
toStandardForm :: (Expr Rational, [OrdRel (Expr Rational)])
-> ((Expr Rational, [(Expr Rational, Rational)]),
    Model Rational -> Model Rational)
toStandardForm prob1 :: (Expr Rational, [OrdRel (Expr Rational)])
prob1@(Expr Rational
obj, [OrdRel (Expr Rational)]
cs) = ((Expr Rational, [(Expr Rational, Rational)])
prob2, Model Rational -> Model Rational
mt)
  where
    vs :: IntSet
vs = Expr Rational -> IntSet
forall a. Variables a => a -> IntSet
vars Expr Rational
obj IntSet -> IntSet -> IntSet
`IS.union` [OrdRel (Expr Rational)] -> IntSet
forall a. Variables a => a -> IntSet
vars [OrdRel (Expr Rational)]
cs
    ((Expr Rational, [(Expr Rational, Rational)])
prob2,VarMap (Expr Rational)
s) = (Expr Rational, [OrdRel (Expr Rational)])
-> ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
toStandardForm' (Expr Rational, [OrdRel (Expr Rational)])
prob1
    mt :: Model Rational -> Model Rational
    mt :: Model Rational -> Model Rational
mt Model Rational
m = [(Key, Rational)] -> Model Rational
forall a. [(Key, a)] -> IntMap a
IM.fromAscList ([(Key, Rational)] -> Model Rational)
-> [(Key, Rational)] -> Model Rational
forall a b. (a -> b) -> a -> b
$ do
      Key
v <- IntSet -> [Key]
IS.toAscList IntSet
vs
      case Key -> VarMap (Expr Rational) -> Maybe (Expr Rational)
forall a. Key -> IntMap a -> Maybe a
IM.lookup Key
v VarMap (Expr Rational)
s of
        Just Expr Rational
def -> (Key, Rational) -> [(Key, Rational)]
forall (m :: * -> *) a. Monad m => a -> m a
return (Key
v, Model Rational -> Expr Rational -> Rational
forall m e v. Eval m e v => m -> e -> v
LA.eval Model Rational
m Expr Rational
def)
        Maybe (Expr Rational)
Nothing  -> (Key, Rational) -> [(Key, Rational)]
forall (m :: * -> *) a. Monad m => a -> m a
return (Key
v, Model Rational
m Model Rational -> Key -> Rational
forall a. IntMap a -> Key -> a
IM.! Key
v)

type M = State Var

toStandardForm'
  :: (LA.Expr Rational, [OrdRel (LA.Expr Rational)])
  -> ( (LA.Expr Rational, [(LA.Expr Rational, Rational)])
     , VarMap (LA.Expr Rational)
     )
toStandardForm' :: (Expr Rational, [OrdRel (Expr Rational)])
-> ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
toStandardForm' (Expr Rational
obj, [OrdRel (Expr Rational)]
cs) = ((Expr Rational, [(Expr Rational, Rational)]),
 VarMap (Expr Rational))
m
  where
    vs :: IntSet
vs = Expr Rational -> IntSet
forall a. Variables a => a -> IntSet
vars Expr Rational
obj IntSet -> IntSet -> IntSet
`IS.union` [OrdRel (Expr Rational)] -> IntSet
forall a. Variables a => a -> IntSet
vars [OrdRel (Expr Rational)]
cs
    v1 :: Key
v1 = if IntSet -> Bool
IS.null IntSet
vs then Key
0 else IntSet -> Key
IS.findMax IntSet
vs Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
1
    initialBounds :: IntMap (Interval Rational)
initialBounds = [(Key, Interval Rational)] -> IntMap (Interval Rational)
forall a. [(Key, a)] -> IntMap a
IM.fromList [(Key
v, Interval Rational
forall r. Ord r => Interval r
Interval.whole) | Key
v <- IntSet -> [Key]
IS.toList IntSet
vs]
    bounds :: IntMap (Interval Rational)
bounds = IntMap (Interval Rational)
-> [OrdRel (Expr Rational)]
-> IntSet
-> Key
-> IntMap (Interval Rational)
forall r.
RealFrac r =>
BoundsEnv r -> [Atom r] -> IntSet -> Key -> BoundsEnv r
BI.inferBounds IntMap (Interval Rational)
initialBounds [OrdRel (Expr Rational)]
cs IntSet
IS.empty Key
10

    gensym :: M Var
    gensym :: M Key
gensym = do
      Key
v <- M Key
forall s (m :: * -> *). MonadState s m => m s
get
      Key -> StateT Key Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Key -> StateT Key Identity ()) -> Key -> StateT Key Identity ()
forall a b. (a -> b) -> a -> b
$ Key
vKey -> Key -> Key
forall a. Num a => a -> a -> a
+Key
1
      Key -> M Key
forall (m :: * -> *) a. Monad m => a -> m a
return Key
v

    m :: ((Expr Rational, [(Expr Rational, Rational)]),
 VarMap (Expr Rational))
m = (State
   Key
   ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
 -> Key
 -> ((Expr Rational, [(Expr Rational, Rational)]),
     VarMap (Expr Rational)))
-> Key
-> State
     Key
     ((Expr Rational, [(Expr Rational, Rational)]),
      VarMap (Expr Rational))
-> ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
forall a b c. (a -> b -> c) -> b -> a -> c
flip State
  Key
  ((Expr Rational, [(Expr Rational, Rational)]),
   VarMap (Expr Rational))
-> Key
-> ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
forall s a. State s a -> s -> a
evalState Key
v1 (State
   Key
   ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
 -> ((Expr Rational, [(Expr Rational, Rational)]),
     VarMap (Expr Rational)))
-> State
     Key
     ((Expr Rational, [(Expr Rational, Rational)]),
      VarMap (Expr Rational))
-> ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
forall a b. (a -> b) -> a -> b
$ do
      VarMap (Expr Rational)
s <- ([VarMap (Expr Rational)] -> VarMap (Expr Rational))
-> StateT Key Identity [VarMap (Expr Rational)]
-> StateT Key Identity (VarMap (Expr Rational))
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [VarMap (Expr Rational)] -> VarMap (Expr Rational)
forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions (StateT Key Identity [VarMap (Expr Rational)]
 -> StateT Key Identity (VarMap (Expr Rational)))
-> StateT Key Identity [VarMap (Expr Rational)]
-> StateT Key Identity (VarMap (Expr Rational))
forall a b. (a -> b) -> a -> b
$ [(Key, Interval Rational)]
-> ((Key, Interval Rational)
    -> StateT Key Identity (VarMap (Expr Rational)))
-> StateT Key Identity [VarMap (Expr Rational)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (IntMap (Interval Rational) -> [(Key, Interval Rational)]
forall a. IntMap a -> [(Key, a)]
IM.toList IntMap (Interval Rational)
bounds) (((Key, Interval Rational)
  -> StateT Key Identity (VarMap (Expr Rational)))
 -> StateT Key Identity [VarMap (Expr Rational)])
-> ((Key, Interval Rational)
    -> StateT Key Identity (VarMap (Expr Rational)))
-> StateT Key Identity [VarMap (Expr Rational)]
forall a b. (a -> b) -> a -> b
$ \(Key
v,Interval Rational
i) -> do
        case Interval Rational -> Extended Rational
forall r. Interval r -> Extended r
Interval.lowerBound Interval Rational
i of
          Extended Rational
Interval.NegInf -> do
            Key
v1 <- M Key
gensym
            Key
v2 <- M Key
gensym
            VarMap (Expr Rational)
-> StateT Key Identity (VarMap (Expr Rational))
forall (m :: * -> *) a. Monad m => a -> m a
return (VarMap (Expr Rational)
 -> StateT Key Identity (VarMap (Expr Rational)))
-> VarMap (Expr Rational)
-> StateT Key Identity (VarMap (Expr Rational))
forall a b. (a -> b) -> a -> b
$ Key -> Expr Rational -> VarMap (Expr Rational)
forall a. Key -> a -> IntMap a
IM.singleton Key
v (Key -> Expr Rational
forall r. Num r => Key -> Expr r
LA.var Key
v1 Expr Rational -> Expr Rational -> Expr Rational
forall v. AdditiveGroup v => v -> v -> v
^-^ Key -> Expr Rational
forall r. Num r => Key -> Expr r
LA.var Key
v2)
          Interval.Finite Rational
lb
            | Rational
lb Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
>= Rational
0   -> VarMap (Expr Rational)
-> StateT Key Identity (VarMap (Expr Rational))
forall (m :: * -> *) a. Monad m => a -> m a
return VarMap (Expr Rational)
forall a. IntMap a
IM.empty
            | Bool
otherwise -> do
                Key
v1 <- M Key
gensym
                VarMap (Expr Rational)
-> StateT Key Identity (VarMap (Expr Rational))
forall (m :: * -> *) a. Monad m => a -> m a
return (VarMap (Expr Rational)
 -> StateT Key Identity (VarMap (Expr Rational)))
-> VarMap (Expr Rational)
-> StateT Key Identity (VarMap (Expr Rational))
forall a b. (a -> b) -> a -> b
$ Key -> Expr Rational -> VarMap (Expr Rational)
forall a. Key -> a -> IntMap a
IM.singleton Key
v (Key -> Expr Rational
forall r. Num r => Key -> Expr r
LA.var Key
v1 Expr Rational -> Expr Rational -> Expr Rational
forall v. AdditiveGroup v => v -> v -> v
^-^ Rational -> Expr Rational
forall r. (Num r, Eq r) => r -> Expr r
LA.constant Rational
lb)
      let obj2 :: Expr Rational
obj2 = VarMap (Expr Rational) -> Expr Rational -> Expr Rational
forall r. (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
LA.applySubst VarMap (Expr Rational)
s Expr Rational
obj

      [(Expr Rational, Rational)]
cs2 <- ([Maybe (Expr Rational, Rational)] -> [(Expr Rational, Rational)])
-> StateT Key Identity [Maybe (Expr Rational, Rational)]
-> StateT Key Identity [(Expr Rational, Rational)]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [Maybe (Expr Rational, Rational)] -> [(Expr Rational, Rational)]
forall a. [Maybe a] -> [a]
catMaybes (StateT Key Identity [Maybe (Expr Rational, Rational)]
 -> StateT Key Identity [(Expr Rational, Rational)])
-> StateT Key Identity [Maybe (Expr Rational, Rational)]
-> StateT Key Identity [(Expr Rational, Rational)]
forall a b. (a -> b) -> a -> b
$ [OrdRel (Expr Rational)]
-> (OrdRel (Expr Rational)
    -> StateT Key Identity (Maybe (Expr Rational, Rational)))
-> StateT Key Identity [Maybe (Expr Rational, Rational)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [OrdRel (Expr Rational)]
cs ((OrdRel (Expr Rational)
  -> StateT Key Identity (Maybe (Expr Rational, Rational)))
 -> StateT Key Identity [Maybe (Expr Rational, Rational)])
-> (OrdRel (Expr Rational)
    -> StateT Key Identity (Maybe (Expr Rational, Rational)))
-> StateT Key Identity [Maybe (Expr Rational, Rational)]
forall a b. (a -> b) -> a -> b
$ \(OrdRel Expr Rational
lhs RelOp
op Expr Rational
rhs) -> do
        case Key -> Expr Rational -> (Rational, Expr Rational)
forall r. Num r => Key -> Expr r -> (r, Expr r)
LA.extract Key
LA.unitVar (VarMap (Expr Rational) -> Expr Rational -> Expr Rational
forall r. (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
LA.applySubst VarMap (Expr Rational)
s (Expr Rational
lhs Expr Rational -> Expr Rational -> Expr Rational
forall v. AdditiveGroup v => v -> v -> v
^-^ Expr Rational
rhs)) of
          (Rational
c,Expr Rational
e) -> do
            let (Expr Rational
lhs2,RelOp
op2,Rational
rhs2) =
                  if -Rational
c Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
>= Rational
0
                  then (Expr Rational
e,RelOp
op,-Rational
c)
                  else (Expr Rational -> Expr Rational
forall v. AdditiveGroup v => v -> v
negateV Expr Rational
e, RelOp -> RelOp
flipOp RelOp
op, Rational
c)
            case RelOp
op2 of
              RelOp
Eql -> Maybe (Expr Rational, Rational)
-> StateT Key Identity (Maybe (Expr Rational, Rational))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Expr Rational, Rational)
 -> StateT Key Identity (Maybe (Expr Rational, Rational)))
-> Maybe (Expr Rational, Rational)
-> StateT Key Identity (Maybe (Expr Rational, Rational))
forall a b. (a -> b) -> a -> b
$ (Expr Rational, Rational) -> Maybe (Expr Rational, Rational)
forall a. a -> Maybe a
Just (Expr Rational
lhs2,Rational
rhs2)
              RelOp
Le  -> do
                Key
v <- M Key
gensym
                Maybe (Expr Rational, Rational)
-> StateT Key Identity (Maybe (Expr Rational, Rational))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Expr Rational, Rational)
 -> StateT Key Identity (Maybe (Expr Rational, Rational)))
-> Maybe (Expr Rational, Rational)
-> StateT Key Identity (Maybe (Expr Rational, Rational))
forall a b. (a -> b) -> a -> b
$ (Expr Rational, Rational) -> Maybe (Expr Rational, Rational)
forall a. a -> Maybe a
Just (Expr Rational
lhs2 Expr Rational -> Expr Rational -> Expr Rational
forall v. AdditiveGroup v => v -> v -> v
^+^ Key -> Expr Rational
forall r. Num r => Key -> Expr r
LA.var Key
v, Rational
rhs2)
              RelOp
Ge  -> do
                case Expr Rational -> [(Rational, Key)]
forall r. Expr r -> [(r, Key)]
LA.terms Expr Rational
lhs2 of
                  [(Rational
1,Key
_)] | Rational
rhs2Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
<=Rational
0 -> Maybe (Expr Rational, Rational)
-> StateT Key Identity (Maybe (Expr Rational, Rational))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Expr Rational, Rational)
forall a. Maybe a
Nothing
                  [(Rational, Key)]
_ -> do
                    Key
v <- M Key
gensym
                    Maybe (Expr Rational, Rational)
-> StateT Key Identity (Maybe (Expr Rational, Rational))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Expr Rational, Rational)
 -> StateT Key Identity (Maybe (Expr Rational, Rational)))
-> Maybe (Expr Rational, Rational)
-> StateT Key Identity (Maybe (Expr Rational, Rational))
forall a b. (a -> b) -> a -> b
$ (Expr Rational, Rational) -> Maybe (Expr Rational, Rational)
forall a. a -> Maybe a
Just (Expr Rational
lhs2 Expr Rational -> Expr Rational -> Expr Rational
forall v. AdditiveGroup v => v -> v -> v
^-^ Key -> Expr Rational
forall r. Num r => Key -> Expr r
LA.var Key
v, Rational
rhs2)
              RelOp
_   -> [Char] -> StateT Key Identity (Maybe (Expr Rational, Rational))
forall a. HasCallStack => [Char] -> a
error ([Char] -> StateT Key Identity (Maybe (Expr Rational, Rational)))
-> [Char] -> StateT Key Identity (Maybe (Expr Rational, Rational))
forall a b. (a -> b) -> a -> b
$ [Char]
"ToySolver.LPUtil.toStandardForm: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ RelOp -> [Char]
forall a. Show a => a -> [Char]
show RelOp
op2 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is not supported"

      Bool -> StateT Key Identity () -> StateT Key Identity ()
forall a. HasCallStack => Bool -> a -> a
assert ([Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [Maybe Rational -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe Rational -> Bool) -> Maybe Rational -> Bool
forall a b. (a -> b) -> a -> b
$ Key -> Expr Rational -> Maybe Rational
forall r. Num r => Key -> Expr r -> Maybe r
LA.lookupCoeff Key
LA.unitVar Expr Rational
c | (Expr Rational
c,Rational
_) <- [(Expr Rational, Rational)]
cs2]) (StateT Key Identity () -> StateT Key Identity ())
-> StateT Key Identity () -> StateT Key Identity ()
forall a b. (a -> b) -> a -> b
$ () -> StateT Key Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

      ((Expr Rational, [(Expr Rational, Rational)]),
 VarMap (Expr Rational))
-> State
     Key
     ((Expr Rational, [(Expr Rational, Rational)]),
      VarMap (Expr Rational))
forall (m :: * -> *) a. Monad m => a -> m a
return ((Expr Rational
obj2,[(Expr Rational, Rational)]
cs2),VarMap (Expr Rational)
s)