{-
Some algorithms in the propositional calculus (see examples lower down for invocation):

  * Parse a propositional formula
  * Convert a propositional formula to Conjunctive Normal Form (CNF)
  * Convert a propositional formula to Disjunctive Normal Form (DNF)
  * Davis - Putnam - Logemann - Loveland (DPLL) 
  * Binary Decision Diagrams (BDD)

-}

import Text.ParserCombinators.Parsec
import Data.Map as M
import Data.Set as S
import Data.List as L
import Control.Monad.State as S

data Atomic = Atomic String |
              AtFalse |
              AtTrue deriving (Ord, Eq)

instance Show Atomic where
    show (Atomic s) = s
    show AtFalse = "F"
    show AtTrue = "T"

data PropFormula = AtomicPropFormula Atomic |
               NotPropFormula PropFormula |
               AndPropFormula PropFormula PropFormula |
               OrPropFormula PropFormula PropFormula |
               ImpliesPropFormula PropFormula PropFormula |
               IffPropFormula PropFormula PropFormula 
                              deriving (Ord, Eq)

instance Show PropFormula where
    show p =
        let showe f = case f of
                        AtomicPropFormula n -> show n
                        x@_ -> "(" ++ (show x) ++ ")"
        in case p of
             (AtomicPropFormula at) -> show at
             (NotPropFormula f) -> "~" ++ (showe f)
             (AndPropFormula f g) -> (showe f)  ++ " /\\ " ++ (showe g)
             (OrPropFormula f g) -> (showe f)  ++ " \\/ " ++ (showe g)
             (ImpliesPropFormula f g) -> (showe f) ++ " => " ++ (showe g)
             (IffPropFormula f g) -> (showe f) ++ " <=> " ++ (showe g)

{- PARSER -}

eatWhite = (skipMany (oneOf " \t"))

readChar c = eatWhite >> char c >> eatWhite
readString s = eatWhite >> string s >> eatWhite

parseAtom :: Parser PropFormula
parseAtom = do
  eatWhite
  s <- many (alphaNum <|> (char '_')) 
  eatWhite
  if (s == "T") then return (AtomicPropFormula AtTrue) 
    else if (s == "F") then return (AtomicPropFormula AtFalse) 
         else return (AtomicPropFormula (Atomic s))

parseFactor :: Parser PropFormula
parseFactor = do
  (((readChar '~') >> parseFactor >>= (\f -> return (NotPropFormula f))) <|>
   ((readChar '(') >> parseFormula >>= (\f -> (readChar ')') >> (return f))) <|>
   parseAtom)

parseTerm :: Parser PropFormula
parseTerm =    
   do
     t1 <- parseFactor
     ((
       (readString "or" >> parseTerm >>= (\t2 -> return (OrPropFormula t1 t2))) <|>
       (readString "and" >> parseTerm >>= (\t2 -> return (AndPropFormula t1 t2)))
      ) <|> (return t1))

parseFormula :: Parser PropFormula
parseFormula = 
   do
     t1 <- parseTerm
     ((
       (readString "=>" >> parseFormula >>= (\t2 -> return (ImpliesPropFormula t1 t2))) <|>
       (readString "<=>" >> parseFormula >>= (\t2 -> return (IffPropFormula t1 t2)))
      ) <|> (return t1))

fromString s = case parse parseFormula "" s of
                 Right p -> p

{- INTERPRETER -}

data Interpretation = Interpretation (M.Map String Bool) deriving (Eq, Ord)

instance Show Interpretation where
    show (Interpretation m) = show m

interp :: Interpretation -> PropFormula -> Bool
interp (Interpretation m) (AtomicPropFormula a) = 
    case a of
      Atomic s -> M.findWithDefault False s m
      AtTrue -> True
      AtFalse -> False
interp i (NotPropFormula f) = not (interp i f)
interp i (AndPropFormula f1 f2) = (interp i f1) && (interp i f2)
interp i (OrPropFormula f1 f2) = (interp i f1) || (interp i f2)
interp i (ImpliesPropFormula f1 f2) = if (interp i f1) then (interp i f2) else True
interp i (IffPropFormula f1 f2) = 
    let i1 = (interp i f1)
        i2 = (interp i f2)
    in (i1 && i2) || ((not i1) && (not i2))

