{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{- |
The following LP problem

maximize @4 x_1 - 3 x_2 + 2 x_3@ subject to

@2 x_1 + x_2 <= 10@

@x_2 + 5 x_3 <= 20@

and

@x_i >= 0@

is used as an example in the doctest comments.


By default all indeterminates are non-negative.
A given bound for a variable completely replaces the default,
so @0 <= x_i <= b@ must be explicitly given as @i >=<. (0,b)@.
Multiple bounds for a variable are not allowed,
instead of @[i >=. a, i <=. b]@ use @i >=<. (a,b)@.
-}
module Numeric.GLPK (
   Term(..),
   Bound(..),
   Inequality(..),
   free, (<=.), (>=.), (==.), (>=<.),
   NoSolutionType(..),
   SolutionType(..),
   Solution,
   Constraints,
   Direction(..),
   Objective,
   Bounds,
   (.*),
   objectiveFromTerms,
   simplex,
   simplexMulti,
   exact,
   interior,
   interiorMulti,
   ) where

import qualified Math.Programming.Glpk.Header as FFI

import qualified Data.Array.Comfort.Storable.Mutable as Mutable
import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.NonEmpty as NonEmpty
import Data.Array.Comfort.Storable (Array)
import Data.Tuple.HT (mapFst, mapSnd)
import Data.Maybe (fromMaybe)
import Data.Foldable (traverse_, for_)

import Control.Monad (void)
import Control.Applicative (liftA2)
import Control.Exception (bracket)
import Control.DeepSeq (NFData, rnf)

import System.IO.Unsafe (unsafePerformIO)

import qualified Foreign
import Foreign.Ptr (nullPtr)
import Foreign.C.Types (CDouble)

{- $setup
>>> import qualified Numeric.GLPK as LP
>>> import Numeric.GLPK ((.*), (<=.), (==.))
>>>
>>> import qualified Data.Array.Comfort.Storable as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>>
>>> import Data.Tuple.HT (mapPair, mapSnd)
>>>
>>> import qualified Test.QuickCheck as QC
>>> import Test.QuickCheck ((===), (.&&.), (.||.))
>>>
>>> type X = Shape.Element
>>> type PairShape = Shape.NestedTuple Shape.TupleIndex (X,X)
>>> type TripletShape = Shape.NestedTuple Shape.TupleIndex (X,X,X)
>>>
>>> pairShape :: PairShape
>>> pairShape = Shape.static
>>>
>>> tripletShape :: TripletShape
>>> tripletShape = Shape.static
>>>
>>> round3 :: Double -> Double
>>> round3 x = fromInteger (round (1000*x)) / 1000
-}


infix 7 .*

(.*) :: Double -> ix -> Term ix
(.*) = Term

data Term ix = Term Double ix
   deriving (Show)


infix 4 <=., >=., >=<., ==.

(<=.), (>=.), (==.) :: x -> Double -> Inequality x
x <=. bnd = Inequality x $ LessEqual bnd
x >=. bnd = Inequality x $ GreaterEqual bnd
x ==. bnd = Inequality x $ Equal bnd

(>=<.) :: x -> (Double,Double) -> Inequality x
x >=<. bnd = Inequality x $ uncurry Between bnd

free :: x -> Inequality x
free x = Inequality x Free

data Inequality x = Inequality x Bound
   deriving Show

data Bound =
     LessEqual Double
   | GreaterEqual Double
   | Between Double Double
   | Equal Double
   | Free
   deriving Show

instance Functor Inequality where
   fmap f (Inequality x bnd)  =  Inequality (f x) bnd

data NoSolutionType =
     Undefined
   | NoFeasible
   | Unbounded
   deriving (Eq, Show)

data SolutionType =
     Feasible
   | Infeasible
   | Optimal
   deriving (Eq, Show)

instance NFData NoSolutionType where
    rnf NoFeasible = ()
    rnf _ = ()

instance NFData SolutionType where
    rnf Optimal = ()
    rnf _ = ()

type Solution sh =
      Either NoSolutionType (SolutionType, (Double, Array sh Double))

type Constraints ix = [Inequality [Term ix]]

data Direction = Minimize | Maximize

type Objective sh = Array sh Double

type Bounds ix = [Inequality ix]



objectiveFromTerms ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) => sh -> [Term ix] -> Objective sh
objectiveFromTerms sh =
   Array.fromAssociations 0 sh . map (\(Term x ix) -> (ix,x))


prepareBounds :: Inequality a -> (a, (FFI.GlpkConstraintType, CDouble, CDouble))
prepareBounds (Inequality x bnd) =
   (,) x $
   (\(bndType,lo,up) -> (bndType, realToFrac lo, realToFrac up)) $
   case bnd of
      LessEqual up    -> (FFI.glpkLT,      0,  up)
      GreaterEqual lo -> (FFI.glpkGT,      lo, 0)
      Between lo up   -> (FFI.glpkBounded, lo, up)
      Equal y         -> (FFI.glpkFixed,   y,  y)
      Free            -> (FFI.glpkFree,    0,  0)

