{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{- |
The monadic interface to GLPK allows to optimize
with respect to multiple objectives, successively.
-}
module Numeric.GLPK.Monad (
   T,
   run,
   simplex,
   exact,
   Direction(..),
   ) where

import qualified Math.Programming.Glpk.Header as FFI
import Numeric.GLPK.Private
         (Term(Term), Constraints, Solution,
          allocaArray, pokeElem, columnIndex, prepareBounds, storeBounds,
          setDirection, setObjective, peekSimplexSolution)
import Numeric.GLPK (Bounds, Direction(..), Objective)

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Foldable (for_)

import qualified Control.Monad.Trans.Reader as MR
import Control.Monad.IO.Class (liftIO)
import Control.Monad (void, when)
import Control.Exception (bracket)

import System.IO.Unsafe (unsafePerformIO)

import Foreign.Ptr (Ptr, nullPtr)


{- $setup
>>> import qualified Numeric.GLPK.Monad as LP
>>> import Test.Numeric.GLPK (TripletShape, tripletShape)
>>> import Numeric.GLPK ((.*), (<=.))
>>>
>>> import qualified Data.Array.Comfort.Storable as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>>
>>> import Data.Tuple.HT (mapSnd)
-}


newtype T sh a = Cons (MR.ReaderT (sh, Ptr FFI.Problem) IO a)
   deriving (Functor, Applicative, Monad)

run ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   sh -> Bounds ix -> T sh a -> a
run shape bounds (Cons act) =
   unsafePerformIO $
   bracket FFI.glp_create_prob FFI.glp_delete_prob $ \lp -> do
      void $ FFI.glp_term_out FFI.glpkOff
      storeBounds lp shape bounds
      liftIO $ MR.runReaderT act (shape, lp)


{- |
Add new constraints to an existing problem
and solve with a new direction and objective.

>>> case Shape.indexTupleFromShape tripletShape of (x,y,z) -> mapSnd (mapSnd Array.toTuple) <$> LP.run tripletShape [] (LP.simplex [[2.*x, 1.*y] <=. 10, [1.*y, (5::Double).*z] <=. 20] (LP.Maximize, Array.fromTuple (4,-3,2) :: Array.Array TripletShape Double))
Right (Optimal,(28.0,(5.0,0.0,4.0)))
-}
simplex ::
   (Eq sh, Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Constraints ix ->
   (Direction, Objective sh) -> T sh (Solution sh)
simplex = solve (flip FFI.glp_simplex nullPtr)

exact ::
   (Eq sh, Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Constraints ix ->
   (Direction, Objective sh) -> T sh (Solution sh)
exact = solve (flip FFI.glp_exact nullPtr)

solve ::
   (Eq sh, Shape.Indexed sh, Shape.Index sh ~ ix) =>
   (Ptr FFI.Problem -> IO FFI.GlpkSimplexStatus) ->
   Constraints ix ->
   (Direction, Objective sh) ->
   T sh (Solution sh)
solve method constrs (dir,obj) = Cons $ do
   (shape, lp) <- MR.ask
   when (shape /= Array.shape obj) $
      error "GLPK.Monad.solve: objective shape mismatch"

   liftIO $ do
      setDirection lp dir
      setObjective lp obj
      newRow <- FFI.glp_add_rows lp $ fromIntegral $ length constrs
      for_ (zip [newRow..] (map prepareBounds constrs)) $
            \(row, (terms,(bnd,lo,up))) -> do
         FFI.glp_set_row_bnds lp row bnd lo up
         let numTerms = length terms
         allocaArray numTerms $ \indicesPtr ->
            allocaArray numTerms $ \coeffsPtr -> do
            for_ (zip [1..] terms) $
               \(k, Term c x) -> do
                  pokeElem indicesPtr k (columnIndex shape x)
                  pokeElem coeffsPtr k (realToFrac c)
            FFI.glp_set_mat_row lp row
               (fromIntegral numTerms) indicesPtr coeffsPtr
      void $ method lp
      peekSimplexSolution shape lp