{- convert a formula to negated normal form -}
nnf :: PropFormula -> PropFormula
nnf f = let 
    edn (NotPropFormula p) = case p of
              NotPropFormula p' -> p'
              (AndPropFormula p1 p2) -> OrPropFormula (edn (NotPropFormula p1)) (edn (NotPropFormula p2))
              (OrPropFormula p1 p2) -> AndPropFormula (edn (NotPropFormula p1)) (edn (NotPropFormula p2))
              _ -> NotPropFormula p 
    in case f of
         (NotPropFormula f) -> edn (NotPropFormula (nnf f))
         (AndPropFormula f1 f2) -> (AndPropFormula (nnf f1) (nnf f2))
         (OrPropFormula f1 f2) -> (OrPropFormula (nnf f1) (nnf f2))
         (ImpliesPropFormula f1 f2) -> nnf (OrPropFormula (NotPropFormula f1) f2)
         (IffPropFormula f1 f2) -> nnf (AndPropFormula (ImpliesPropFormula f1 f2) (ImpliesPropFormula f2 f1))
         _ -> f

{- e.g. (A \/ B \/ ~C) /\ (D \/ ~F \/ ~J) /\ ~Q is represented as {({A,B},{C}),({D},{F,J}),({},{Q})}-}
data CNormalForm = CNormalForm (S.Set (S.Set Atomic, S.Set Atomic)) deriving (Eq,Ord)

instance Show CNormalForm where
    show (CNormalForm ss) = concat (intersperse " /\\ " (L.map (\s -> "(" ++ (show1 s) ++ ")") (S.toList ss)))
        where show1 (x,y) = concat (intersperse " \\/ " ((L.map show (S.toList x)) ++ 
                                      (L.map (\s -> "~" ++ (show s)) (S.toList y))))

{-convert NNF to equivalent CNormalForm by pushing disjunctions-}
cnfStep1 :: PropFormula -> CNormalForm
cnfStep1 f =
    let combineOr (CNormalForm ss1) (CNormalForm ss2) = 
            let combineOr1 (x,y) ss = S.map (\(x',y') -> ((S.union x x'), (S.union y y'))) ss
            in CNormalForm (S.unions (S.toList (S.map (\s -> combineOr1 s ss2) ss1)))
        combineAnd (CNormalForm s1) (CNormalForm s2) = CNormalForm (S.union s1 s2)
        returnAt at = CNormalForm (S.singleton (S.singleton at, S.empty))
        returnNegAt at = CNormalForm (S.singleton (S.empty, S.singleton at))
    in case f of
         AtomicPropFormula at -> returnAt at 
         NotPropFormula (AtomicPropFormula at) -> returnNegAt at
         OrPropFormula f1 f2 -> combineOr (cnfStep1 f1) (cnfStep1 f2)
         AndPropFormula f1 f2 -> combineAnd (cnfStep1 f1) (cnfStep1 f2)

whileChanged :: (Eq a) => (a -> a) -> a -> a 
whileChanged f a = let a' = f a 
                   in if a' == a then a else whileChanged f a'
                   
{-Simplification steps-}

cnfStep2_1 ss = S.filter (not . (\(x,y) -> (not (S.null (S.intersection x y))))) ss 

cnfStep2_2 ss = S.difference ss (S.filter (isSupersetOfAny ss) ss)
    where isSupersetOfAny ss (x,y) = L.any (\(x',y') -> ((x',y') /= (x,y)) 
                                                        && (S.isSubsetOf x' x) && (S.isSubsetOf y' y)) (S.toList ss)

-- find all pairs satisfying test, return them and the remainder of the set without them
findPairs :: (Eq a, Ord a) => Set a -> (a -> a -> Bool) -> ([(a,a)], Set a)
findPairs s test = let sl = S.toList s
                       slp = [(x,x') | x <- sl, x' <- sl, test x x']
                   in (slp, foldl (\s' (x,x') -> S.delete x' (S.delete x s')) s slp)

