-- |
-- Module      : Data.Boolean.SatSolver
-- Copyright   : Sebastian Fischer
-- License     : BSD3
-- 
-- Maintainer  : Sebastian Fischer (sebf@informatik.uni-kiel.de)
-- Stability   : experimental
-- Portability : portable
-- 
-- This Haskell library provides an implementation of the
-- Davis-Putnam-Logemann-Loveland algorithm
-- (cf. <http://en.wikipedia.org/wiki/DPLL_algorithm>) for the boolean
-- satisfiability problem. It not only allows to solve boolean
-- formulas in one go but also to add constraints and query bindings
-- of variables incrementally.
-- 
-- The implementation is not sophisticated at all but uses the basic
-- DPLL algorithm with unit propagation.
-- 
module Data.Boolean.SatSolver (

  Boolean(..), SatSolver,

  newSatSolver, isSolved, 

  lookupVar, assertTrue, branchOnVar, selectBranchVar, solve

  ) where

import Data.List
import Data.Boolean

import Control.Monad.Writer

import qualified Data.IntMap as IM

-- | A @SatSolver@ can be used to solve boolean formulas.
-- 
data SatSolver = SatSolver { clauses :: CNF, bindings :: IM.IntMap Bool }
 deriving Show

-- | A new SAT solver without stored constraints.
-- 
newSatSolver :: SatSolver
newSatSolver = SatSolver [] IM.empty

-- | This predicate tells whether all constraints are solved.
-- 
isSolved :: SatSolver -> Bool
isSolved = null . clauses

-- |
-- We can lookup the binding of a variable according to the currently
-- stored constraints. If the variable is unbound, the result is
-- @Nothing@.
-- 
lookupVar :: Int -> SatSolver -> Maybe Bool
lookupVar name = IM.lookup name . bindings

-- | 
-- We can assert boolean formulas to update a @SatSolver@. The
-- assertion may fail if the resulting constraints are unsatisfiable.
-- 
assertTrue :: MonadPlus m => Boolean -> SatSolver -> m SatSolver
assertTrue formula solver =
  simplify (solver { clauses = booleanToCNF formula ++ clauses solver })

-- |
-- This function guesses a value for the given variable, if it is
-- currently unbound. As this is a non-deterministic operation, the
-- resulting solvers are returned in an instance of @MonadPlus@.
-- 
branchOnVar :: MonadPlus m => Int -> SatSolver -> m SatSolver
branchOnVar name solver =
  maybe (branchOnUnbound name solver)
        (const (return solver))
        (lookupVar name solver)

-- |
-- We select a variable from the shortest clause hoping to produce a
-- unit clause.
--
selectBranchVar :: SatSolver -> Int
selectBranchVar = literalVar . head . head . sortBy shorter . clauses

-- | 
-- This function guesses values for variables such that the stored
-- constraints are satisfied. The result may be non-deterministic and
-- is, hence, returned in an instance of @MonadPlus@.
-- 
solve :: MonadPlus m => SatSolver -> m SatSolver
solve solver
  | isSolved solver = return solver
  | otherwise = branchOnUnbound (selectBranchVar solver) solver >>= solve


-- private helper functions

updateSolver :: MonadPlus m => CNF -> [(Int,Bool)] -> SatSolver -> m SatSolver
updateSolver cs bs solver = do
  bs' <- foldr (uncurry insertBinding) (return (bindings solver)) bs
  return $ solver { clauses = cs, bindings = bs' }

insertBinding :: MonadPlus m
              => Int -> Bool -> m (IM.IntMap Bool) -> m (IM.IntMap Bool)
insertBinding name newValue binds = do
  bs <- binds
  maybe (return (IM.insert name newValue bs))
        (\oldValue -> do guard (oldValue==newValue); return bs)
        (IM.lookup name bs)

simplify :: MonadPlus m => SatSolver -> m SatSolver
simplify solver = do
  (cs,bs) <- runWriterT . simplifyClauses . clauses $ solver
  updateSolver cs bs solver

simplifyClauses :: MonadPlus m => CNF -> WriterT [(Int,Bool)] m CNF
simplifyClauses [] = return []
simplifyClauses allClauses = do
  let shortestClause = head . sortBy shorter $ allClauses
  guard (not (null shortestClause))
  if null (tail shortestClause)
   then propagate (head shortestClause) allClauses >>= simplifyClauses
   else return allClauses

propagate :: MonadPlus m => Literal -> CNF -> WriterT [(Int,Bool)] m CNF
propagate literal allClauses = do
  tell [(literalVar literal, isPositiveLiteral literal)]
  return (foldr prop [] allClauses)
 where
  prop c cs | literal `elem` c = cs
            | otherwise        = filter (invLiteral literal/=) c : cs

branchOnUnbound :: MonadPlus m => Int -> SatSolver -> m SatSolver
branchOnUnbound name solver =
  guess (Pos name) solver `mplus` guess (Neg name) solver

guess :: MonadPlus m => Literal -> SatSolver -> m SatSolver
guess literal solver = do
  (cs,bs) <- runWriterT (propagate literal (clauses solver) >>= simplifyClauses)
  updateSolver cs bs solver

shorter :: [a] -> [a] -> Ordering
shorter []     []     = EQ
shorter []     _      = LT
shorter _      []     = GT
shorter (_:xs) (_:ys) = shorter xs ys