{-# LANGUAGE TemplateHaskell #-}
module Main (main) where

import Control.Monad
import Data.List
import Data.Ratio
import Data.VectorSpace
import Test.HUnit hiding (Test)
import Test.Framework (Test, defaultMain, testGroup)
import Test.Framework.TH
import Test.Framework.Providers.HUnit
import Text.Printf

import qualified ToySolver.Data.LA as LA
import ToySolver.Simplex2

case_test1 :: IO ()
case_test1 = do
  solver <- newSolver
  x <- newVar solver
  y <- newVar solver
  z <- newVar solver
  assertAtom solver (LA.fromTerms [(7,x), (12,y), (31,z)] .==. LA.constant 17)
  assertAtom solver (LA.fromTerms [(3,x), (5,y), (14,z)]  .==. LA.constant 7)
  assertAtom solver (LA.var x .>=. LA.constant 1)
  assertAtom solver (LA.var x .<=. LA.constant 40)
  assertAtom solver (LA.var y .>=. LA.constant (-50))
  assertAtom solver (LA.var y .<=. LA.constant 50)

  ret <- check solver
  ret @?= True

  vx <- getValue solver x
  vy <- getValue solver y
  vz <- getValue solver z
  7*vx + 12*vy + 31*vz @?= 17
  3*vx + 5*vy + 14*vz @?= 7
  assertBool (printf "vx should be >=1 but %s"   (show vx)) $ vx >= 1
  assertBool (printf "vx should be <=40 but %s"  (show vx)) $ vx <= 40
  assertBool (printf "vx should be >=-50 but %s" (show vy)) $ vy >= -50
  assertBool (printf "vx should be <=50 but %s"  (show vy)) $ vy <= 50

case_test2 :: IO ()
case_test2 = do
  solver <- newSolver
  x <- newVar solver
  y <- newVar solver
  assertAtom solver (LA.fromTerms [(11,x), (13,y)] .>=. LA.constant 27)
  assertAtom solver (LA.fromTerms [(11,x), (13,y)] .<=. LA.constant 45)
  assertAtom solver (LA.fromTerms [(7,x), (-9,y)] .>=. LA.constant (-10))
  assertAtom solver (LA.fromTerms [(7,x), (-9,y)] .<=. LA.constant 4)

  ret <- check solver
  ret @?= True

  vx <- getValue solver x
  vy <- getValue solver y
  let v1 = 11*vx + 13*vy
      v2 = 7*vx - 9*vy
  assertBool (printf "11*vx + 13*vy should be >=27 but %s" (show v1)) $ 27 <= v1
  assertBool (printf "11*vx + 13*vy should be <=45 but %s" (show v1)) $ v1 <= 45
  assertBool (printf "7*vx - 9*vy should be >=-10 but %s" (show v2)) $ -10 <= v2
  assertBool (printf "7*vx - 9*vy should be >=-10 but %s" (show v2)) $ v2 <= 4


{-
Minimize
 obj: - x1 - 2 x2 - 3 x3 - x4
Subject To
 c1: - x1 + x2 + x3 + 10 x4 <= 20
 c2: x1 - 3 x2 + x3 <= 30
 c3: x2 - 3.5 x4 = 0
Bounds
 0 <= x1 <= 40
 2 <= x4 <= 3
End
-}
case_test3 :: IO ()
case_test3 = do
  solver <- newSolver

  _ <- newVar solver
  x1 <- newVar solver
  x2 <- newVar solver
  x3 <- newVar solver
  x4 <- newVar solver

  setObj solver (LA.fromTerms [(-1,x1), (-2,x2), (-3,x3), (-1,x4)])

  assertAtom solver (LA.fromTerms [(-1,x1), (1,x2), (1,x3), (10,x4)] .<=. LA.constant 20)
  assertAtom solver (LA.fromTerms [(1,x1), (-3,x2), (1,x3)] .<=. LA.constant 30)
  assertAtom solver (LA.fromTerms [(1,x2), (-3.5,x4)] .==. LA.constant 0)

  assertAtom solver (LA.fromTerms [(1,x1)] .>=. LA.constant 0)
  assertAtom solver (LA.fromTerms [(1,x1)] .<=. LA.constant 40)
  assertAtom solver (LA.fromTerms [(1,x2)] .>=. LA.constant 0)
  assertAtom solver (LA.fromTerms [(1,x3)] .>=. LA.constant 0)
  assertAtom solver (LA.fromTerms [(1,x4)] .>=. LA.constant 2)
  assertAtom solver (LA.fromTerms [(1,x4)] .<=. LA.constant 3)

  ret1 <- check solver
  ret1 @?= True

  ret2 <- optimize solver defaultOptions
  ret2 @?= Optimum

{-
http://www.math.cuhk.edu.hk/~wei/lpch5.pdf
example 5.7

minimize 3 x1 + 4 x2 + 5 x3
subject to 
1 x1 + 2 x2 + 3 x3 >= 5
2 x1 + 2 x2 + 1 x3 >= 6

optimal value is 11
-}
case_test6 :: IO ()
case_test6 = do
  solver <- newSolver

  _  <- newVar solver
  x1 <- newVar solver
  x2 <- newVar solver
  x3 <- newVar solver

  assertLower solver x1 0
  assertLower solver x2 0
  assertLower solver x3 0
  assertAtom solver (LA.fromTerms [(1,x1),(2,x2),(3,x3)] .>=. LA.constant 5)
  assertAtom solver (LA.fromTerms [(2,x1),(2,x2),(1,x3)] .>=. LA.constant 6)

  setObj solver (LA.fromTerms [(3,x1),(4,x2),(5,x3)])
  setOptDir solver OptMin
  b <- isOptimal solver
  assertBool "should be optimal" $ b

  ret <- dualSimplex solver defaultOptions
  ret @?= Optimum

  val <- getObjValue solver
  val @?= 11

{-
http://www.math.cuhk.edu.hk/~wei/lpch5.pdf
example 5.7

maximize -3 x1 -4 x2 -5 x3
subject to 
-1 x1 -2 x2 -3 x3 <= -5
-2 x1 -2 x2 -1 x3 <= -6

optimal value should be -11
-}
case_test7 :: IO ()
case_test7 = do
  solver <- newSolver

  _  <- newVar solver
  x1 <- newVar solver
  x2 <- newVar solver
  x3 <- newVar solver

  assertLower solver x1 0
  assertLower solver x2 0
  assertLower solver x3 0
  assertAtom solver (LA.fromTerms [(-1,x1),(-2,x2),(-3,x3)] .<=. LA.constant (-5))
  assertAtom solver (LA.fromTerms [(-2,x1),(-2,x2),(-1,x3)] .<=. LA.constant (-6))

  setObj solver (LA.fromTerms [(-3,x1),(-4,x2),(-5,x3)])
  setOptDir solver OptMax
  b <- isOptimal solver
  assertBool "should be optimal" $ b

  ret <- dualSimplex solver defaultOptions
  ret @?= Optimum

  val <- getObjValue solver
  val @?= -11

case_AssertAtom :: IO ()
case_AssertAtom = do
  solver <- newSolver
  x0 <- newVar solver
  assertAtom solver (LA.constant 1 .<=. LA.var x0)
  ret <- getLB solver x0
  ret @?= Just 1

  solver <- newSolver
  x0 <- newVar solver
  assertAtom solver (LA.var x0 .>=. LA.constant 1)
  ret <- getLB solver x0
  ret @?= Just 1

  solver <- newSolver
  x0 <- newVar solver
  assertAtom solver (LA.constant 1 .>=. LA.var x0)
  ret <- getUB solver x0
  ret @?= Just 1

  solver <- newSolver
  x0 <- newVar solver
  assertAtom solver (LA.var x0 .<=. LA.constant 1)
  ret <- getUB solver x0
  ret @?= Just 1

------------------------------------------------------------------------

case_example_3_2 = do
  solver <- newSolver
  [x1,x2,x3] <- replicateM 3 (newVar solver)
  setOptDir solver OptMax
  setObj solver $ LA.fromTerms [(3,x1), (2,x2), (3,x3)]
  mapM_ (assertAtom solver) $
    [ LA.fromTerms [(2,x1), (1,x2), (1,x3)] .<=. LA.constant 2
    , LA.fromTerms [(1,x1), (2,x2), (3,x3)] .<=. LA.constant 5
    , LA.fromTerms [(2,x1), (2,x2), (1,x3)] .<=. LA.constant 6
    , LA.var x1 .>=. LA.constant 0
    , LA.var x2 .>=. LA.constant 0
    , LA.var x3 .>=. LA.constant 0
    ]

  ret <- optimize solver defaultOptions
  ret @?= Optimum
  val <- getObjValue solver
  val @?= 27/5

  forM_ [(x1,1/5),(x2,0),(x3,8/5)] $ \(var,expected) -> do
    val <- getValue solver var
    val @?= expected

case_example_3_5 = do
  solver <- newSolver
  [x1,x2,x3,x4,x5] <- replicateM 5 (newVar solver)
  setOptDir solver OptMin
  setObj solver $ LA.fromTerms [(-2,x1), (4,x2), (7,x3), (1,x4), (5,x5)]
  mapM_ (assertAtom solver) $
    [ LA.fromTerms [(-1,x1), (1,x2), (2,x3), (1,x4), (2,x5)] .==. LA.constant 7
    , LA.fromTerms [(-1,x1), (2,x2), (3,x3), (1,x4), (1,x5)] .==. LA.constant 6
    , LA.fromTerms [(-1,x1), (1,x2), (1,x3), (2,x4), (1,x5)] .==. LA.constant 4
    , LA.var x2 .>=. LA.constant 0
    , LA.var x3 .>=. LA.constant 0
    , LA.var x4 .>=. LA.constant 0
    , LA.var x5 .>=. LA.constant 0
    ]

  ret <- optimize solver defaultOptions
  ret @?= Optimum
  val <- getObjValue solver
  val @?= 19

  forM_ [(x1,-1),(x2,0),(x3,1),(x4,0),(x5,2)] $ \(var,expected) -> do
    val <- getValue solver var
    val @?= expected

case_example_4_1 = do
  solver <- newSolver
  [x1,x2] <- replicateM 2 (newVar solver)
  setOptDir solver OptMin
  setObj solver $ LA.fromTerms [(2,x1), (1,x2)]
  mapM_ (assertAtom solver) $
    [ LA.fromTerms [(-1,x1), (1,x2)] .>=. LA.constant 2
    , LA.fromTerms [( 1,x1), (1,x2)] .<=. LA.constant 1
    , LA.var x1 .>=. LA.constant 0
    , LA.var x2 .>=. LA.constant 0
    ]
  ret <- optimize solver defaultOptions
  ret @?= Unsat

case_example_4_2 = do
  solver <- newSolver
  [x1,x2] <- replicateM 2 (newVar solver)
  setOptDir solver OptMax
  setObj solver $ LA.fromTerms [(2,x1), (1,x2)]
  mapM_ (assertAtom solver) $
    [ LA.fromTerms [(-1,x1), (-1,x2)] .<=. LA.constant 10
    , LA.fromTerms [( 2,x1), (-1,x2)] .<=. LA.constant 40
    , LA.var x1 .>=. LA.constant 0
    , LA.var x2 .>=. LA.constant 0
    ]
  ret <- optimize solver defaultOptions
  ret @?= Unbounded

case_example_4_3 = do
  solver <- newSolver
  [x1,x2] <- replicateM 2 (newVar solver)
  setOptDir solver OptMax
  setObj solver $ LA.fromTerms [(6,x1), (-2,x2)]
  mapM_ (assertAtom solver) $
    [ LA.fromTerms [(2,x1), (-1,x2)] .<=. LA.constant 2
    , LA.var x1 .<=. LA.constant 4
    , LA.var x1 .>=. LA.constant 0
    , LA.var x2 .>=. LA.constant 0
    ]

  ret <- optimize solver defaultOptions
  ret @?= Optimum
  val <- getObjValue solver
  val @?= 12

  forM_ [(x1,4),(x2,6)] $ \(var,expected) -> do
    val <- getValue solver var
    val @?= expected

case_example_4_5 = do
  solver <- newSolver
  [x1,x2] <- replicateM 2 (newVar solver)
  setOptDir solver OptMax
  setObj solver $ LA.fromTerms [(2,x1), (1,x2)]
  mapM_ (assertAtom solver) $
    [ LA.fromTerms [(4,x1), ( 3,x2)] .<=. LA.constant 12
    , LA.fromTerms [(4,x1), ( 1,x2)] .<=. LA.constant 8
    , LA.fromTerms [(4,x1), (-1,x2)] .<=. LA.constant 8
    , LA.var x1 .>=. LA.constant 0
    , LA.var x2 .>=. LA.constant 0
    ]

  ret <- optimize solver defaultOptions
  ret @?= Optimum
  val <- getObjValue solver
  val @?= 5

  forM_ [(x1,3/2),(x2,2)] $ \(var,expected) -> do
    val <- getValue solver var
    val @?= expected

case_example_4_6 = do
  solver <- newSolver
  [x1,x2,x3,x4] <- replicateM 4 (newVar solver)
  setOptDir solver OptMax
  setObj solver $ LA.fromTerms [(20,x1), (1/2,x2), (-6,x3), (3/4,x4)]
  mapM_ (assertAtom solver) $
    [ LA.var x1 .<=. LA.constant 2
    , LA.fromTerms [( 8,x1), (  -1,x2), (9,x3), (1/4, x4)] .<=. LA.constant 16
    , LA.fromTerms [(12,x1), (-1/2,x2), (3,x3), (1/2, x4)] .<=. LA.constant 24
    , LA.var x2 .<=. LA.constant 1
    , LA.var x1 .>=. LA.constant 0
    , LA.var x2 .>=. LA.constant 0
    , LA.var x3 .>=. LA.constant 0
    , LA.var x4 .>=. LA.constant 0
    ]

  ret <- optimize solver defaultOptions
  ret @?= Optimum
  val <- getObjValue solver
  val @?= 165/4

  forM_ [(x1,2),(x2,1),(x3,0),(x4,1)] $ \(var,expected) -> do
    val <- getValue solver var
    val @?= expected

case_example_4_7 = do
  solver <- newSolver
  [x1,x2,x3,x4] <- replicateM 4 (newVar solver)
  setOptDir solver OptMax
  setObj solver $ LA.fromTerms [(1,x1), (1.5,x2), (5,x3), (2,x4)]
  mapM_ (assertAtom solver) $
    [ LA.fromTerms [(3,x1), (2,x2), ( 1,x3), (4,x4)] .<=. LA.constant 6
    , LA.fromTerms [(2,x1), (1,x2), ( 5,x3), (1,x4)] .<=. LA.constant 4
    , LA.fromTerms [(2,x1), (6,x2), (-4,x3), (8,x4)] .==. LA.constant 0
    , LA.fromTerms [(1,x1), (3,x2), (-2,x3), (4,x4)] .==. LA.constant 0
    , LA.var x1 .>=. LA.constant 0
    , LA.var x2 .>=. LA.constant 0
    , LA.var x3 .>=. LA.constant 0
    , LA.var x4 .>=. LA.constant 0
    ]

  ret <- optimize solver defaultOptions
  ret @?= Optimum
  val <- getObjValue solver
  val @?= 48/11

  forM_ [(x1,0),(x2,0),(x3,8/11),(x4,4/11)] $ \(var,expected) -> do
    val <- getValue solver var
    val @?= expected

-- 退化して巡回の起こるKuhnの7変数3制約の例
case_kuhn_7_3 = do
  solver <- newSolver
  [x1,x2,x3,x4,x5,x6,x7] <- replicateM 7 (newVar solver)
  setOptDir solver OptMin
  setObj solver $ LA.fromTerms [(-2,x4),(-3,x5),(1,x6),(12,x7)]
  mapM_ (assertAtom solver) $
    [ LA.fromTerms [(1,x1), ( -2,x4), (-9,x5), (   1,x6), (  9,x7)] .==. LA.constant 0
    , LA.fromTerms [(1,x2), (1/3,x4), ( 1,x5), (-1/3,x6), ( -2,x7)] .==. LA.constant 0
    , LA.fromTerms [(1,x3), (  2,x4), ( 3,x5), (  -1,x6), (-12,x7)] .==. LA.constant 2
    , LA.var x1 .>=. LA.constant 0
    , LA.var x2 .>=. LA.constant 0
    , LA.var x3 .>=. LA.constant 0
    , LA.var x4 .>=. LA.constant 0
    , LA.var x5 .>=. LA.constant 0
    , LA.var x6 .>=. LA.constant 0
    , LA.var x7 .>=. LA.constant 0
    ]

  ret <- optimize solver defaultOptions
  ret @?= Optimum
  val <- getObjValue solver
  val @?= -2

  forM_ [(x1,2),(x2,0),(x3,0),(x4,2),(x5,0),(x6,2),(x7,0)] $ \(var,expected) -> do
    val <- getValue solver var
    val @?= expected

------------------------------------------------------------------------
-- Test harness

main :: IO ()
main = $(defaultMainGenerator)