cnfStep2_3 ss = 
    let test (x,y) (x',y') = (S.size (S.intersection x y') == 1) && (S.difference x y' == x') && (S.difference y' x == y)
        collapse (x,y) (x',y') = (S.difference x (S.intersection x y'),y)
    in whileChanged (\ss -> let (p,ss') = findPairs ss test
                            in foldl (\ss' (x,x') -> S.insert (collapse x x') ss') ss' p) ss

{-Simplify CNormalForm-}
cnfStep2 :: CNormalForm -> CNormalForm
cnfStep2 (CNormalForm nf) = CNormalForm $ whileChanged (cnfStep2_3 . cnfStep2_2 . cnfStep2_1) nf

elimTF (CNormalForm nf) = CNormalForm $ S.filter (not . (\(x,y) -> (S.member AtTrue x) || (S.member AtFalse y))) 
                            (S.map (\(x,y) -> (S.delete AtFalse x, S.delete AtTrue y)) nf)

cnf2prop :: CNormalForm -> PropFormula
cnf2prop (CNormalForm nf) = 
    let orn f y = (case (S.toList y) of
                     [] -> Nothing
                     (at:[]) -> Just ((f (AtomicPropFormula at)))
                     (at:atl) -> Just (foldl (\p at -> OrPropFormula p (f (AtomicPropFormula at))) (f (AtomicPropFormula at)) atl))
        or_together (x,y) = case (orn id x, orn NotPropFormula y) of
                              (Nothing, Just l) -> l
                              (Just l, Nothing) -> l
                              (Just l1, Just l2) -> OrPropFormula l1 l2
    in case (S.toList nf) of
         [] -> AtomicPropFormula AtTrue
         ((x,y):[]) -> if ((S.null x) && (S.null y)) then AtomicPropFormula AtFalse else (or_together (x,y))
         ((x,y):nf') -> foldl (\p (x,y) -> AndPropFormula p (or_together (x,y))) (or_together (x,y)) nf'

{- Convert a propositional formula to CNF; could post-compose with cnf2prop to present it as a formula explicitly -}
cnf :: PropFormula -> CNormalForm
cnf = cnfStep2 . elimTF . cnfStep1 . nnf 

{- 
Example use:

  cnf (fromString "~(((A => B) => A) => (B <=> (~A and B))) <=> (B => ~A)")

evaluates to

  (~A \/ ~B) /\ (A) /\ (B)
-}

{-  DNF  -}

{- e.g. (A /\ B /\ ~C) \/ (D /\ ~F /\ ~J) \/ ~Q is represented as {({A,B},{C}),({D},{F,J}),({},{Q})}-}
data DNormalForm = DNormalForm (S.Set (S.Set Atomic, S.Set Atomic)) deriving (Eq,Ord)

instance Show DNormalForm where
    show (DNormalForm ss) = concat (intersperse " \\/ " (L.map (\s -> "(" ++ (show1 s) ++ ")") (S.toList ss)))
        where show1 (x,y) = concat (intersperse " /\\ " ((L.map show (S.toList x)) ++ 
                                      (L.map (\s -> "~" ++ (show s)) (S.toList y))))

{-convert NNF to equivalent DNormalForm by pushing conjunctions-}
dnfStep1 :: PropFormula -> DNormalForm
dnfStep1 f =
    let combineOr (DNormalForm ss1) (DNormalForm ss2) = 
            let combineOr1 (x,y) ss = S.map (\(x',y') -> ((S.union x x'), (S.union y y'))) ss
            in DNormalForm (S.unions (S.toList (S.map (\s -> combineOr1 s ss2) ss1)))
        combineAnd (DNormalForm s1) (DNormalForm s2) = DNormalForm (S.union s1 s2)
        returnAt at = DNormalForm (S.singleton (S.singleton at, S.empty))
        returnNegAt at = DNormalForm (S.singleton (S.empty, S.singleton at))
    in case f of
         AtomicPropFormula at -> returnAt at 
         NotPropFormula (AtomicPropFormula at) -> returnNegAt at
         AndPropFormula f1 f2 -> combineOr (dnfStep1 f1) (dnfStep1 f2)
         OrPropFormula f1 f2 -> combineAnd (dnfStep1 f1) (dnfStep1 f2)

{-Simplification steps-}

dnfStep2_1 ss = S.filter (not . (\(x,y) -> (not (S.null (S.intersection x y))))) ss 

dnfStep2_2 ss = S.difference ss (S.filter (isSupersetOfAny ss) ss)
    where isSupersetOfAny ss (x,y) = L.any (\(x',y') -> ((x',y') /= (x,y)) 
                                                        && (S.isSubsetOf x' x) && (S.isSubsetOf y' y)) (S.toList ss)

dnfStep2_3 ss = 
    let test (x,y) (x',y') = (S.size (S.intersection x y') == 1) && (S.difference x y' == x') && (S.difference y' x == y)
        collapse (x,y) (x',y') = (S.difference x (S.intersection x y'),y)
    in whileChanged (\ss -> let (p,ss') = findPairs ss test
                            in foldl (\ss' (x,x') -> S.insert (collapse x x') ss') ss' p) ss

{-Simplify DNormalForm-}
dnfStep2 :: DNormalForm -> DNormalForm
dnfStep2 (DNormalForm nf) = DNormalForm $ whileChanged (dnfStep2_3 . dnfStep2_2 . dnfStep2_1) nf

delimTF (DNormalForm nf) = DNormalForm $ S.filter (not . (\(x,y) -> (S.member AtFalse x) || (S.member AtTrue y))) 
                            (S.map (\(x,y) -> (S.delete AtTrue x, S.delete AtFalse y)) nf)

dnf2prop :: DNormalForm -> PropFormula
dnf2prop (DNormalForm nf) = 
    let orn f y = (case (S.toList y) of
                     [] -> Nothing
                     (at:[]) -> Just ((f (AtomicPropFormula at)))
                     (at:atl) -> Just (foldl (\p at -> AndPropFormula p (f (AtomicPropFormula at))) (f (AtomicPropFormula at)) atl))
        or_together (x,y) = case (orn id x, orn NotPropFormula y) of
                              (Nothing, Just l) -> l
                              (Just l, Nothing) -> l
                              (Just l1, Just l2) -> AndPropFormula l1 l2
    in case (S.toList nf) of
         [] -> AtomicPropFormula AtTrue
         ((x,y):[]) -> if ((S.null x) && (S.null y)) then AtomicPropFormula AtFalse else (or_together (x,y))
         ((x,y):nf') -> foldl (\p (x,y) -> OrPropFormula p (or_together (x,y))) (or_together (x,y)) nf'

{- Convert a propositional formula to DNF-}
dnf :: PropFormula -> DNormalForm
dnf = dnfStep2 . delimTF . dnfStep1 . nnf 

{-DPLL-}
data DPLL = Satisfiable Interpretation | Inconsistent deriving Show

whileChangeM_ m = do {s <- get; m; s' <- get; if (s /= s') then whileChangeM_ m else return ()}

cnormalFormToList (CNormalForm cnf) = L.map (\(x,y) -> (S.toList x,S.toList y)) (S.toList cnf)

getClausesAndAssignments = do
  (i, cnf) <- get
  return (cnormalFormToList cnf, i)

emptyClause = (S.empty, S.empty)

emptyClauseGenerated = do
  (_, (CNormalForm cnf)) <- get
  if S.member emptyClause cnf then return True else return False

deleteClausesContaining f p = do
  modify (\(i, (CNormalForm cnf)) -> 
              (i, (CNormalForm (S.filter (\x -> not (S.member p (f x))) cnf))))
deleteClausesContainingL p = deleteClausesContaining fst p
deleteClausesContainingR p = deleteClausesContaining snd p

removePropFromClausesL p = do
  modify (\(i, (CNormalForm cnf)) -> 
              (i, (CNormalForm (S.map (\(x,y) -> (S.delete p x,y)) cnf))))

removePropFromClausesR :: Atomic -> S.State (Interpretation, CNormalForm) ()
removePropFromClausesR p = do
  modify (\(i, (CNormalForm cnf)) -> 
              (i, (CNormalForm (S.map (\(x,y) -> (x,S.delete p y)) cnf))))

getUnitsL = do
  (_, (CNormalForm cnf)) <- get
  return (S.toList (S.map (\(x,y) -> case (S.toList x) of [a] -> a) 
                          (S.filter (\(x,y) -> (S.null y) && ((S.size x) == 1)) cnf)))

getUnitsR = do
  (_, (CNormalForm cnf)) <- get
  return (S.toList (S.map (\(x,y) -> case (S.toList y) of [a] -> a) 
                          (S.filter (\(x,y) -> (S.null x) && ((S.size y) == 1)) cnf)))

unionL cnf = S.unions (S.toList (S.map fst cnf))
unionR cnf = S.unions (S.toList (S.map snd cnf))

getPureLiteralsL = do
  (_, (CNormalForm cnf)) <- get
  return (S.toList (S.difference (unionL cnf) (unionR cnf)))

getPureLiteralsR = do
  (_, (CNormalForm cnf)) <- get
  return (S.toList ((S.difference (unionR cnf) (unionL cnf))))

assignToInterp b p@(Atomic name) = do
  modify (\((Interpretation m), cnf) -> 
              ((Interpretation (M.insert name b m)), cnf))
  case b of
    True -> deleteClausesContainingL p >> removePropFromClausesR p
    False -> deleteClausesContainingR p >> removePropFromClausesL p

pureLiterals = do
  pureL <- getPureLiteralsL
  mapM_ (assignToInterp True) pureL
  pureR <- getPureLiteralsR
  mapM_ (assignToInterp False) pureR

unitClauses = do
  unitsL <- getUnitsL
  mapM_ (assignToInterp True) unitsL
  unitsR <- getUnitsR
  mapM_ (assignToInterp False) unitsR

caseSplit p = do
  s <- get
  assignToInterp True p
  r <- dpllM
  case r of
    Inconsistent -> do 
                modify (const s)
                assignToInterp False p
                dpllM
    _ -> return r

dpllM = do
  whileChangeM_ (unitClauses >> pureLiterals)
  ec <- emptyClauseGenerated
  if ec then return Inconsistent
   else 
       do
         (c,a) <- getClausesAndAssignments
         case c of 
           [] -> return (Satisfiable a)
           (x:_) -> (case x of
                       (p:_,_) -> caseSplit p
                       (_,p:_) -> caseSplit p)

dpll f = evalState dpllM (Interpretation (M.empty), cnf f) 

{-
Example: Pierce's Law:

   dpll (fromString "~(((A => B) => A) => A)")  

-}

{-BDD-}

data Node = NTrue | NFalse | Node String Int deriving (Eq,Ord)

instance Show Node where
    show NTrue = "T"
    show NFalse = "F"
    show (Node s i) = s ++ (show i)

data BDD = BDD Node (M.Map Node (Node,Node)) deriving (Eq,Ord,Show)

trueBdd = BDD NTrue M.empty
falseBdd = BDD NFalse M.empty

negateBdd :: BDD -> BDD
negateBdd (BDD at m) = 
    let swap a = case a of
                   NFalse -> NTrue
                   NTrue -> NFalse
                   _ -> a
    in
      case at of
        NTrue -> falseBdd
        NFalse -> trueBdd
        _ -> (BDD at (M.map (\(a,a') -> (swap a, swap a')) m))

orCombine p p' z z' = 
    case (p,p') of
      (NTrue,_) -> trueBdd
      (NFalse,_) -> z'
      (_,NTrue) -> trueBdd
      (_,NFalse) -> z

andCombine p p' z z' = 
    case (p,p') of
      (NTrue,_) -> z'
      (NFalse,_) -> falseBdd
      (_,NTrue) -> z
      (_,NFalse) -> falseBdd

thenCombine p p' z z' = 
    case (p,p') of
      (NTrue,_) -> z'
      (NFalse,_) -> trueBdd
      (_,NTrue) -> trueBdd
      (_,NFalse) -> negateBdd z

iffCombine p p' z z' =
    case (p,p') of
      (NTrue,_) -> z'
      (NFalse,_) -> negateBdd z'
      (_,NTrue) -> z
      (_,NFalse) -> negateBdd z

isomorphic :: BDD -> BDD -> Bool
isomorphic (BDD p m) (BDD p' m') = 
    case (p,p') of
      (NTrue, NTrue) -> True
      (NFalse, NFalse) -> True
      (Node s _, Node s' _) -> if not (s == s') then False
                               else let (Just (l,r)) = M.lookup p m
                                        (Just (l',r')) = M.lookup p' m'
                                    in (isomorphic (BDD l m) (BDD l' m')) && (isomorphic (BDD r m) (BDD r' m'))
      _ -> False

combineBdds :: (String -> String -> Ordering) -> (Node -> Node -> BDD -> BDD -> BDD) -> BDD -> BDD -> BDD
combineBdds t b z@(BDD p m) z'@(BDD p' m') = 
    let below p m = case M.lookup p m of
                      Nothing -> M.empty
                      Just (l,r) -> M.insert p (l,r) (M.union (below l m) (below r m))
        truePart (BDD p m) = case M.lookup p m of
                               Just (l,r) -> BDD l (below l m)
        falsePart (BDD p m) = case M.lookup p m of
                                Just (l,r) -> BDD r (below r m)
        findIsoNode z m = case M.toList (M.filterWithKey (\p _ -> isomorphic z (BDD p m)) m) of 
                            [] -> Nothing
                            ((x,_):_) -> Just x
        newNode p = do
          m <- get
          case L.map (\(Node _ n) -> n) (L.filter (\(Node s _) -> s == p) (M.keys m)) of
            [] -> return (Node p 0)
            x -> return (Node p ((maximum x) + 1))
        merge z@(BDD p m') = do
          case p of
            Node s _ -> do
                   m <- get 
                   case findIsoNode z m of
                     Just p' -> return p'
                     Nothing -> do
                               p' <- newNode s
                               case M.lookup p m' of
                                 Just (l,r) -> do {l' <- merge (BDD l m') ; 
                                                   r' <- merge (BDD r m') ; 
                                                   modify (\m -> M.insert p' (l',r') m);
                                                   return p'}
            x@_ -> return x
        integrate p z@(BDD l ml) z' = 
            if (isomorphic z z') then z
            else let (r',ml') = runState (merge z') ml
                 in BDD p (M.insert p (l,r') ml')
    in case (p,p') of
         (p@(Node s _), p'@(Node s' _)) -> 
             case t s s' of
               LT -> integrate p (combineBdds t b (truePart z) z') (combineBdds t b (falsePart z) z')
               GT -> integrate p' (combineBdds t b z (truePart z')) (combineBdds t b z (falsePart z'))
               EQ -> integrate p (combineBdds t b (truePart z) (truePart z')) (combineBdds t b (falsePart z) (falsePart z'))
         _ -> b p p' z z'

bdd :: (String -> String -> Ordering) -> PropFormula -> BDD
bdd t (AtomicPropFormula at) = case at of
                                 Atomic s -> let n = (Node s 0) in BDD n (M.fromList [(n, (NTrue, NFalse))])
                                 AtTrue -> trueBdd
                                 AtFalse -> falseBdd
bdd t (NotPropFormula p) = negateBdd (bdd t p)
bdd t (AndPropFormula p q) = combineBdds t andCombine (bdd t p) (bdd t q)
bdd t (OrPropFormula p q) = combineBdds t orCombine (bdd t p) (bdd t q)
bdd t (ImpliesPropFormula p q) = combineBdds t thenCombine (bdd t p) (bdd t q)
bdd t (IffPropFormula p q) = combineBdds t iffCombine (bdd t p) (bdd t q)


{-`compare` here is the order relation on strings-}
makeBdd = bdd compare

{-- example 

makeBdd (fromString "(P or R) and (Q or R)")
makeBdd (fromString "(P and Q) or R")

-- an example where we need four `P` nodes:

makeBdd (fromString "(A <=> (Q => B)) <=> P")

-- test functions against eachother:

let x = (fromString "(A <=> (Q => B)) <=> P") in isomorphic (makeBdd (dnf2prop (dnf x))) (makeBdd x)

--}