columnIndex :: (Shape.Indexed sh, Shape.Index sh ~ ix) => sh -> ix -> FFI.Column
columnIndex shape var = 1 + fromIntegral (Shape.offset shape var)

deferredColumnIndex :: Shape.DeferredIndex ix -> FFI.Column
deferredColumnIndex (Shape.DeferredIndex var) = 1 + fromIntegral var

allocaArray :: (Foreign.Storable a) => Int -> (FFI.GlpkArray a -> IO b) -> IO b
allocaArray n f = Foreign.allocaArray (n+1) $ f . FFI.GlpkArray

pokeElem :: (Foreign.Storable a) => FFI.GlpkArray a -> Int -> a -> IO ()
pokeElem (FFI.GlpkArray ptr) k a = Foreign.pokeElemOff ptr k a


{- |
>>> case Shape.indexTupleFromShape tripletShape of (x1,x2,x3) -> mapSnd (mapSnd Array.toTuple) <$> LP.simplex [] [[2.*x1, 1.*x2] <=. 10, [1.*x2, 5.*x3] <=. 20] (LP.Maximize, Array.fromTuple (4,-3,2) :: Array.Array TripletShape Double)
Right (Optimal,(28.0,(5.0,0.0,4.0)))

prop> \(QC.Positive posWeight) (QC.Positive negWeight) target -> case Shape.indexTupleFromShape pairShape of (pos,neg) -> case mapSnd (mapSnd Array.toTuple) <$> LP.simplex [] [[1.*pos, (-1).*neg] ==. target] (LP.Minimize, Array.fromTuple (posWeight,negWeight) :: Array.Array PairShape Double) of (Right (LP.Optimal,(absol,(posResult,negResult)))) -> QC.property (absol>=0) .&&. (posResult === 0 .||. negResult === 0); _ -> QC.property False
-}
simplex ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Bounds ix -> Constraints ix ->
   (Direction, Objective sh) -> Solution sh
simplex = solve (flip FFI.glp_simplex nullPtr)

{- |
Optimize for one objective after another.
That is, if the first optimization succeeds
then the optimum is fixed as constraint
and the optimization is continued with respect to the second objective and so on.
The iteration fails if one optimization fails.
The obtained objective values are returned as well.
Their number equals the number of attempted optimizations.

The last objective value is included in the Solution value.
This is a bit inconsistent,
but this way you have a warranty that there is an objective value
if the optimization is successful.

The objectives are expected as 'Term's
because after successful optimization step
they are used as (sparse) constraints.
It's also easy to assert that the same array shape
is used for all objectives.

The function does not work reliably,
because an added objective can make the system infeasible
due to rounding errors.
E.g. a non-negative objective can become very small but negative.
-}
simplexMulti, interiorMulti ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Bounds ix -> Constraints ix ->
   sh -> NonEmpty.T [] (Direction, [Term ix]) -> ([Double], Solution sh)
simplexMulti = solveMulti . simplex
interiorMulti = solveMulti . interior

solveMulti ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   (Constraints ix -> (Direction, Objective sh) -> Solution sh) ->
   Constraints ix ->
   sh -> NonEmpty.T [] (Direction, [Term ix]) -> ([Double], Solution sh)
solveMulti solver constrs sh (NonEmpty.Cons obj0 objs0) =
   let go curObj ((dir,obj):objs) (Right (Optimal, (opt,_))) =
         mapFst (opt:) $
         go obj objs $
         solver
            ((curObj==.opt) : constrs)
            (dir, objectiveFromTerms sh obj)
       go _ _ sol = ([], sol)
   in go (snd obj0) objs0 $
      solver constrs $ mapSnd (objectiveFromTerms sh) obj0


{- |
>>> case Shape.indexTupleFromShape tripletShape of (x1,x2,x3) -> mapSnd (mapSnd Array.toTuple) <$> LP.exact [] [[2.*x1, 1.*x2] <=. 10, [1.*x2, 5.*x3] <=. 20] (LP.Maximize, Array.fromTuple (4,-3,2) :: Array.Array TripletShape Double)
Right (Optimal,(28.0,(5.0,0.0,4.0)))
-}
exact ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Bounds ix -> Constraints ix ->
   (Direction, Objective sh) -> Solution sh
exact = solve (flip FFI.glp_exact nullPtr)


