{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Numeric.HiGHS.LP.Private where

import qualified Numeric.HiGHS.LP.FFI as Highs
import Numeric.HiGHS.LP.Enumeration (ModelStatus, modelStatusFromC)
import Numeric.HiGHS.LP.FFI (Highs, HighsInt)
import Numeric.LinearProgramming.Common
         (Term(..), Bound(..), Inequality(Inequality),
          Bounds, Constraints, Direction(..))

import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Storable.Unchecked.Monadic as ArrayMonadic
import qualified Data.Array.Comfort.Storable.Unchecked as ArrayUnchecked
import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.Traversable as Trav
import Data.Array.Comfort.Storable (Array)
import Data.Foldable (for_)
import Data.Tuple.HT (mapPair)

import qualified Control.Monad.Trans.Reader as MR
import qualified Control.Monad.Trans.Cont as MC
import qualified Control.Applicative.HT as AppHT
import qualified Control.Functor.HT as FuncHT
import Control.Monad (guard, when)
import Control.Applicative (liftA2, liftA3)

import qualified Foreign.Marshal.Array.Guarded as ForeignArray
import Foreign.Storable (pokeElemOff, peekElemOff)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr)
import Foreign.C.String (withCString)
import Foreign.C.Types (CDouble)



withBuffer :: Array sh a -> MC.ContT r IO (Ptr a)
withBuffer arr =
   MC.ContT $ withForeignPtr (ArrayUnchecked.buffer arr)

runContT :: MC.ContT a IO a -> IO a
runContT act = MC.runContT act return



positiveInfinity, negativeInfinity :: CDouble
positiveInfinity =  1/0
negativeInfinity = -1/0

prepareBounds :: Inequality a -> (a, (CDouble, CDouble))
prepareBounds (Inequality x bnd) =
   (,) x $
   case bnd of
      LessEqual up    -> (negativeInfinity, realToFrac up)
      GreaterEqual lo -> (realToFrac lo,    positiveInfinity)
      Between lo up   -> (realToFrac lo,    realToFrac up)
      Equal y         -> (realToFrac y,     realToFrac y)
      Free            -> (negativeInfinity, positiveInfinity)

prepareColumnBoundsArrays ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   sh -> Bounds ix -> (Array sh CDouble, Array sh CDouble)
prepareColumnBoundsArrays shape =
   mapPair (Array.fromBoxed, Array.fromBoxed) .
   FuncHT.unzip .
   BoxedArray.fromAssociations (0, positiveInfinity) shape .
   map prepareBounds


type ShapeInt = Shape.ZeroBased Int

prepareRowBoundsArrays ::
   Bounds ix -> (Array ShapeInt CDouble, Array ShapeInt CDouble)
prepareRowBoundsArrays constrs =
   let shape = Shape.ZeroBased $ length constrs in
   mapPair (Array.fromList shape, Array.fromList shape) $
   unzip $ map (snd . prepareBounds) constrs

storeBounds ::
   (Array sh CDouble, Array sh CDouble) ->
   MC.ContT r IO (Ptr CDouble, Ptr CDouble)
storeBounds = AppHT.mapPair (withBuffer, withBuffer)


prepareConstraints ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   sh -> Constraints Double ix ->
   (Array ShapeInt CDouble, Array ShapeInt HighsInt, Array ShapeInt HighsInt)
prepareConstraints shape constrs =
   let {-
       Highs.passLp returns Warning when there are zero coefficients.
       I think zero coefficients are reasonably ok.
       -}
       constrsNonZero = map (fmap (filter (\(Term c _x) -> c/=0))) constrs
       rowStarts =
         Array.vectorFromList $ scanl (+) 0 $
         map (\(Inequality terms _bnd) -> fromIntegral $ length terms)
            constrsNonZero
       shapeOffset = Shape.offset shape
       coefficients =
         concatMap (\(Inequality terms _bnd) -> terms) constrsNonZero
       indexArr =
         Array.vectorFromList $
         map (\(Term _ ix) -> fromIntegral $ shapeOffset ix) coefficients
       coefficientArr =
         Array.vectorFromList $
         map (\(Term c _) -> realToFrac c) coefficients
   in (coefficientArr, indexArr, rowStarts)

storeConstraints ::
   (Array ShapeInt CDouble, Array ShapeInt HighsInt, Array ShapeInt HighsInt) ->
   MC.ContT r IO (Ptr CDouble, Ptr HighsInt, Ptr HighsInt)
storeConstraints (coefficients, indices, rowStarts) =
   liftA3 (,,)
      (withBuffer coefficients)
      (withBuffer indices)
      (withBuffer rowStarts)


objectiveSense :: Direction -> Highs.ObjSense
objectiveSense dir =
   case dir of
      Minimize -> Highs.objSenseMinimize
      Maximize -> Highs.objSenseMaximize


setBoolOptionValue :: Ptr Highs -> String -> Highs.Bool -> IO ()
setBoolOptionValue model key b =
   checkStatus $ withCString key $ \cstr ->
      Highs.setBoolOptionValue model cstr b

setIntOptionValue :: Ptr Highs -> String -> HighsInt -> IO ()
setIntOptionValue model key n =
   checkStatus $ withCString key $ \cstr ->
      Highs.setIntOptionValue model cstr n

