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

import Control.Monad
import Control.Monad.State
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import Data.List
import Data.Ratio
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.Data.LA ((.<=.))
import ToySolver.Arith.Simplex
import qualified ToySolver.Arith.LPSolver as LP

example_3_2 :: Tableau Rational
example_3_2 = IntMap.fromList
  [ (4, (IntMap.fromList [(1,2), (2,1), (3,1)], 2))
  , (5, (IntMap.fromList [(1,1), (2,2), (3,3)], 5))
  , (6, (IntMap.fromList [(1,2), (2,2), (3,1)], 6))
  , (objRowIndex, (IntMap.fromList [(1,-3), (2,-2), (3,-3)], 0))
  ]

case_example_3_2_simplex :: IO ()
case_example_3_2_simplex = do
  assertBool "simplex failed" ret
  assertBool "invalid tableau" (isValidTableau result)
  assertBool "infeasible tableau" (isFeasible result)
  assertBool "unoptimal tableau" (isOptimal OptMax result)
  currentObjValue result @?= 27/5
  where
    ret :: Bool
    result :: Tableau Rational
    (ret,result) = simplex OptMax example_3_2

case_example_3_2_primalDualSimplex :: IO ()
case_example_3_2_primalDualSimplex = do
  assertBool "simplex failed" ret
  assertBool "invalid tableau" (isValidTableau result)
  assertBool "infeasible tableau" (isFeasible result)
  assertBool "unoptimal tableau" (isOptimal OptMax result)
  currentObjValue result @?= 27/5
  where
    ret :: Bool
    result :: Tableau Rational
    (ret,result) = primalDualSimplex OptMax example_3_2

-- from http://www.math.cuhk.edu.hk/~wei/lpch5.pdf
exampe_5_3_phase1 :: Tableau Rational
exampe_5_3_phase1 = IntMap.fromList
  [ (6, (IntMap.fromList [(2,-1), (3,-1), (5,1), (6,1)], 1))
  , (7, (IntMap.fromList [(3,1), (4,-1), (5,1), (7,1)], 0))
  ]

case_exampe_5_3_phase1 :: IO ()
case_exampe_5_3_phase1 = do
  let (ret,result) = phaseI exampe_5_3_phase1 (IntSet.fromList [6,7])
  assertBool "phase1 failed" ret
  assertBool "invalid tableau" (isValidTableau result)
  assertBool "infeasible tableau" (isFeasible result)    

-- 退化して巡回の起こるKuhnの7変数3制約の例
kuhn_7_3 :: Tableau Rational
kuhn_7_3 = IntMap.fromList
  [ (1, (IntMap.fromList [(4,-2), (5,-9), (6,1), (7,9)],       0))
  , (2, (IntMap.fromList [(4,1/3), (5,1), (6,-1/3), (7,-2)],   0))
  , (3, (IntMap.fromList [(4,2), (5,3), (6,-1), (7,-12)],      2))
  , (objRowIndex, (IntMap.fromList [(4,2), (5,3), (6,-1), (7,-12)], 0))
  ]

case_kuhn_7_3 :: IO ()
case_kuhn_7_3 = do
  assertBool "simplex failed" ret
  assertBool "invalid tableau" (isValidTableau result)
  currentObjValue result @?= -2
  where
    ret :: Bool
    result :: Tableau Rational
    (ret,result) = simplex OptMin kuhn_7_3

-- case_pd_kuhn_7_3 :: IO ()
-- case_pd_kuhn_7_3 = do
--   assertBool "simplex failed" ret
--   assertBool "invalid tableau" (isValidTableau result)
--   currentObjValue result @?= -2
--   where
--     ret :: Bool
--     result :: Tableau Rational
--     (ret,result) = primalDualSimplex OptMin kuhn_7_3

-- from http://www.math.cuhk.edu.hk/~wei/lpch5.pdf
example_5_7 :: Tableau Rational
example_5_7 = IntMap.fromList
  [ (4, (IntMap.fromList [(1,-1), (2,-2), (3,-3)], -5))
  , (5, (IntMap.fromList [(1,-2), (2,-2), (3,-1)], -6))
  , (objRowIndex, (IntMap.fromList [(1,3),(2,4),(3,5)], 0))
  ]

case_example_5_7 :: IO ()
case_example_5_7 = do
  assertBool "dual simplex failed" ret
  assertBool "invalid tableau" (isValidTableau result)
  currentObjValue result @?= -11
  where
    ret :: Bool
    result :: Tableau Rational
    (ret,result) = dualSimplex OptMax example_5_7

case_pd_example_5_7 :: IO ()
case_pd_example_5_7 = do
  assertBool "dual simplex failed" ret
  assertBool "invalid tableau" (isValidTableau result)
  currentObjValue result @?= -11
  where
    ret :: Bool
    result :: Tableau Rational
    (ret,result) = primalDualSimplex OptMax example_5_7

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

case_lp_example_5_7_twoPhaseSimplex :: IO ()
case_lp_example_5_7_twoPhaseSimplex = do  
  ret @?= LP.Optimum
  oval @?= -11
  assertBool "invalid tableau" (isValidTableau tbl)
  assertBool "infeasible tableau" (isFeasible tbl)
  assertBool "non-optimal tableau" (isOptimal OptMax tbl)
  where
    oval :: Rational
    ((ret,tbl,oval),result) = flip runState (LP.emptySolver IntSet.empty) $ do
      _ <- LP.newVar
      x1 <- LP.newVar 
      x2 <- LP.newVar
      x3 <- LP.newVar
      LP.addConstraint (LA.fromTerms [(-1,x1),(-2,x2),(-3,x3)] .<=. LA.constant (-5))
      LP.addConstraint (LA.fromTerms [(-2,x1),(-2,x2),(-1,x3)] .<=. LA.constant (-6))
      let obj = LA.fromTerms [(-3,x1), (-4,x2),(-5,x3)]
      ret <- LP.twoPhaseSimplex OptMax obj
      tbl <- LP.getTableau
      m <- LP.getModel (IntSet.fromList [x1,x2,x3])
      let oval = LA.evalExpr m obj
      return (ret,tbl,oval)

case_lp_example_5_7_primalDualSimplex :: IO ()
case_lp_example_5_7_primalDualSimplex = do  
  ret @?= LP.Optimum
  oval @?= -11
  assertBool "invalid tableau" (isValidTableau tbl)
  assertBool "infeasible tableau" (isFeasible tbl)
  assertBool "non-optimal tableau" (isOptimal OptMax tbl)
  where
    oval :: Rational
    ((ret,tbl,oval),result) = flip runState (LP.emptySolver IntSet.empty) $ do
      _ <- LP.newVar
      x1 <- LP.newVar 
      x2 <- LP.newVar
      x3 <- LP.newVar
      LP.addConstraint (LA.fromTerms [(-1,x1),(-2,x2),(-3,x3)] .<=. LA.constant (-5))
      LP.addConstraint (LA.fromTerms [(-2,x1),(-2,x2),(-1,x3)] .<=. LA.constant (-6))
      let obj = LA.fromTerms [(-3,x1), (-4,x2),(-5,x3)]
      ret <- LP.primalDualSimplex OptMax obj
      tbl <- LP.getTableau
      m <- LP.getModel (IntSet.fromList [x1,x2,x3])
      let oval = LA.evalExpr m obj
      return (ret,tbl,oval)

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

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