{-|
Module      : LPPaver.Constraint.Util
Description : Utility functions for the 'Constraint' type
Copyright   : (c) Junaid Rasheed, 2021-2022
License     : MPL
Maintainer  : jrasheed178@gmail.com
Stability   : experimental
Module defining useful utility functions for the 'Constraint' data type.
-}
{-# LANGUAGE LambdaCase #-}

module LPPaver.Constraint.Util where

import MixedTypesNumPrelude
import LPPaver.Constraint.Type
import Data.List (nub)
import qualified Linear.Simplex.Types as LT
import qualified Data.Map as M
import Data.Maybe
import Linear.Simplex.Types (PolyConstraint)

-- | Get the LHS of a 'Constraint'
constraintLeftSide :: Constraint -> [(String, Rational)]
constraintLeftSide :: Constraint -> [(String, Rational)]
constraintLeftSide (GEQ [(String, Rational)]
lhs Rational
_) = [(String, Rational)]
lhs
constraintLeftSide (LEQ [(String, Rational)]
lhs Rational
_) = [(String, Rational)]
lhs

-- | Get the LHS of a 'Constraint'
constraintRightSide :: Constraint -> Rational
constraintRightSide :: Constraint -> Rational
constraintRightSide (GEQ [(String, Rational)]
_ Rational
rhs) = Rational
rhs
constraintRightSide (LEQ [(String, Rational)]
_ Rational
rhs) = Rational
rhs

-- | Get a list of all variables in a list of 'Constraint's
constraintVars :: [Constraint] -> [String]
constraintVars :: [Constraint] -> [String]
constraintVars [Constraint]
cs = [String] -> [String]
forall a. Eq a => [a] -> [a]
nub ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ [Constraint] -> [String]
aux [Constraint]
cs
  where
    aux :: [Constraint] -> [String]
    aux :: [Constraint] -> [String]
aux [] = []
    aux (Constraint
x : [Constraint]
xs) = ((String, Rational) -> String) -> [(String, Rational)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (String, Rational) -> String
forall a b. (a, b) -> a
fst (Constraint -> [(String, Rational)]
constraintLeftSide Constraint
x) [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [Constraint] -> [String]
aux [Constraint]
xs

-- | Convert a list of 'Constraint's into a list of 'LT.PolyConstraint's.
-- The function returns a pair with the first item being a list of 'LT.PolyConstraint's and the second item being a mapping of 'String' variables (from 'Constraint's) to 'Integer' variables (for 'LT.PolyConstraint's).
constraintsToSimplexConstraints :: [Constraint] -> ([LT.PolyConstraint], M.Map String Integer)
constraintsToSimplexConstraints :: [Constraint] -> ([PolyConstraint], Map String Integer)
constraintsToSimplexConstraints [Constraint]
constraints =
  (
    (Constraint -> PolyConstraint) -> [Constraint] -> [PolyConstraint]
forall a b. (a -> b) -> [a] -> [b]
map
    (\case
      GEQ [(String, Rational)]
varsWithCoeffs Rational
rhs -> VarConstMap -> Rational -> PolyConstraint
LT.GEQ (((String, Rational) -> (Integer, Rational))
-> [(String, Rational)] -> VarConstMap
forall a b. (a -> b) -> [a] -> [b]
map (\(String
stringVar, Rational
coeff) -> (Maybe Integer -> Integer
forall a. HasCallStack => Maybe a -> a
fromJust (String -> Map String Integer -> Maybe Integer
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
stringVar Map String Integer
stringIntVarMap), Rational
coeff)) [(String, Rational)]
varsWithCoeffs) Rational
rhs
      LEQ [(String, Rational)]
varsWithCoeffs Rational
rhs -> VarConstMap -> Rational -> PolyConstraint
LT.LEQ (((String, Rational) -> (Integer, Rational))
-> [(String, Rational)] -> VarConstMap
forall a b. (a -> b) -> [a] -> [b]
map (\(String
stringVar, Rational
coeff) -> (Maybe Integer -> Integer
forall a. HasCallStack => Maybe a -> a
fromJust (String -> Map String Integer -> Maybe Integer
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
stringVar Map String Integer
stringIntVarMap), Rational
coeff)) [(String, Rational)]
varsWithCoeffs) Rational
rhs
    )
    [Constraint]
constraints,
    Map String Integer
stringIntVarMap
  )
  where
    stringVars :: [String]
stringVars = [Constraint] -> [String]
constraintVars [Constraint]
constraints
    stringIntVarMap :: Map String Integer
stringIntVarMap = [(String, Integer)] -> Map String Integer
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(String, Integer)] -> Map String Integer)
-> [(String, Integer)] -> Map String Integer
forall a b. (a -> b) -> a -> b
$ [String] -> [Integer] -> [(String, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [String]
stringVars [Integer
1..]