-----------------------------------------------------------------------------
-- |
-- 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 = forall a. Variables a => a -> IntSet
vars Expr Rational
obj IntSet -> IntSet -> IntSet
`IS.union` 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 = forall a. [(Var, a)] -> IntMap a
IM.fromAscList forall a b. (a -> b) -> a -> b
$ do
      Var
v <- IntSet -> [Var]
IS.toAscList IntSet
vs
      case forall a. Var -> IntMap a -> Maybe a
IM.lookup Var
v VarMap (Expr Rational)
s of
        Just Expr Rational
def -> forall (m :: * -> *) a. Monad m => a -> m a
return (Var
v, forall m e v. Eval m e v => m -> e -> v
LA.eval Model Rational
m Expr Rational
def)
        Maybe (Expr Rational)
Nothing  -> forall (m :: * -> *) a. Monad m => a -> m a
return (Var
v, Model Rational
m 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 = forall a. Variables a => a -> IntSet
vars Expr Rational
obj IntSet -> IntSet -> IntSet
`IS.union` 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 forall a. Num a => a -> a -> a
+ Var
1
    initialBounds :: IntMap (Interval Rational)
initialBounds = forall a. [(Var, a)] -> IntMap a
IM.fromList [(Var
v, forall r. Ord r => Interval r
Interval.whole) | Var
v <- IntSet -> [Var]
IS.toList IntSet
vs]
    bounds :: IntMap (Interval Rational)
bounds = 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 <- forall s (m :: * -> *). MonadState s m => m s
get
      forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ Var
vforall a. Num a => a -> a -> a
+Var
1
      forall (m :: * -> *) a. Monad m => a -> m a
return Var
v

    m :: ((Expr Rational, [(Expr Rational, Rational)]),
 VarMap (Expr Rational))
m = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> a
evalState Var
v1 forall a b. (a -> b) -> a -> b
$ do
      VarMap (Expr Rational)
s <- forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a. IntMap a -> [(Var, a)]
IM.toList IntMap (Interval Rational)
bounds) forall a b. (a -> b) -> a -> b
$ \(Var
v,Interval Rational
i) -> do
        case 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
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Var -> a -> IntMap a
IM.singleton Var
v (forall r. Num r => Var -> Expr r
LA.var Var
v1 forall v. AdditiveGroup v => v -> v -> v
^-^ forall r. Num r => Var -> Expr r
LA.var Var
v2)
          Interval.Finite Rational
lb
            | Rational
lb forall a. Ord a => a -> a -> Bool
>= Rational
0   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. IntMap a
IM.empty
            | Bool
otherwise -> do
                Var
v1 <- M Var
gensym
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Var -> a -> IntMap a
IM.singleton Var
v (forall r. Num r => Var -> Expr r
LA.var Var
v1 forall v. AdditiveGroup v => v -> v -> v
^-^ forall r. (Num r, Eq r) => r -> Expr r
LA.constant Rational
lb)
      let obj2 :: Expr Rational
obj2 = 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 <- forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [OrdRel (Expr Rational)]
cs forall a b. (a -> b) -> a -> b
$ \(OrdRel Expr Rational
lhs RelOp
op Expr Rational
rhs) -> do
        case forall r. Num r => Var -> Expr r -> (r, Expr r)
LA.extract Var
LA.unitVar (forall r. (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
LA.applySubst VarMap (Expr Rational)
s (Expr Rational
lhs 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 forall a. Ord a => a -> a -> Bool
>= Rational
0
                  then (Expr Rational
e,RelOp
op,-Rational
c)
                  else (forall v. AdditiveGroup v => v -> v
negateV Expr Rational
e, RelOp -> RelOp
flipOp RelOp
op, Rational
c)
            case RelOp
op2 of
              RelOp
Eql -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Expr Rational
lhs2,Rational
rhs2)
              RelOp
Le  -> do
                Var
v <- M Var
gensym
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Expr Rational
lhs2 forall v. AdditiveGroup v => v -> v -> v
^+^ forall r. Num r => Var -> Expr r
LA.var Var
v, Rational
rhs2)
              RelOp
Ge  -> do
                case forall r. Expr r -> [(r, Var)]
LA.terms Expr Rational
lhs2 of
                  [(Rational
1,Var
_)] | Rational
rhs2forall a. Ord a => a -> a -> Bool
<=Rational
0 -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
                  [(Rational, Var)]
_ -> do
                    Var
v <- M Var
gensym
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Expr Rational
lhs2 forall v. AdditiveGroup v => v -> v -> v
^-^ forall r. Num r => Var -> Expr r
LA.var Var
v, Rational
rhs2)
              RelOp
_   -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"ToySolver.LPUtil.toStandardForm: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show RelOp
op2 forall a. [a] -> [a] -> [a]
++ [Char]
" is not supported"

      forall a. HasCallStack => Bool -> a -> a
assert (forall (t :: * -> *). Foldable t => t Bool -> Bool
and [forall a. Maybe a -> Bool
isNothing forall a b. (a -> b) -> a -> b
$ 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]) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return ()

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