-----------------------------------------------------------------------------
-- |
-- 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 = [(Var, Rational)] -> Model Rational
forall a. [(Var, a)] -> IntMap a
IM.fromAscList ([(Var, Rational)] -> Model Rational)
-> [(Var, Rational)] -> Model Rational
forall a b. (a -> b) -> a -> b
$ do
      Var
v <- IntSet -> [Var]
IS.toAscList IntSet
vs
      case Var -> VarMap (Expr Rational) -> Maybe (Expr Rational)
forall a. Var -> IntMap a -> Maybe a
IM.lookup Var
v VarMap (Expr Rational)
s of
        Just Expr Rational
def -> (Var, Rational) -> [(Var, Rational)]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return (Var
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  -> (Var, Rational) -> [(Var, Rational)]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return (Var
v, Model Rational
m Model Rational -> Var -> Rational
forall a. IntMap a -> Var -> a
IM.! Var
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 :: Var
v1 = if IntSet -> Bool
IS.null IntSet
vs then Var
0 else IntSet -> Var
IS.findMax IntSet
vs Var -> Var -> Var
forall a. Num a => a -> a -> a
+ Var
1
    initialBounds :: IntMap (Interval Rational)
initialBounds = [(Var, Interval Rational)] -> IntMap (Interval Rational)
forall a. [(Var, a)] -> IntMap a
IM.fromList [(Var
v, Interval Rational
forall r. Ord r => Interval r
Interval.whole) | Var
v <- IntSet -> [Var]
IS.toList IntSet
vs]
    bounds :: IntMap (Interval Rational)
bounds = IntMap (Interval Rational)
-> [OrdRel (Expr Rational)]
-> IntSet
-> Var
-> IntMap (Interval Rational)
forall r.
RealFrac r =>
BoundsEnv r -> [Atom r] -> IntSet -> Var -> BoundsEnv r
BI.inferBounds IntMap (Interval Rational)
initialBounds [OrdRel (Expr Rational)]
cs IntSet
IS.empty Var
10

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

    m :: ((Expr Rational, [(Expr Rational, Rational)]),
 VarMap (Expr Rational))
m = (State
   Var
   ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
 -> Var
 -> ((Expr Rational, [(Expr Rational, Rational)]),
     VarMap (Expr Rational)))
-> Var
-> State
     Var
     ((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
  Var
  ((Expr Rational, [(Expr Rational, Rational)]),
   VarMap (Expr Rational))
-> Var
-> ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
forall s a. State s a -> s -> a
evalState Var
v1 (State
   Var
   ((Expr Rational, [(Expr Rational, Rational)]),
    VarMap (Expr Rational))
 -> ((Expr Rational, [(Expr Rational, Rational)]),
     VarMap (Expr Rational)))
-> State
     Var
     ((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 Var Identity [VarMap (Expr Rational)]
-> StateT Var 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 Var Identity [VarMap (Expr Rational)]
 -> StateT Var Identity (VarMap (Expr Rational)))
-> StateT Var Identity [VarMap (Expr Rational)]
-> StateT Var Identity (VarMap (Expr Rational))
forall a b. (a -> b) -> a -> b
$ [(Var, Interval Rational)]
-> ((Var, Interval Rational)
    -> StateT Var Identity (VarMap (Expr Rational)))
-> StateT Var 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) -> [(Var, Interval Rational)]
forall a. IntMap a -> [(Var, a)]
IM.toList IntMap (Interval Rational)
bounds) (((Var, Interval Rational)
  -> StateT Var Identity (VarMap (Expr Rational)))
 -> StateT Var Identity [VarMap (Expr Rational)])
-> ((Var, Interval Rational)
    -> StateT Var Identity (VarMap (Expr Rational)))
-> StateT Var Identity [VarMap (Expr Rational)]
forall a b. (a -> b) -> a -> b
$ \(Var
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
            Var
v1 <- M Var
gensym
            Var
v2 <- M Var
gensym
            VarMap (Expr Rational)
-> StateT Var Identity (VarMap (Expr Rational))
forall a. a -> StateT Var Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (VarMap (Expr Rational)
 -> StateT Var Identity (VarMap (Expr Rational)))
-> VarMap (Expr Rational)
-> StateT Var Identity (VarMap (Expr Rational))
forall a b. (a -> b) -> a -> b
$ Var -> Expr Rational -> VarMap (Expr Rational)
forall a. Var -> a -> IntMap a
IM.singleton Var
v (Var -> Expr Rational
forall r. Num r => Var -> Expr r
LA.var Var
v1 Expr Rational -> Expr Rational -> Expr Rational
forall v. AdditiveGroup v => v -> v -> v
^-^ Var -> Expr Rational
forall r. Num r => Var -> Expr r
LA.var Var
v2)
          Interval.Finite Rational
lb
            | Rational
lb Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
>= Rational
0   -> VarMap (Expr Rational)
-> StateT Var Identity (VarMap (Expr Rational))
forall a. a -> StateT Var Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return VarMap (Expr Rational)
forall a. IntMap a
IM.empty
            | Bool
otherwise -> do
                Var
v1 <- M Var
gensym
                VarMap (Expr Rational)
-> StateT Var Identity (VarMap (Expr Rational))
forall a. a -> StateT Var Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (VarMap (Expr Rational)
 -> StateT Var Identity (VarMap (Expr Rational)))
-> VarMap (Expr Rational)
-> StateT Var Identity (VarMap (Expr Rational))
forall a b. (a -> b) -> a -> b
$ Var -> Expr Rational -> VarMap (Expr Rational)
forall a. Var -> a -> IntMap a
IM.singleton Var
v (Var -> Expr Rational
forall r. Num r => Var -> Expr r
LA.var Var
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 Var Identity [Maybe (Expr Rational, Rational)]
-> StateT Var 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 Var Identity [Maybe (Expr Rational, Rational)]
 -> StateT Var Identity [(Expr Rational, Rational)])
-> StateT Var Identity [Maybe (Expr Rational, Rational)]
-> StateT Var Identity [(Expr Rational, Rational)]
forall a b. (a -> b) -> a -> b
$ [OrdRel (Expr Rational)]
-> (OrdRel (Expr Rational)
    -> StateT Var Identity (Maybe (Expr Rational, Rational)))
-> StateT Var 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 Var Identity (Maybe (Expr Rational, Rational)))
 -> StateT Var Identity [Maybe (Expr Rational, Rational)])
-> (OrdRel (Expr Rational)
    -> StateT Var Identity (Maybe (Expr Rational, Rational)))
-> StateT Var Identity [Maybe (Expr Rational, Rational)]
forall a b. (a -> b) -> a -> b
$ \(OrdRel Expr Rational
lhs RelOp
op Expr Rational
rhs) -> do
        case Var -> Expr Rational -> (Rational, Expr Rational)
forall r. Num r => Var -> Expr r -> (r, Expr r)
LA.extract Var
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 Var Identity (Maybe (Expr Rational, Rational))
forall a. a -> StateT Var Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Expr Rational, Rational)
 -> StateT Var Identity (Maybe (Expr Rational, Rational)))
