{-#  LANGUAGE ForeignFunctionInterface, PatternGuards, RankNTypes, TypeFamilies #-}
{- |

Copyright: (C) 2013 Adam Vogt
Maintainer: Adam Vogt <vogt.adam@gmail.com>
Stability: unstable

Binding to ipopt <http://projects.coin-or.org/Ipopt>. Uses "Numeric.AD" to compute
derivatives required by ipopt. Current limitations include:

* derivatives are computed and stored without taking advantage of sparsity

* copying is done in converting between "Data.Vector.Storable" and "Data.Vector" might be unnecessary. Currently it is done because AD needs a Traversable structure, but Storable vectors are not traversable.

* probably doesn't work if @coin\/IpStdCInterface.h@ has Number =/= 'CDouble'

* no binding to SetIntermediateCallback

* garbage collection of 'IpProblem' won't free C-side resources

* specifying problems might be made easier by following (extending) the approach taken by glpk-hs

Refer to @Test1.hs@ for an example where the derivatives are computed by hand,
and @Test2.hs@ for the use of 'createIpoptProblemAD'.

-}
module Ipopt (
    -- * specifying problem
    createIpoptProblemAD,

    -- ** solve
    ipoptSolve,
    IpOptSolved(..),

    -- ** solver options
    addIpoptNumOption,
    addIpoptStrOption,
    addIpoptIntOption,
    openIpoptOutputFile,

    -- * types
    Vec,

    IpNumber(..),
    IpIndex(..),
    IpInt(..),
    IpBool(..),

    IpF(..),
    IpGradF(..),
    IpG(..),
    IpJacG(..),
    IpH(..),

    IpProblem(..),

    ApplicationReturnStatus(..),

    -- * lower-level parts of the binding
    createIpoptProblem,

    freeIpoptProblem,
    setIpoptProblemScaling,


    -- ** marshalling functions
    wrapIpF,
    wrapIpGradF,
    wrapIpG,
    wrapIpJacG,
    wrapIpH,

    wrapIpF1,
    wrapIpGradF1,
    wrapIpG1,
    wrapIpJacG1,
    wrapIpH1,

    wrapIpF2,
    wrapIpGradF2,
    wrapIpG2,
    wrapIpJacG2,
    wrapIpH2,

    ) where

import C2HS
import Control.Exception
import Control.Monad
import Data.IORef
import Foreign.C
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.Storable
import Numeric.AD
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VM



#include "coin/IpStdCInterface.h"

type IpNumber = {# type Number #}
type IpIndex = {# type Index #}
type IpInt = {# type Int #}
type IpBool = {# type Bool #}

type IpF = {# type Eval_F_CB #}
type IpGradF = {# type Eval_Grad_F_CB #}
type IpG = {# type Eval_G_CB #}
type IpJacG = {# type Eval_Jac_G_CB #}
type IpH = {# type Eval_H_CB #}

{#enum ApplicationReturnStatus as ^ {underscoreToCase} deriving (Show) #}
{#enum AlgorithmMode as ^ {underscoreToCase} deriving (Show) #}


newtype IpProblem = IpProblem { unIpProblem :: Ptr ()}

type family UnFunPtr a
type instance UnFunPtr (FunPtr a) = a

ipTrue = 1 :: IpBool
ipFalse = 0 :: IpBool

-- | likely an unsafe method for getting a "Data.Vector.Storable.Mutable" out of a 'Ptr'
ptrToVS n p = do
    fp <- newForeignPtr_ p
    return (VM.unsafeFromForeignPtr0 fp (fromIntegral n))

foreign import ccall "wrapper" wrapIpF1 :: UnFunPtr IpF -> IO IpF
foreign import ccall "wrapper" wrapIpG1 :: UnFunPtr IpG -> IO IpG
foreign import ccall "wrapper" wrapIpGradF1 :: UnFunPtr IpGradF -> IO IpGradF
foreign import ccall "wrapper" wrapIpJacG1 :: UnFunPtr IpJacG -> IO IpJacG
foreign import ccall "wrapper" wrapIpH1 :: UnFunPtr IpH -> IO IpH


toB x =  either (\ e@SomeException {} -> print e >> return ipFalse)
                (\ _ -> return ipTrue ) =<< try x

wrapIpF2' fun n xin new_x obj_val _userData = do
    toB $ poke obj_val =<< fun =<< ptrToVS n xin

wrapIpF2 fun n xin new_x obj_val _userData = do
    toB $ poke obj_val =<< fun =<< ptrToVS n xin

wrapIpG2 fun n xin new_x m gout _userData = do
    toB $ join $ liftM2 VM.copy (ptrToVS m gout) (fun =<< ptrToVS n xin)

wrapIpGradF2 fun n x new_x grad_f _userData = do
    toB $ join $ liftM2 VM.copy (ptrToVS n grad_f) (fun =<< ptrToVS n x)

wrapIpJacG2 fun1 fun2 n x new_x m nj iRow jCol jacs _userData
    | jacs == nullPtr = do
            toB $ join $ liftM2 fun1 (ptrToVS nj iRow) (ptrToVS nj jCol)
    | otherwise = do
            toB $ join $ liftM2 fun2 (ptrToVS n x) (ptrToVS nj jacs)

wrapIpH2 funSparsity funEval n x new_x obj_factor m lambda new_lambda nHess iRow jCol values _userData
    | iRow == nullPtr = do
            toB $ join $ liftM3 (funEval obj_factor)
                        (ptrToVS m lambda)
                        (ptrToVS n x)
                        (ptrToVS nHess values)
    | otherwise = do
            toB $ join $ liftM2 funSparsity
                        (ptrToVS nHess iRow)
                        (ptrToVS nHess jCol)

wrapIpF f = wrapIpF1 (wrapIpF2 f)
wrapIpG f = wrapIpG1 (wrapIpG2 f)
wrapIpGradF f = wrapIpGradF1 (wrapIpGradF2 f)
wrapIpJacG f1 f2 = wrapIpJacG1 (wrapIpJacG2 f1 f2)
wrapIpH fSparsity fEval = wrapIpH1 (wrapIpH2 fSparsity fEval)


vmUnsafeWith = VM.unsafeWith

-- | Vector of numbers
type Vec = VM.IOVector IpNumber

createIpoptProblem :: Vec -> Vec -> Vec -> Vec
    -> Int -> Int -> IpF -> IpG -> IpGradF -> IpJacG -> IpH -> IO IpProblem
createIpoptProblem xL xU gL gU nJac nHess f g gradF jacG hess
    | lx <- VM.length xL,
      lx == VM.length xU,
      lg <- VM.length gL,
      lg == VM.length gU = createIpoptProblem3 lx xL xU lg gL gU nJac nHess 0 f g gradF jacG hess
   | otherwise = error "dimensions wrong!"

{#fun CreateIpoptProblem as createIpoptProblem3
        { `Int', vmUnsafeWith* `Vec', vmUnsafeWith* `Vec',
          `Int', vmUnsafeWith* `Vec', vmUnsafeWith* `Vec',
          `Int', `Int', `Int', id `IpF', id `IpG', id `IpGradF',
          id `IpJacG', id `IpH' } -> `IpProblem' IpProblem #}
_ = {#fun AddIpoptNumOption as ^
    { unIpProblem `IpProblem', `String', `Double' } -> `Bool' #}

_ = {#fun AddIpoptStrOption as ^
    { unIpProblem `IpProblem', `String', `String' } -> `Bool' #}

_ = {#fun AddIpoptIntOption as ^
    { unIpProblem `IpProblem', `String', `Int' } -> `Bool' #}

_ = {#fun FreeIpoptProblem as ^
    { unIpProblem `IpProblem' } -> `()' #}

_ = {#fun OpenIpoptOutputFile as ^
    { unIpProblem `IpProblem', `String', `Int' } -> `Bool' #}

_ = {#fun SetIpoptProblemScaling as ^
    { unIpProblem `IpProblem',
    `Double',
     vmUnsafeWith* `Vec',
     vmUnsafeWith* `Vec'
     } -> `Bool' #}


data IpOptSolved = IpOptSolved
  { ipOptSolved_status :: ApplicationReturnStatus,
    ipOptSolved_objective :: Double,
    ipOptSolved_x,
    ipOptSolved_g,
    ipOptSolved_mult_g,
    ipOptSolved_mult_x_L,
    ipOptSolved_mult_x_U :: Vec }

ipoptSolve :: IpProblem
    -> Vec -- ^ starting point @x@. Note that the value is overwritte with the final @x@.
    -> IO IpOptSolved
ipoptSolve problem x = do
    g <- VM.new (VM.length x)
    mult_g <- VM.new (VM.length x)
    mult_x_L <- VM.new (VM.length x)
    mult_x_U <- VM.new (VM.length x)

    out <- ipoptSolve2
     problem
     x
     g
     mult_g
     mult_x_L
     mult_x_U
     nullPtr
    return $ IpOptSolved
      (fst out)
      (snd out)
      x
      g
      mult_g
      mult_x_L
      mult_x_U

_ = {#fun IpoptSolve as ipoptSolve2
        { unIpProblem `IpProblem',
        vmUnsafeWith* `Vec',
        vmUnsafeWith* `Vec',
        alloca- `Double' peekFloatConv*,
        vmUnsafeWith* `Vec',
        vmUnsafeWith* `Vec',
        vmUnsafeWith* `Vec',
        id `Ptr ()' } -> `ApplicationReturnStatus' cToEnum #}


{- | Set-up an 'IpProblem' to be solved later. Only objective function (@f@)
and constraint functions (@g@) need to be specified. Derivatives needed by ipopt
are computed by "Numeric.AD".

To solve the optimization problem:

>              min f(x)
>   such that
>           xL <=  x     <= xU
>           gL <=  g(x)  <= gU

First create an opaque 'IpProblem' object (nlp):

> nlp <- createIpOptProblemAD xL xU gL gU f g

Then pass it off to 'ipoptSolve'.

> ipoptSolve nlp x0

Refer to the example @Test2.hs@ for details of setting up the vectors supplied.
-}
createIpoptProblemAD
    :: Vec -- ^ @xL@ 'VM.Vector' of lower bounds for decision variables with length @n@
    -> Vec -- ^ @xU@ 'VM.Vector' of upper bounds for decision variables
    -> Vec -- ^ @gL@ 'VM.Vector' of lower bounds for constraint functions @g(x)@ with length @m@
    -> Vec -- ^ @gU@ 'VM.Vector' of upper bounds for constraint functions @g(x)@
    -> (forall a. Num a => V.Vector a -> a) -- ^ objective function @f : R^n -> R@
    -> (forall a. Num a => V.Vector a -> V.Vector a) -- ^ constraint functions @g : R^n -> R^m@
    -> IO IpProblem
createIpoptProblemAD xL xU gL gU f g
    | n <- VM.length xL,
      n == VM.length xU,
      m <- VM.length gL,
      m == VM.length gU = do
    (eval_f, eval_grad_f, eval_g, eval_jac_g, eval_h) <- mkFs n m f g
    createIpoptProblem xL xU gL gU (n*m) (((n+1)*n) `div` 2)
            eval_f eval_g eval_grad_f eval_jac_g eval_h

mkFs :: Int -- ^ @n@ number of variables
    -> Int -- ^ @m@ number of constraints
    -> (forall a. Num a => V.Vector a -> a) -- ^ objective function @R^n -> R@
    -> (forall a. Num a => V.Vector a -> V.Vector a) -- ^ constraint functions @R^n -> R^m@
    -> IO (IpF, IpGradF, IpG, IpJacG, IpH)
mkFs n m f g = do
  ipF <- wrapIpF $ \x -> do
    x <- VG.convert `fmap` VS.unsafeFreeze x
    return $ f x

  ipGradF <- wrapIpGradF $ \x -> do
    x <- VG.convert `fmap` VS.unsafeFreeze x
    VS.unsafeThaw $ VG.convert (grad f x)

  ipG <- wrapIpG $ \x -> do
    x <- VG.convert `fmap` VS.unsafeFreeze x
    VS.unsafeThaw $ VG.convert (g x)

  ipJacG <- wrapIpJacG (denseIJ n m) $ \x y -> do
    x <- VG.convert `fmap` VS.unsafeFreeze x
    jac <- VS.unsafeThaw $ VG.convert $ VG.concat $ VG.toList $ jacobian g x
    VM.copy y jac

  ipH <- wrapIpH (denseIJh n m)
     ( \ obj_factor lambda x values -> do

        x <- VG.convert `fmap` VS.unsafeFreeze x
        lambda <- VG.convert `fmap` VS.unsafeFreeze lambda

        let tri = VG.concat . VG.toList . V.imap (\n -> V.take (n+1))
            obj = V.map (*obj_factor) $ tri $ hessian f x
            gj = V.zipWith (\l v -> V.map (l*) v) lambda (V.map tri (hessianF g x))
            lagrangian = V.foldl (V.zipWith (+)) obj gj

        VM.copy values =<< VS.unsafeThaw (VG.convert lagrangian)
        )

  return (ipF, ipGradF, ipG, ipJacG, ipH)


-- | indexes the same as http://www.coin-or.org/Ipopt/documentation/node40.html
denseIJ n m iRow jCol = do
    VM.copy iRow =<< VS.unsafeThaw (VS.generate (n*m) (\x -> fromIntegral $ x `div` n))
    VM.copy jCol =<< VS.unsafeThaw (VS.generate (n*m) (\x -> fromIntegral $ x `mod` n))

-- | indexes the same as http://www.coin-or.org/Ipopt/documentation/node41.html
denseIJh n m iRow jCol = do
    i <- newIORef 0
    forM_ [0 .. fromIntegral n-1] $ \ row ->
        forM_ [ 0 .. row ] $ \col -> do
            ii <- readIORef i
            VM.write iRow ii row
            VM.write jCol ii col
            writeIORef i (ii+1)