{-# INLINE solve #-}
solve ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   (Foreign.Ptr FFI.Problem -> IO FFI.GlpkSimplexStatus) ->
   Bounds ix -> Constraints ix ->
   (Direction, Objective sh) -> Solution sh
solve solver bounds constrs (dir,obj) = unsafePerformIO $
   bracket FFI.glp_create_prob FFI.glp_delete_prob $ \lp -> do
   storeProblem bounds constrs (dir,obj) lp
   void $ solver lp
   let examine =
         liftA2 (,)
            (realToFrac <$> FFI.glp_get_obj_val lp)
            (readGLPArray (Array.shape obj) $ \arr ix ->
               Mutable.write arr ix . realToFrac
                  =<< FFI.glp_get_col_prim lp (deferredColumnIndex ix))
   status <- FFI.glp_get_status lp
   either (return . Left) (\typ -> Right . (,) typ <$> examine) $
      analyzeStatus status


{- |
>>> case Shape.indexTupleFromShape tripletShape of (x1,x2,x3) -> mapSnd (mapPair (round3, Array.toTuple . Array.map round3)) <$> LP.interior [] [[2.*x1, 1.*x2] <=. 10, [1.*x2, 5.*x3] <=. 20] (LP.Maximize, Array.fromTuple (4,-3,2) :: Array.Array TripletShape Double)
Right (Optimal,(28.0,(5.0,0.0,4.0)))
-}
interior ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Bounds ix -> Constraints ix ->
   (Direction, Objective sh) -> Solution sh
interior bounds constrs (dir,obj) = unsafePerformIO $
   bracket FFI.glp_create_prob FFI.glp_delete_prob $ \lp -> do
   storeProblem bounds constrs (dir,obj) lp
   void $ FFI.glp_interior lp nullPtr
   let examine =
         liftA2 (,)
            (realToFrac <$> FFI.glp_ipt_obj_val lp)
            (readGLPArray (Array.shape obj) $ \arr ix ->
               Mutable.write arr ix . realToFrac
                  =<< FFI.glp_ipt_col_prim lp (deferredColumnIndex ix))
   status <- FFI.glp_ipt_status lp
   either (return . Left) (\typ -> Right . (,) typ <$> examine) $
      analyzeStatus status


storeProblem ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Bounds ix -> Constraints ix ->
   (Direction, Objective sh) -> Foreign.Ptr FFI.Problem -> IO ()
storeProblem bounds constrs (dir,obj) lp = do
   let shape = Array.shape obj
   void $ FFI.glp_term_out FFI.glpkOff
   FFI.glp_set_obj_dir lp $
      case dir of
         Minimize -> FFI.glpkMin
         Maximize -> FFI.glpkMax
   firstRow <- FFI.glp_add_rows lp $ fromIntegral $ length constrs
   for_ (zip [firstRow..] $
      map prepareBounds constrs) $ \(row,(_x,(bnd,lo,up))) ->
      FFI.glp_set_row_bnds lp row bnd lo up
   _firstCol <- FFI.glp_add_cols lp $ fromIntegral $ Shape.size shape
   for_ (Shape.indices $ Shape.Deferred shape) $ \x ->
      FFI.glp_set_col_bnds lp (deferredColumnIndex x) FFI.glpkGT 0 0
   for_ (map prepareBounds bounds) $ \(x,(bnd,lo,up)) ->
      FFI.glp_set_col_bnds lp (columnIndex shape x) bnd lo up
   for_ (Array.toAssociations obj) $ \(x,c) ->
      FFI.glp_set_obj_coef lp (columnIndex shape x) (realToFrac c)
   let numRows = length $ concatMap (fst . prepareBounds) constrs
   allocaArray numRows $ \ia ->
      allocaArray numRows $ \ja ->
      allocaArray numRows $ \ar -> do
      for_ (zip [1..] $ concat $
            zipWith (map . (,)) [firstRow..] $
            map (fst . prepareBounds) constrs) $
         \(k, (row, Term c x)) -> do
            pokeElem ia k row
            pokeElem ja k (columnIndex shape x)
            pokeElem ar k (realToFrac c)
      FFI.glp_load_matrix lp (fromIntegral numRows) ia ja ar

{-# INLINE readGLPArray #-}
readGLPArray ::
   (Shape.C sh, Foreign.Storable a, Num a) =>
   sh ->
   (Mutable.Array IO (Shape.Deferred sh) a ->
    Shape.DeferredIndex sh -> IO ()) ->
   IO (Array sh a)
readGLPArray shape act = do
   let defShape = Shape.Deferred shape
   arr <- Mutable.new defShape 0
   traverse_ (act arr) (Shape.indices defShape)
   Array.reshape shape <$> Mutable.freeze arr

analyzeStatus :: FFI.GlpkSolutionStatus -> Either NoSolutionType SolutionType
analyzeStatus status =
   fromMaybe (error "glpk-solver: solution type unknown") $ lookup status $
      (FFI.glpkUndefined,  Left Undefined) :
      (FFI.glpkFeasible,   Right Feasible) :
      (FFI.glpkInfeasible, Right Infeasible) :
      (FFI.glpkNoFeasible, Left NoFeasible) :
      (FFI.glpkOptimal,    Right Optimal) :
      (FFI.glpkUnbounded,  Left Unbounded) :
      []