{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Numeric.HiGHS.LP (
   -- * simple solver
   solve,
   LP.Direction(..),
   Term(..), (LP..*),
   Constraints,
   LP.free, (LP.<=.), (LP.>=.), (LP.==.), (LP.>=<.),
   Method, Priv.simplex, Priv.choose, Priv.ipm,
   LPEnum.ModelStatus,
   Result,

   -- * solve with extra queries on the result
   solveWith,
   Query, Priv.getObjectiveValue,
   Priv.getOptimalVector, Priv.getSolutionVectors,
   Priv.getBasisStatus, Highs.BasisStatus,
   ) where

-- ToDo: extend .* to functions of indices using a type class
import qualified Numeric.HiGHS.LP.Enumeration as LPEnum
import qualified Numeric.HiGHS.LP.FFI as Highs
import qualified Numeric.HiGHS.LP.Debug as Debug
import qualified Numeric.HiGHS.LP.Private as Priv
import Numeric.HiGHS.LP.Private
         (Method, Result, Query, checkStatus, runContT, withBuffer,
          storeBounds, prepareRowBoundsArrays, prepareColumnBoundsArrays,
          storeConstraints, prepareConstraints,
          setMethod, objectiveSense, examineStatus)

import qualified Numeric.LinearProgramming.Common as LP
import Numeric.LinearProgramming.Common
         (Bounds, Term(Term), Constraints, Direction(..), Objective)

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape

import Control.Monad.IO.Class (liftIO)
import Control.Exception (bracket)

import System.IO.Unsafe (unsafePerformIO)



{- $setup
>>> import qualified Numeric.HiGHS.LP as LP
>>> import qualified Numeric.LinearProgramming.Test as TestLP
>>> import Numeric.HiGHS.LP ((.*), (==.), (<=.), (>=.), (>=<.))
>>>
>>> import qualified Data.Array.Comfort.Storable as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>> import Data.Tuple.HT (mapPair, mapSnd)
>>>
>>> import Control.Applicative (liftA2)
>>>
>>> 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
>>>
>>> approxReal :: (Ord a, Num a) => a -> a -> a -> Bool
>>> approxReal tol x y = abs (x-y) <= tol
>>>
>>> forAllMethod ::
>>>    (QC.Testable prop) => (LP.Method -> prop) -> QC.Property
>>> forAllMethod = QC.forAll (QC.elements [LP.simplex, LP.choose, LP.ipm])
-}



{- |
>>> :{
   case Shape.indexTupleFromShape tripletShape of
      (x,y,z) ->
         fmap (mapSnd Array.toTuple) $ snd $
         LP.solve LP.simplex []
            [[2.*x, 1.*y] <=. 10, [1.*y, (5::Double).*z] <=. 20]
            (LP.Maximize, Array.fromTuple (4,-3,2)
               :: Array.Array TripletShape Double)
:}
Just (28.0,(5.0,0.0,4.0))

>>> :{
   case Shape.indexTupleFromShape tripletShape of
      (x,y,z) ->
         fmap (mapSnd Array.toTuple) $ snd $
         LP.solve LP.choose [y >=<. (-12,12)]
            [[1.*x, (-1).*y] <=. 10, [(-1).*y, (1::Double).*z] <=. 20]
            (LP.Maximize, Array.fromTuple (4,-3,2)
               :: Array.Array TripletShape Double)
:}
Just (116.0,(22.0,12.0,32.0))

>>> :{
   case Shape.indexTupleFromShape tripletShape of
      (x,y,z) ->
         mapSnd (fmap (mapSnd Array.toTuple)) $
         LP.solve LP.choose [y >=<. (-12,12)]
            [[1.*x, 1.*y] <=. 10, [1.*y, (-1::Double).*z] >=. 20]
            (LP.Maximize, Array.fromTuple (4,3,2)
               :: Array.Array TripletShape Double)
:}
(ModelStatusInfeasible,...)

>>> :{
   case Shape.indexTupleFromShape tripletShape of
      (x,y,z) ->
         mapSnd (fmap (mapSnd Array.toTuple)) $
         LP.solve LP.choose [y >=<. (-12,12)]
            [[1.*x, 1.*y] <=. 10, [1.*y, (1::Double).*z] >=. 20]
            (LP.Maximize, Array.fromTuple (4,3,2)
               :: Array.Array TripletShape Double)
:}
(ModelStatusUnbounded,...)

prop> :{
   forAllMethod $ \method
      (QC.Positive posWeight) (QC.Positive negWeight) target ->
   case Shape.indexTupleFromShape pairShape of
      (pos,neg) ->
         case fmap (mapSnd Array.toTuple) $ snd $ LP.solve method []
                  [[1.*pos, (-1::Double).*neg] ==. target]
                  (LP.Minimize,
                   Array.fromTuple (posWeight,negWeight)
                     :: Array.Array PairShape Double) of
            Nothing -> QC.property False
            Just (absol,(posResult,negResult)) ->
               QC.property (absol>=0)
               .&&.
               (posResult === 0 .||. negResult === 0)
:}
prop> :{
   forAllMethod $ \method target ->
   case Shape.indexTupleFromShape pairShape of
      (pos,neg) ->
         case fmap (mapSnd Array.toTuple) $ snd $ LP.solve method []
                  [[1.*pos, (-1::Double).*neg] ==. target]
                  (LP.Minimize, Array.fromTuple (1,1)
                     :: Array.Array PairShape Double) of
            Nothing -> QC.property False
            Just (absol,(posResult,negResult)) ->
               QC.counterexample (show(absol,(posResult,negResult))) $
               QC.property (approxReal 0.001 absol (abs target))
               .&&.
               (posResult === 0 .||. negResult === 0)
:}

prop> :{
   forAllMethod $ \method ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllProblem origin $ \bounds constrs ->
   QC.forAll (TestLP.genObjective origin) $ \(dir,obj) ->
   case snd $ LP.solve method bounds constrs (dir,obj) of
      Nothing -> False
      Just _ -> True
:}
prop> :{
   forAllMethod $ \method ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllProblem origin $ \bounds constrs ->
   QC.forAll (TestLP.genObjective origin) $ \(dir,obj) ->
   case snd $ LP.solve method bounds constrs (dir,obj) of
      Nothing -> QC.property False
      Just (_,sol) -> TestLP.checkFeasibility 0.1 bounds constrs sol
:}
prop> :{
   forAllMethod $ \method ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllProblem origin $ \bounds constrs ->
   QC.forAll (TestLP.genObjective origin) $ \(dir,obj) ->
   case snd $ LP.solve method bounds constrs (dir,obj) of
      Nothing -> QC.property False
      Just (_,sol) ->
         QC.forAll (QC.choose (0,1)) $ \lambda ->
         TestLP.checkFeasibility 0.1 bounds constrs $
         TestLP.affineCombination lambda sol (Array.map fromIntegral origin)
:}
prop> :{
   forAllMethod $ \method ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllProblem origin $ \bounds constrs ->
   QC.forAll (TestLP.genObjective origin) $ \(dir,obj) ->
   case snd $ LP.solve method bounds constrs (dir,obj) of
      Nothing -> QC.property False
      Just (opt,sol) ->
         QC.forAll (QC.choose (0,1)) $ \lambda ->
            let val = TestLP.scalarProduct obj $
                        TestLP.affineCombination lambda sol $
                        Array.map fromIntegral origin
            in QC.counterexample (show (dir,opt,val)) $
               case dir of
                  LP.Minimize -> opt-0.01 <= val
                  LP.Maximize -> opt+0.01 >= val
:}
prop> :{
   forAllMethod $ \method ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllBoundedProblem origin $ \bounds constrs ->
   QC.forAll (TestLP.genObjective origin) $ \dirObjA ->
   QC.forAll (TestLP.genObjective origin) $ \dirObjB ->
   let solA = snd $ LP.solve method bounds constrs dirObjA in
   let solB = snd $ LP.solve method bounds constrs dirObjB in
   QC.counterexample (show (fmap fst solA, fmap fst solB)) $
   case (solA, solB) of
      (Just _, Just _) -> True
      (Nothing, Nothing) -> True
      _ -> False
:}
prop> :{
   forAllMethod $ \method ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllProblem origin $ \bounds constrs ->
   QC.forAll (TestLP.genObjective origin) $ \(_dir,obj) ->
   case (snd $ LP.solve method bounds constrs (LP.Minimize,obj),
         snd $ LP.solve method bounds constrs (LP.Maximize,obj)) of
      (Just (optMin,_), Just (optMax,_)) ->
         QC.counterexample (show (optMin, optMax)) $ optMin <= optMax + 0.01
      _ -> QC.property False
:}
prop> :{
   forAllMethod $ \method ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllProblem origin $ \bounds allConstrs ->
   QC.forAll (QC.sublistOf allConstrs) $ \someConstrs ->
   QC.forAll (TestLP.genObjective origin) $ \(dir,obj) ->
   case (snd $ LP.solve method bounds allConstrs (dir,obj),
         snd $ LP.solve method bounds someConstrs (dir,obj)) of
      (Just (optAll,_), Just (optSome,_)) ->
         QC.counterexample (show (optAll, optSome)) $
         case dir of
            LP.Minimize -> optAll >= optSome-0.01
            LP.Maximize -> optAll <= optSome+0.01
      _ -> QC.property False
:}
prop> :{
   forAllMethod $ \methodA ->
   forAllMethod $ \methodB ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllProblem origin $ \bounds constrs ->
   QC.forAll (TestLP.genObjective origin) $ \dirObj ->
   case (snd $ LP.solve methodA bounds constrs dirObj,
         snd $ LP.solve methodB bounds constrs dirObj) of
      (Just (optA,_), Just (optB,_)) ->
         QC.counterexample (show (optA, optB)) $
         approxReal 0.01 optA optB
      _ -> QC.property False
:}
-}
solve ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Method -> Bounds ix -> Constraints Double ix ->
   (Direction, Objective sh) -> Result sh
solve = solveWith Priv.getResult

{- |
>>> :{
   case Shape.indexTupleFromShape tripletShape of
      (x,y,z) ->
         fmap (mapSnd (mapPair (Array.toTuple, Array.toList))) $ snd $
         LP.solveWith
            (liftA2 (,) LP.getObjectiveValue LP.getBasisStatus)
            LP.simplex []
            [[2.*x, 1.*y] <=. 10, [1.*y, (5::Double).*z] <=. 20]
            (LP.Maximize, Array.fromTuple (4,-3,2) :: Array.Array TripletShape Double)
:}
Just (28.0,((Highs.basisStatusBasic,Highs.basisStatusLower,Highs.basisStatusBasic),[Highs.basisStatusUpper,Highs.basisStatusUpper]))

prop> :{
   forAllMethod $ \method ->
   TestLP.forAllOrigin $ \origin ->
   TestLP.forAllProblem origin $ \bounds constrs ->
   QC.forAll (TestLP.genObjective origin) $ \(dir,obj) ->
   case (snd $ LP.solve method bounds constrs (dir,obj),
         snd $ LP.solveWith LP.getSolutionVectors
                  method bounds constrs (dir,obj)) of
      (Just (_,sol0), Just ((sol1,_),_)) -> sol0 == sol1
      _ -> False
:}
-}
solveWith ::
   (Shape.Indexed sh, Shape.Index sh ~ ix) =>
   Query sh result -> Method -> Bounds ix -> Constraints Double ix ->
   (Direction, Objective sh) -> (LPEnum.ModelStatus, Maybe result)
solveWith query method bounds constrs (dir,obj) =
   unsafePerformIO $
   let shape = Array.shape obj in
   let numCols = Shape.size shape in
   let numRows = length constrs in

   runContT $ do
   objPtr <- withBuffer $ Array.map realToFrac obj
   (collbPtr,colubPtr) <-
      storeBounds $ prepareColumnBoundsArrays shape bounds
   (rowlbPtr,rowubPtr) <- storeBounds $ prepareRowBoundsArrays constrs

   let (coefficients, indices, rowStarts) = prepareConstraints shape constrs
   (coefficientsPtr, indexPtr, startPtr)
      <- storeConstraints (coefficients, indices, rowStarts)

   liftIO $ bracket Highs.create Highs.destroy $ \model -> do
      Debug.initLog model
      setMethod model method
      checkStatus $ Highs.passLp model
         (fromIntegral numCols)
         (fromIntegral numRows)
         (fromIntegral $ Shape.size $ Array.shape coefficients)
         Highs.matrixFormatRowwise
         (objectiveSense dir)
         0 objPtr
         collbPtr colubPtr
         rowlbPtr rowubPtr
         startPtr indexPtr coefficientsPtr

      examineStatus query shape model =<< Highs.run model