-> Maybe (Expr Rational, Rational)
-> StateT Var 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
                Var
v <- M Var
gensym
                Maybe (Expr Rational, Rational)
-> StateT Var Identity (Maybe (Expr Rational, Rational))
forall a. a -> StateT Var Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Expr Rational, Rational)
 -> StateT Var Identity (Maybe (Expr Rational, Rational)))
-> Maybe (Expr Rational, Rational)
-> StateT Var 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
^+^ Var -> Expr Rational
forall r. Num r => Var -> Expr r
LA.var Var
v, Rational
rhs2)
              RelOp
Ge  -> do
                case Expr Rational -> [(Rational, Var)]
forall r. Expr r -> [(r, Var)]
LA.terms Expr Rational
lhs2 of
                  [(Rational
1,Var
_)] | Rational
rhs2Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
<=Rational
0 -> Maybe (Expr Rational, Rational)
-> StateT Var Identity (Maybe (Expr Rational, Rational))
forall a. a -> StateT Var Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Expr Rational, Rational)
forall a. Maybe a
Nothing
                  [(Rational, Var)]
_ -> do
                    Var
v <- M Var
gensym
                    Maybe (Expr Rational, Rational)
-> StateT Var Identity (Maybe (Expr Rational, Rational))
forall a. a -> StateT Var Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Expr Rational, Rational)
 -> StateT Var Identity (Maybe (Expr Rational, Rational)))
-> Maybe (Expr Rational, Rational)
-> StateT Var 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
^-^ Var -> Expr Rational
forall r. Num r => Var -> Expr r
LA.var Var
v, Rational
rhs2)
              RelOp
_   -> [Char] -> StateT Var Identity (Maybe (Expr Rational, Rational))
forall a. HasCallStack => [Char] -> a
error ([Char] -> StateT Var Identity (Maybe (Expr Rational, Rational)))
-> [Char] -> StateT Var 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 Var Identity () -> StateT Var 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
$ Var -> Expr Rational -> Maybe Rational
forall r. Num r => Var -> Expr r -> Maybe r
LA.lookupCoeff Var
LA.unitVar Expr Rational
c | (Expr Rational
c,Rational
_) <- [(Expr Rational, Rational)]
cs2]) (StateT Var Identity () -> StateT Var Identity ())
-> StateT Var Identity () -> StateT Var Identity ()
forall a b. (a -> b) -> a -> b
$ () -> StateT Var Identity ()
forall a. a -> StateT Var Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

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