{-# LANGUAGE DeriveDataTypeable #-}
{-# OPTIONS_GHC -fglasgow-exts #-}

-- A Haskell implementation of the basic algorithm (with non-chronological
-- backtracking) from "SAT-MICRO: petit mais costaud!" by Sylvain Conchon,
-- Johannes Kanig, and Stephane Lescuyer.

-- To compile:
-- ghc --make -O -W -o ms -main-is MicroSat MicroSat.hs

-- Backtracking uses the control stack, so, you may want to invoke with
-- something like
--    ms <cnf-file> +RTS -K768M -RTS
-- depending on the size of the SAT instance.

module MicroSat where

import Control.Monad.Cont
import Control.Monad.State.Strict hiding ((>=>))
import Data.Foldable hiding (sequence_)
import Data.List hiding (elem, concat, foldl', any, all)
import Data.Map (Map)
import Data.Set (Set)
import Data.Typeable
import Debug.Trace()
import Prelude hiding (or, and, all, any, elem, minimum, foldr, splitAt, concatMap, foldl, catch)
import System(getArgs)
import System.Exit
import Text.PrettyPrint.HughesPJ
import qualified Data.Foldable as Foldable
import qualified Data.List as L
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified ParseCnf


usage = "Usage:\nms <cnf-filename>"

main =
    getArgs >>= \args ->
    case args of
      [] -> putStrLn usage
      "-verify":_ -> do putStrLn "*Testing*"
                        undefined
      path:_ -> readFile path >>= parseAndSolve
         where
           parseAndSolve contents =
              let cnf = asCnf $ ParseCnf.parseCNF path contents
              in do putStrLn ("Solving " ++ path ++ "...")
                    let result = dpll cnf
                    when (not (verifyResult result cnf)) $ do
                      print "VERIFICATION ERROR"
                      exitWith (ExitFailure 1)
                    print result
       
type Cnf = [[Lit]]

asCnf :: ParseCnf.CNF -> Cnf
asCnf (ParseCnf.CNF _ _ is) = map (nub . map fromInteger) $ is

data Result = Sat [Lit] | Unsat
instance Show Result where
   show (Sat lits) = "satisfiable: " ++ intercalate " " (map show lits)
   show Unsat      = "unsatisfiable"


newtype Lit = L {unLit :: Int} deriving (Eq, Ord, Typeable)
inLit f = L . f . unLit

instance Show Lit where
    show = show . unLit
instance Read Lit where
    readsPrec i s = map (\(i,s) -> (L i, s)) (readsPrec i s :: [(Int, String)])

vari :: Lit -> Int
vari = abs . unLit

instance Num Lit where
    _ + _ = error "+ doesn't make sense for literals"
    _ - _ = error "- doesn't make sense for literals"
    _ * _ = error "* doesn't make sense for literals"
    signum _ = error "signum doesn't make sense for literals"
    negate   = inLit negate
    abs      = inLit abs
    fromInteger l | l == 0    = error "0 is not a literal"
                  | otherwise = L $ fromInteger l


-- gamma = annotated assignment literals
-- delta = annotated cnf
data StateContents = S {
      gamma :: Map Lit (Set Lit),
      delta :: [([Lit], Set Lit)]
    }
instance Show StateContents where
    show = render . stateDoc
      where
stateDoc (S {gamma=gamma, delta=delta}) =
    brackets (hcat . intersperse space . map (text . show) $ Map.keys gamma)
    <+> braces (hcat
                . intersperse (comma <> space)
                . map (\(c, a) -> braces (hcat
                                          . intersperse space
                                          . map (text . show) $ c)
                                  <> tups (hcat
                                           . intersperse comma
                                           . map (text . show)
                                           $ Set.toList a))
                $ delta)
        where tups p = char '<' <> p <> char '>'


dpll :: Cnf -> Result
dpll f = (`runCont` id) $ do
    r <- callCC $ \bj -> do
           (Right env) <- bcp bj (initialState f)
           unsat env return
    either (const $ return Unsat) (return . Sat) r

dispatch d = map (\l -> (l, d))
initialState f = S Map.empty (dispatch Set.empty f)
fromRight (Right a) = a
fromRight (Left _)  = error "fromRight: Left"

getGamma l e = Map.findWithDefault (error $ show l ++ ": annotation not found")
               l (gamma e)

-- bcp either:
--  1. finds a conflict and returns annotation literals (Left)
--  2. computes a new environment (Right)

assume bj env (l, s) = -- update only if not present
    if l `Map.member` gamma env
    then return (Right env)
    else bcp bj env{gamma = Map.insert l s (gamma env)}

-- | If @bcp@ finds a conflict, calls @bj@ with set of literals, the
-- contributors to the conflict.
bcp bj env = do
    env' <-
        foldM (\env' (cl, a) -> do
                 let (cl_neg, cl_pos) =
                         partition (\l -> negate l `Map.member` gamma env') cl
                 if any (`Map.member` gamma env') cl
                  then return env'
                  else do
                    -- update clause annotation
                    let a' = foldl'
                             (\set l -> set `Set.union` getGamma (negate l) env')
                             a cl_neg
                    case cl_pos of
                         []  -> bj (Left a')
                         [f] -> assume bj env' (f, a') >>= return . fromRight
                         _   -> return $ env'{delta = (cl_pos, a'):(delta env')})
        (env{delta = []})
        (delta env)
    return $ Right env'

-- unsat either:
--  1. returns annotation literals (Left)
--  2. finds sat assignment

-- bj is the 'back jump' function, allowing conflicts to backtrack to the
-- point where the last involved literal was decided.                           
unsat env bj =
    case delta env of
      [] -> return $ Right $ Map.keys (gamma env)
      ([_],_):_ -> error "unpropagated unit literal"
      ([],_):_  -> error "conflict unresolved"
      ((a:_),_):_ -> do
        r <- callCC $ \innerBj -> do
               (Right env') <- assume innerBj env (a, Set.singleton a)
               -- done propagating, no conflicts: continue
               unsat env' return
        case r of
          Left d ->
              if not $ a `elem` d
              then bj (Left d)
              else (callCC $ \innerBj -> do
                      (Right env') <-
                          assume innerBj env (negate a, Set.delete a d)
                      unsat env' bj)
                   >>= either (bj . Left) (return . Right)
          Right _ -> return r




paper1 :: Cnf =
  [[-1, -3, -4]
  ,[-1, -3, 4]
  ,[2, 3, 5]
  ,[3, 5]
  ,[3, -5]]



verifyResult :: Result -> Cnf -> Bool
verifyResult (Sat m) cnf =
   -- m is well-formed
   all (\l -> not $ negate l `elem` m) m
   && all (\cl -> any (`elem` cl) m) cnf
verifyResult Unsat _ = True