content: {-# OPTIONS_GHC -Wall #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell #-}
  module Test.SAT.TheorySolver (satTheorySolverTestGroup) where
  
  import Control.Monad
  import Data.IORef
  import Data.Map (Map)
  import qualified Data.Map as Map
  import Data.IntMap (IntMap)
  import qualified Data.IntMap as IntMap
  import qualified Data.IntSet as IntSet
  import qualified Data.Traversable as Traversable
  
  import Test.Tasty
  import Test.Tasty.QuickCheck hiding ((.&&.), (.||.))
  import Test.Tasty.HUnit
  import Test.Tasty.TH
  import qualified Test.QuickCheck.Monadic as QM
  
  import ToySolver.Data.BoolExpr
  import ToySolver.Data.Boolean
  import qualified ToySolver.SAT as SAT
  import ToySolver.SAT.TheorySolver
  import qualified ToySolver.SAT.Encoder.Tseitin as Tseitin
  import qualified ToySolver.FileFormat.CNF as CNF
  import ToySolver.Data.OrdRel
  import qualified ToySolver.Data.LA as LA
  import qualified ToySolver.Arith.Simplex as Simplex
  import qualified ToySolver.EUF.EUFSolver as EUF
  
  import Test.SAT.Utils
  
  
  newTheorySolver :: CNF.CNF -> IO TheorySolver
  newTheorySolver cnf = do
    let nv = CNF.cnfNumVars cnf
        cs = CNF.cnfClauses cnf
    solver <- SAT.newSolver
    SAT.newVars_ solver nv
    forM_ cs $ \c -> SAT.addClause solver (SAT.unpackClause c)
  
    ref <- newIORef []
    let tsolver =
          TheorySolver
          { thAssertLit = \_ l -> do
              if abs l > nv then
                return True
              else do
                m <- readIORef ref
                case m of
                  [] -> SAT.addClause solver [l]
                  xs : xss -> writeIORef ref ((l : xs) : xss)
                return True
          , thCheck = \_ -> do
              xs <- liftM concat $ readIORef ref
              SAT.solveWith solver xs
          , thExplain = \m -> do
              case m of
                Nothing -> liftM IntSet.toList $ SAT.getFailedAssumptions solver
                Just _ -> return []
          , thPushBacktrackPoint = modifyIORef ref ([] :)
          , thPopBacktrackPoint = modifyIORef ref tail
          , thConstructModel = return ()
          }
    return tsolver
  
  prop_solveCNF_using_BooleanTheory :: Property
  prop_solveCNF_using_BooleanTheory = QM.monadicIO $ do
    cnf <- QM.pick arbitraryCNF
    let nv = CNF.cnfNumVars cnf
        nc = CNF.cnfNumClauses cnf
        cs = CNF.cnfClauses cnf
        cnf1 = cnf{ CNF.cnfClauses = [c | (i,c) <- zip [(0::Int)..] cs, i `mod` 2 == 0], CNF.cnfNumClauses = nc - (nc `div` 2) }
        cnf2 = cnf{ CNF.cnfClauses = [c | (i,c) <- zip [(0::Int)..] cs, i `mod` 2 /= 0], CNF.cnfNumClauses = nc `div` 2 }
  
    solver <- arbitrarySolver
  
    ret <- QM.run $ do
      SAT.newVars_ solver nv
  
      tsolver <- newTheorySolver cnf1
      SAT.setTheory solver tsolver
  
      forM_ (CNF.cnfClauses cnf2) $ \c -> SAT.addClause solver (SAT.unpackClause c)
      ret <- SAT.solve solver
      if ret then do
        m <- SAT.getModel solver
        return (Just m)
      else do
        return Nothing
  
    case ret of
      Just m -> QM.assert $ evalCNF m cnf
      Nothing -> do
        forM_ (allAssignments nv) $ \m -> do
          QM.assert $ not (evalCNF m cnf)
  
  case_QF_LRA :: Assertion
  case_QF_LRA = do
    satSolver <- SAT.newSolver
    lraSolver <- Simplex.newSolver
  
    tblRef <- newIORef $ Map.empty
    defsRef <- newIORef $ IntMap.empty
    let abstractLAAtom :: LA.Atom Rational -> IO SAT.Lit
        abstractLAAtom atom = do
          (v,op,rhs) <- Simplex.simplifyAtom lraSolver atom
          tbl <- readIORef tblRef
          (vLt, vEq, vGt) <-
            case Map.lookup (v,rhs) tbl of
              Just (vLt, vEq, vGt) -> return (vLt, vEq, vGt)
              Nothing -> do
                vLt <- SAT.newVar satSolver
                vEq <- SAT.newVar satSolver
                vGt <- SAT.newVar satSolver
                SAT.addClause satSolver [vLt,vEq,vGt]
                SAT.addClause satSolver [-vLt, -vEq]
                SAT.addClause satSolver [-vLt, -vGt]
                SAT.addClause satSolver [-vEq, -vGt]
                writeIORef tblRef (Map.insert (v,rhs) (vLt, vEq, vGt) tbl)
                let xs = IntMap.fromList
                         [ (vEq,  LA.var v .==. LA.constant rhs)
                         , (vLt,  LA.var v .<.  LA.constant rhs)
                         , (vGt,  LA.var v .>.  LA.constant rhs)
                         , (-vLt, LA.var v .>=. LA.constant rhs)
                         , (-vGt, LA.var v .<=. LA.constant rhs)
                         ]
                modifyIORef defsRef (IntMap.union xs)
                return (vLt, vEq, vGt)
          case op of
            Lt  -> return vLt
            Gt  -> return vGt
            Eql -> return vEq
            Le  -> return (-vGt)
            Ge  -> return (-vLt)
            NEq -> return (-vEq)
  
        abstract :: BoolExpr (Either SAT.Lit (LA.Atom Rational)) -> IO Tseitin.Formula
        abstract e = liftM (fold Tseitin.Atom) $ Traversable.mapM f e
          where
            f (Left lit) = return lit
            f (Right atom) = abstractLAAtom atom
  
    let tsolver =
          TheorySolver
          { thAssertLit = \_ l -> do
              defs <- readIORef defsRef
              case IntMap.lookup l defs of
                Nothing -> return True
                Just atom -> do
                  Simplex.assertAtomEx' lraSolver atom (Just l)
                  return True
          , thCheck = \_ -> do
              Simplex.check lraSolver
          , thExplain = \m -> do
              case m of
                Nothing -> liftM IntSet.toList $ Simplex.explain lraSolver
                Just _ -> return []
          , thPushBacktrackPoint = do
              Simplex.pushBacktrackPoint lraSolver
          , thPopBacktrackPoint = do
              Simplex.popBacktrackPoint lraSolver
          , thConstructModel = do
              return ()
          }
    SAT.setTheory satSolver tsolver
  
    enc <- Tseitin.newEncoder satSolver
    let addFormula :: BoolExpr (Either SAT.Lit (LA.Atom Rational)) -> IO ()
        addFormula c = Tseitin.addFormula enc =<< abstract c
  
    a <- SAT.newVar satSolver
    x <- Simplex.newVar lraSolver
    y <- Simplex.newVar lraSolver
  
    let le1 = LA.fromTerms [(2,x), (1/3,y)] .<=. LA.constant (-4) -- 2 x + (1/3) y <= -4
        eq2 = LA.fromTerms [(1.5,x)] .==. LA.fromTerms [(-2,x)] -- 1.5 y = -2 x
        gt3 = LA.var x .>. LA.var y -- x > y
        lt4 = LA.fromTerms [(3,x)] .<. LA.fromTerms [(-1,LA.unitVar), (1/5,x), (1/5,y)] -- 3 x < -1 + (1/5) (x + y)
  
        c1, c2 :: BoolExpr (Either SAT.Lit (LA.Atom Rational))
        c1 = ite (Atom (Left a) :: BoolExpr (Either SAT.Lit (LA.Atom Rational))) (Atom $ Right le1) (Atom $ Right eq2)
        c2 = Atom (Right gt3) .||. (Atom (Left a) .<=>. Atom (Right lt4))
  
    addFormula c1
    addFormula c2
  
    ret <- SAT.solve satSolver
    ret @?= True
  
    m1 <- SAT.getModel satSolver
    m2 <- Simplex.getModel lraSolver
    let f (Left lit) = SAT.evalLit m1 lit
        f (Right atom) = LA.eval m2 atom
    fold f c1 @?= True
    fold f c2 @?= True
  
  
  case_QF_EUF :: Assertion
  case_QF_EUF = do
    satSolver <- SAT.newSolver
    eufSolver <- EUF.newSolver
    enc <- Tseitin.newEncoder satSolver
  
    tblRef <- newIORef (Map.empty :: Map (EUF.Term, EUF.Term) SAT.Var)
    defsRef <- newIORef (IntMap.empty :: IntMap (EUF.Term, EUF.Term))
    eufModelRef <- newIORef (undefined :: EUF.Model)
  
    let abstractEUFAtom :: (EUF.Term, EUF.Term) -> IO SAT.Lit
        abstractEUFAtom (t1,t2) | t1 >= t2 = abstractEUFAtom (t2,t1)
        abstractEUFAtom (t1,t2) = do
          tbl <- readIORef tblRef
          case Map.lookup (t1,t2) tbl of
            Just v -> return v
            Nothing -> do
              v <- SAT.newVar satSolver
              writeIORef tblRef $! Map.insert (t1,t2) v tbl
              modifyIORef' defsRef $! IntMap.insert v (t1,t2)
              return v
  
        abstract :: BoolExpr (Either SAT.Lit (EUF.Term, EUF.Term)) -> IO Tseitin.Formula
        abstract e = liftM (fold Tseitin.Atom) $ Traversable.mapM f e
          where
            f (Left lit) = return lit
            f (Right atom) = abstractEUFAtom atom
  
    let tsolver =
          TheorySolver
          { thAssertLit = \_ l -> do
              defs <- readIORef defsRef
              case IntMap.lookup (SAT.litVar l) defs of
                Nothing -> return True
                Just (t1,t2) -> do
                  if SAT.litPolarity l then
                    EUF.assertEqual' eufSolver t1 t2 (Just l)
                  else
                    EUF.assertNotEqual' eufSolver t1 t2 (Just l)
                  return True
          , thCheck = \callback -> do
              b <- EUF.check eufSolver
              when b $ do
                defs <- readIORef defsRef
                forM_ (IntMap.toList defs) $ \(v, (t1, t2)) -> do
                  b2 <- EUF.areEqual eufSolver t1 t2
                  when b2 $ do
                    _ <- callback v
                    return ()
              return b
          , thExplain = \m -> do
              case m of
                Nothing -> liftM IntSet.toList $ EUF.explain eufSolver Nothing
                Just v -> do
                  defs <- readIORef defsRef
                  case IntMap.lookup v defs of
                    Nothing -> error "should not happen"
                    Just (t1,t2) -> do
                      liftM IntSet.toList $ EUF.explain eufSolver (Just (t1,t2))
          , thPushBacktrackPoint = do
              EUF.pushBacktrackPoint eufSolver
          , thPopBacktrackPoint = do
              EUF.popBacktrackPoint eufSolver
          , thConstructModel = do
              writeIORef eufModelRef =<< EUF.getModel eufSolver
              return ()
          }
    SAT.setTheory satSolver tsolver
  
    cTrue  <- EUF.newConst eufSolver
    cFalse <- EUF.newConst eufSolver
    EUF.assertNotEqual eufSolver cTrue cFalse
    boolToTermRef <- newIORef (IntMap.empty :: IntMap EUF.Term)
    termToBoolRef <- newIORef (Map.empty :: Map EUF.Term SAT.Lit)
  
    let connectBoolTerm :: SAT.Lit -> EUF.Term -> IO ()
        connectBoolTerm lit t = do
          lit1 <- abstractEUFAtom (t, cTrue)
          lit2 <- abstractEUFAtom (t, cFalse)
          SAT.addClause satSolver [-lit, lit1]  --  lit  ->  lit1
          SAT.addClause satSolver [-lit1, lit]  --  lit1 ->  lit
          SAT.addClause satSolver [lit, lit2]   -- -lit  ->  lit2
          SAT.addClause satSolver [-lit2, -lit] --  lit2 -> -lit
          modifyIORef' boolToTermRef $ IntMap.insert lit t
          modifyIORef' termToBoolRef $ Map.insert t lit
  
        boolToTerm :: SAT.Lit -> IO EUF.Term
        boolToTerm lit = do
          tbl <- readIORef boolToTermRef
          case IntMap.lookup lit tbl of
            Just t -> return t
            Nothing -> do
              t <- EUF.newConst eufSolver
              connectBoolTerm lit t
              return t
  
        termToBool :: EUF.Term -> IO SAT.Lit
        termToBool t = do
          tbl <- readIORef termToBoolRef
          case Map.lookup t tbl of
            Just lit -> return lit
            Nothing -> do
              lit <- SAT.newVar satSolver
              connectBoolTerm lit t
              return lit
  
    let addFormula :: BoolExpr (Either SAT.Lit (EUF.Term, EUF.Term)) -> IO ()
        addFormula c = Tseitin.addFormula enc =<< abstract c
  
    do
      x <- SAT.newVar satSolver
      x' <- boolToTerm x
      f <- EUF.newFun eufSolver
      fx <- termToBool (f x')
      ftt <- abstractEUFAtom (f cTrue, cTrue)
      ret <- SAT.solveWith satSolver [-fx, ftt]
      ret @?= True
  
      m1 <- SAT.getModel satSolver
      m2 <- readIORef eufModelRef
      let e (Left lit) = SAT.evalLit m1 lit
          e (Right (lhs,rhs)) = EUF.eval m2 lhs == EUF.eval m2 rhs
      fold e (notB (Atom (Left fx)) .||. (Atom (Right (f cTrue, cTrue)))) @?= True
      SAT.evalLit m1 x @?= False
  
      ret <- SAT.solveWith satSolver [-fx, ftt, x]
      ret @?= False
  
    do
      -- a : Bool
      -- f : U -> U
      -- x : U
      -- y : U
      -- (a or x=y)
      -- f x /= f y
      a <- SAT.newVar satSolver
      f <- EUF.newFun eufSolver
      x <- EUF.newConst eufSolver
      y <- EUF.newConst eufSolver
      let c1, c2 :: BoolExpr (Either SAT.Lit (EUF.Term, EUF.Term))
          c1 = Atom (Left a) .||. Atom (Right (x,y))
          c2 = notB $ Atom (Right (f x, f y))
      addFormula c1
      addFormula c2
      ret <- SAT.solve satSolver
      ret @?= True
      m1 <- SAT.getModel satSolver
      m2 <- readIORef eufModelRef
      let e (Left lit) = SAT.evalLit m1 lit
          e (Right (lhs,rhs)) = EUF.eval m2 lhs == EUF.eval m2 rhs
      fold e c1 @?= True
      fold e c2 @?= True
  
      ret <- SAT.solveWith satSolver [-a]
      ret @?= False
  
  
  satTheorySolverTestGroup :: TestTree
  satTheorySolverTestGroup = $(testGroupGenerator)
  