setDoubleOptionValue :: Ptr Highs -> String -> CDouble -> IO ()
setDoubleOptionValue model key x =
   checkStatus $ withCString key $ \cstr ->
      Highs.setDoubleOptionValue model cstr x

setStringOptionValue :: Ptr Highs -> String -> String -> IO ()
setStringOptionValue model key value =
   checkStatus $
      withCString key $ \keyPtr ->
      withCString value $ \valuePtr ->
      Highs.setStringOptionValue model keyPtr valuePtr



newtype Method = Method {deMethod :: String}
   deriving (Show)

simplex, choose, ipm :: Method
simplex = Method "simplex"
choose = Method "choose"
ipm = Method "ipm"

setMethod :: Ptr Highs -> Method -> IO ()
setMethod model (Method method) =
   setStringOptionValue model "solver" method

checkStatus :: IO Highs.Status -> IO ()
checkStatus act = do
   status <- act
   when (status == Highs.statusError) $ fail "Highs function failed"


newtype Query sh a = Query (MR.ReaderT (sh, Ptr Highs) IO a)
   deriving (Functor, Applicative, Monad)

getResult :: (Shape.C sh) => Query sh (Double, Array sh Double)
getResult = liftA2 (,) getObjectiveValue getOptimalVector

getObjectiveValue :: Query sh Double
getObjectiveValue =
   Query $ MR.ReaderT $ fmap realToFrac . Highs.getObjectiveValue . snd

doubleFromCDoubleBuffer :: Int -> Ptr CDouble -> Ptr Double -> IO ()
doubleFromCDoubleBuffer n srcPtr dstPtr =
   for_ (take n [0..]) $ \k ->
      pokeElemOff dstPtr k . realToFrac =<< peekElemOff srcPtr k

getOptimalVector :: (Shape.C sh) => Query sh (Array sh Double)
getOptimalVector =
   Query $ MR.ReaderT $ \(shape,model) ->
   ArrayMonadic.unsafeCreateWithSize shape $ \numCols arrPtr ->
      (fmap fromIntegral (Highs.getNumRow model) >>= ) $ \numRows ->
      ForeignArray.alloca numCols $ \colValuePtr ->
      ForeignArray.alloca numCols $ \colDualPtr  ->
      ForeignArray.alloca numRows $ \rowValuePtr ->
      ForeignArray.alloca numRows $ \rowDualPtr  -> do

      checkStatus $
         Highs.getSolution model colValuePtr colDualPtr rowValuePtr rowDualPtr

      doubleFromCDoubleBuffer numCols colValuePtr arrPtr

getSolutionVectors ::
   (Shape.C sh) =>
   Query sh
      ((Array sh Double, Array sh Double),
       (Array ShapeInt Double, Array ShapeInt Double))
getSolutionVectors =
   Query $ MR.ReaderT $ \(shape,model) ->
   fmap (\(colPrimal,(colDual,row2)) -> ((colPrimal,colDual),row2)) $
   ArrayMonadic.unsafeCreateWithSizeAndResult shape $ \numCols colValueArr ->
   ArrayMonadic.unsafeCreateWithSizeAndResult shape $ \_numCols colDualArr ->
   (fmap fromIntegral (Highs.getNumRow model) >>= ) $ \numRows ->
   let constraintsShape = Shape.ZeroBased numRows in
   ArrayMonadic.unsafeCreateWithSizeAndResult constraintsShape $
      \_numRows rowValueArr ->
   ArrayMonadic.unsafeCreate constraintsShape $ \rowDualArr ->
      ForeignArray.alloca numCols $ \colValuePtr ->
      ForeignArray.alloca numCols $ \colDualPtr  ->
      ForeignArray.alloca numRows $ \rowValuePtr ->
      ForeignArray.alloca numRows $ \rowDualPtr  -> do

      checkStatus $
         Highs.getSolution model colValuePtr colDualPtr rowValuePtr rowDualPtr

      doubleFromCDoubleBuffer numCols colValuePtr colValueArr
      doubleFromCDoubleBuffer numCols colDualPtr colDualArr
      doubleFromCDoubleBuffer numRows rowValuePtr rowValueArr
      doubleFromCDoubleBuffer numRows rowDualPtr rowDualArr

getBasisStatus ::
   (Shape.C sh) =>
   Query sh (Array sh Highs.BasisStatus, Array ShapeInt Highs.BasisStatus)
getBasisStatus =
   Query $ MR.ReaderT $ \(shape,model) ->
   ArrayMonadic.unsafeCreateWithSizeAndResult shape $ \_numCols colStatusPtr ->
   (fmap fromIntegral (Highs.getNumRow model) >>= ) $ \numRows ->
   ArrayMonadic.unsafeCreate (Shape.ZeroBased numRows) $ \rowStatusPtr ->
      checkStatus $ Highs.getBasis model colStatusPtr rowStatusPtr



type Result sh = (ModelStatus, Maybe (Double, Array sh Double))

examineStatus :: (Shape.C sh) =>
   Query sh a -> sh -> Ptr Highs -> Highs.Status -> IO (ModelStatus, Maybe a)
examineStatus (Query query) shape model status =
   liftA2 (,) (fmap modelStatusFromC $ Highs.getModelStatus model) $
   Trav.for (guard (status /= Highs.statusError)) $ \() ->
      MR.runReaderT query (shape,model)