module Exercise_9 where

{-WETT-}
import Data.Maybe
import Data.List
import qualified Data.HashSet as HashSet
import qualified Data.HashMap.Strict as HashMap

type Name = String
type Valuation = [(Name,Bool)]
data Atom = T | Var Name
  deriving (Eq, Show)
data Literal = Pos Atom | Neg Atom
  deriving (Eq, Show)
data Form = L Literal | Form :&: Form | Form :|: Form
  deriving (Eq, Show)
type Clause = [Literal]
type ConjForm = [Clause]
{-T9.3.2-}
top :: Literal
top = Pos T

bottom :: Literal
bottom = Neg T

{-T9.3.3-}
clauseToForm :: Clause -> Form
clauseToForm [] = L bottom
clauseToForm ls = foldr ((:|:) . L) (L $ last ls) (init ls)

conjToForm :: ConjForm -> Form
conjToForm [] = L top
conjToForm ds = foldr ((:&:) . clauseToForm) (clauseToForm $ last ds) (init ds)

{-T9.3.4-}
substLiteral :: Valuation -> Literal -> Literal
substLiteral v l@(Pos (Var n)) = case lookup n v of
  Just b -> if b then top else bottom
  Nothing -> l
substLiteral v l@(Neg (Var n)) = case lookup n v of
  Just b -> if b then bottom else top
  Nothing -> l
substLiteral v l = l

substClause :: Valuation -> Clause -> Clause
substClause = map . substLiteral

substConj :: Valuation -> ConjForm -> ConjForm
substConj = map . substClause

{-H9.2.1-}
containsTop :: Clause -> Bool
containsTop [] = False
containsTop (l:ls) = l == top || containsTop ls

simpConj :: ConjForm -> ConjForm
simpConj [] = []
simpConj cs = if elem [] simple then [[]] else simple where simple = [filter (/=bottom) c | c <- cs, not $ containsTop c]

{-H9.2.2-}
nnf2cnf :: Form -> Form
nnf2cnf (L literal) = (L literal)
nnf2cnf (f1 :&: f2) = nnf2cnf f1 :&: nnf2cnf f2
nnf2cnf (f1 :|: f2) = or (nnf2cnf f1) (nnf2cnf f2)
  where
    or (f11 :&: f12) f2 = (or f11 f2) :&: (or f12 f2)
    or f1 (f21 :&: f22) = (or f1 f21) :&: (or f1 f22)
    or f1 f2 = f1 :|: f2

cnf2conj :: Form -> ConjForm
cnf2conj (L literal) = [[literal]]
cnf2conj (f1 :&: f2) = cnf2conj f1 ++ cnf2conj f2
cnf2conj (f1 :|: f2) = [c1 ++ c2 | (c1,c2) <- zip (cnf2conj f1) (cnf2conj f2)]

cnf :: Form -> ConjForm
cnf = cnf2conj . nnf2cnf

{-H9.2.3-}
selectV :: ConjForm -> Maybe (Name, Bool)
selectV (((Pos (Var name)):_):_) = Just (name, True)
selectV (((Neg (Var name)):_):_) = Just (name, False)
selectV ((_:ls):cs) = if isJust x then x else selectV cs where x = selectV [ls]
selectV ([]:cs) = selectV cs
selectV _ = Nothing

{-H9.2.5-}
subst :: ConjForm -> (Name, Bool) -> ConjForm
subst [] _ = []
subst (c:cs) v@(name, True)  = if elem (Pos (Var name)) c then subst cs v else [l | l <- c, l /= Neg (Var name)] : subst cs v
subst (c:cs) v@(name, False) = if elem (Neg (Var name)) c then subst cs v else [l | l <- c, l /= Pos (Var name)] : subst cs v

showInt :: Int -> (Name, Bool)
showInt i
  | i > 0 = (show i, True)
  | i < 0 = (show (-i), False)

plrCandidates :: [HashSet.HashSet Int] -> [Int]
plrCandidates sets = HashSet.toList $ HashSet.map (\x -> negate x) $ HashSet.difference posSetInverted fullSet
  where
    posSetInverted = HashSet.map (\x -> negate x) fullSet
    fullSet = HashSet.unions sets

simplifySets :: [HashSet.HashSet Int] -> [HashSet.HashSet Int]
simplifySets [] = []
simplifySets (set:sets)
  | HashSet.member 0 set = simplifySets sets -- Remove sets containg truth value
  | HashSet.size set == 1 && HashSet.member 100000000 set = [HashSet.empty]
  | HashSet.member 100000000 set = ((head . simplifySets)  [HashSet.delete 100000000 set]) : simplifySets sets
  | otherwise = set : simplifySets sets

satConj :: ConjForm -> Maybe Valuation
satConj f = satConjSet conj (plrCandidates conj)
  where conj = (simplifySets . cnfToSets) f

satConjSet :: [HashSet.HashSet Int] -> [Int] -> Maybe Valuation
satConjSet [] _ = Just []
satConjSet sets plrCandidates
  | elem HashSet.empty sets = Nothing
  | isJust posEval = Just (showInt v:(fromJust posEval))
  | isJust negEval = Just (showInt (-v):(fromJust negEval))
  | otherwise = Nothing
    where
      v = fst selectVRes
      selectVRes = selectVSet sets plrCandidates
      posEval = satConjSet (substSet sets v) $ snd selectVRes
      negEval = satConjSet (substSet sets (-v)) $ snd selectVRes

selectVSet :: [HashSet.HashSet Int] -> [Int] -> (Int, [Int])
selectVSet sets plrCandidates
  | isJust olrCandidate = (head $ HashSet.toList (fromJust olrCandidate), plrCandidates)
  | not $ null plrCandidates = (head plrCandidates, tail plrCandidates)
  | otherwise = (selectFirstVar sets, plrCandidates)
    where
      olrCandidate = find (\set -> HashSet.size set == 1) sets

keyOfMaxVal :: HashMap.HashMap Int Int -> Int
keyOfMaxVal map = max 0 Nothing (HashMap.toList map)
  where
    max maxKey _ [] = maxKey
    max maxKey Nothing ((k,v):entries) = max k (Just v) entries
    max maxKey (Just curMax) ((k,v):entries)
      | v <= curMax = max maxKey (Just curMax) entries
      | v > curMax = max k (Just v) entries

selectBestVar :: [HashSet.HashSet Int] -> HashMap.HashMap Int Int -> Int
selectBestVar [] countMap = keyOfMaxVal countMap
selectBestVar (set:sets) countMap = selectBestVar sets (HashMap.unionWith (\k v -> v+1) countMap ((HashMap.fromList . convList . HashSet.toList) set))
  where 
    convList [] = []
    convList (k:ks) = (k, 0) : convList ks

selectFirstVar :: [HashSet.HashSet Int] -> Int
selectFirstVar [] = 0
selectFirstVar (set:sets) = head $ HashSet.toList set

substSet :: [HashSet.HashSet Int] -> Int -> [HashSet.HashSet Int]
substSet [] _ = []
substSet (set:sets) val
  | HashSet.member val set = substSet sets val
  | otherwise = (HashSet.delete (-val) set) : substSet sets val

literalToInt :: Literal -> Int
literalToInt (Pos T) = 0
literalToInt (Neg T) = 100000000
literalToInt (Pos (Var name)) = read name
literalToInt (Neg (Var name)) = -(read name)

cnfToSets :: ConjForm -> [HashSet.HashSet Int]
cnfToSets [] = []
cnfToSets (c:cs) = (HashSet.fromList c') : cnfToSets cs where c' = map (\l -> literalToInt l) c

{-H9.2.6-}
sat :: Form -> Maybe Valuation
sat = satConj . simpConj . cnf


-- This is my CDCL implementation. I implemented it badly, making it slower than my dpll on most formulas. Would probably be a lot faster with lazy data structures and better checks for satisfaction of clauses
{-
simplifySets :: [HashSet.HashSet Int] -> [HashSet.HashSet Int]
simplifySets [] = []
simplifySets (set:sets)
  | HashSet.member 0 set = simplifySets sets -- Remove sets containg truth value
  | HashSet.size set == 1 && HashSet.member 100000000 set = [HashSet.empty]
  | HashSet.member 100000000 set = ((head . simplifySets)  [HashSet.delete 100000000 set]) : simplifySets sets
  | otherwise = set : simplifySets sets

  
satConj :: ConjForm -> Maybe Valuation
satConj f 
  | snd conjPropagated == HashSet.empty && snd result = Just [(show k, getAssign ((IntMap.!) (fst result) k)) | k <- IntMap.keys (fst result)]
  | otherwise = Nothing
    where
      result = satConjSet conj (fst conjPropagated) 1
      conjPropagated = unitPropagation conj IntMap.empty 0
      conj = (simplifySets . cnfToSets) f

-- value, antecendent, decision level
type Variable = (Bool,HashSet.HashSet Int, Int)
type Assignment = IntMap.IntMap Variable

getAssign :: Variable -> Bool
getAssign (assign,_,_) = assign

getAnt :: Variable -> HashSet.HashSet Int
getAnt (_,antecendent,_) = antecendent

getDl :: Variable -> Int
getDl (_,_,dl) = dl

sigma :: HashSet.HashSet Int -> Assignment -> Int -> Int
sigma clause assignment dl = HashSet.size $ HashSet.filter (\v -> (getDl . fromJust) (IntMap.lookup (abs v) assignment) <= dl) clause

resolution :: HashSet.HashSet Int -> HashSet.HashSet Int -> HashSet.HashSet Int
resolution c1 c2 = HashSet.filter (\v -> ((not $ HashSet.member v tautologies) && (not $ HashSet.member (-v) tautologies))) unionSet
  where 
    unionSet = HashSet.union c1 c2
    tautologies = HashSet.intersection c1Neg c2
    c1Neg = HashSet.map negate c1

selectVar :: Assignment -> [Int] -> Int -> Variable
selectVar assignment (l:ls) dl
  | dl == dl2 = var
  | otherwise = selectVar assignment ls dl
    where 
      dl2 = getDl var
      var = (IntMap.!) assignment (abs l)

learnClause :: (Assignment, HashSet.HashSet Int) -> Int -> HashSet.HashSet Int
learnClause (assignment, clause) dl
  | sigma clause assignment dl == 1 = clause
  | otherwise = resolution clause (getAnt var)
    where var = selectVar assignment (HashSet.toList clause) dl

findMinDl :: Assignment -> HashSet.HashSet Int -> Int
findMinDl assignment set =  HashSet.foldr (\n m -> if n < m then n else m) 1000000 set
  where dls = HashSet.map (\v -> getDl $ (IntMap.!) assignment v) set

backtrack :: Assignment -> Int -> Assignment
backtrack assignment dl = IntMap.filter(\(_,_,dl2) -> dl2 <= dl) assignment

-- dl: decisionLevel
satConjSet :: [HashSet.HashSet Int] -> Assignment -> Int -> (Assignment, Bool)
satConjSet sets assignment dl
  | isNothing v = (assignment, True) -- satisfied
  | dl == 0 && (not . HashSet.null . snd $ propagated) = (IntMap.empty, False)
  | not . HashSet.null . snd $ propagated = satConjSet (learntClause:sets) (backtrack assignment newDl) newDl -- conflict found; learn clause and backtrack
  | otherwise = satConjSet sets (fst propagated) (dl+1) -- no conflicts found after branching; continuing...
    where
      newDl = findMinDl assignment learntClause
      learntClause = learnClause propagated dl
      v = selectFirstVar sets assignment
      justV = fromJust v
      newAssignment = (IntMap.insert (abs justV) (justV>0, HashSet.empty, dl) assignment)
      propagated = unitPropagation sets newAssignment dl

resolve :: Int -> Assignment -> Int
resolve val assignment
  | isJust assigned && ((getAssign $ fromJust assigned) == (val>0)) = 0
  | isJust assigned = 100000000
  | otherwise = val
    where assigned = IntMap.lookup (abs val) assignment

isSatisfied :: HashSet.HashSet Int -> Assignment -> Maybe Bool
isSatisfied set assignment
  | HashSet.size resolved == 0 || HashSet.member 0 resolved = Just True
  | HashSet.size resolved == 0 || (HashSet.size resolved == 1 && HashSet.member 100000000 resolved) = Just False
  | otherwise = Nothing
  where 
    resolved = HashSet.map (\v -> resolve v assignment) set

findUnitClause :: [HashSet.HashSet Int] -> Assignment -> HashSet.HashSet Int -> Maybe (HashSet.HashSet Int, Int)
findUnitClause [] _ _ = Nothing
findUnitClause (set:sets) assignment assignedVars
  | HashSet.size setAbs /= HashSet.size set = findUnitClause sets assignment assignedVars
  | HashSet.size unassignedVars == 1 && isNothing satisfaction = Just (set, olrVarSign)
  | HashSet.size unassignedVars == 0 && not (fromJust $ satisfaction) = Just (set, 100000000) -- conflict appeared
  | otherwise = findUnitClause sets assignment assignedVars
    where
      satisfaction = isSatisfied set assignment
      olrVarSign = if HashSet.member olrVar set then olrVar else -olrVar
      olrVar = head $ HashSet.toList unassignedVars
      unassignedVars = HashSet.difference setAbs assignedVars
      setAbs = HashSet.map abs set

-- satisfy all unit clauses; return nothing on conflict; else return new assignments
unitPropagation :: [HashSet.HashSet Int] -> Assignment -> Int -> (Assignment, HashSet.HashSet Int)
unitPropagation sets assignment dl
  | isJust firstUnitClause && var == 100000000 = (assignment, fst . fromJust $ firstUnitClause)
  | isJust firstUnitClause = unitPropagation sets (IntMap.insert (abs var) (var>0, fst (fromJust firstUnitClause), dl) assignment) dl
  | otherwise = (assignment, HashSet.empty)
    where
      var = snd $ fromJust firstUnitClause
      firstUnitClause = findUnitClause sets assignment $ HashSet.fromList (IntMap.keys assignment)


selectFirstVar :: [HashSet.HashSet Int] -> Assignment -> Maybe Int
selectFirstVar [] _ = Nothing
selectFirstVar (set:sets) assignment
  | length vars /= 0 = Just $ head vars
  | otherwise = selectFirstVar sets assignment
    where vars = HashSet.toList $ HashSet.filter (\v -> not $ IntMap.member (abs v) assignment) set

literalToInt :: Literal -> Int
literalToInt (Pos T) = 0
literalToInt (Neg T) = 100000000
literalToInt (Pos (Var name)) = read name
literalToInt (Neg (Var name)) = -(read name)

cnfToSets :: ConjForm -> [HashSet.HashSet Int]
cnfToSets [] = []
cnfToSets (c:cs) = (HashSet.fromList c') : cnfToSets cs where c' = map (\l -> literalToInt l) c
-}

{-debug-}
printCNF :: ConjForm -> String
printCNF [] = []
printCNF (c:cs) = printClause c ++ "\n" ++ printCNF cs

printForm :: Form -> String
printForm (L literal) = printLiteral literal
printForm (f1 :|: f2) = '(' : printForm f1 ++ " || " ++ printForm f2 ++ ")"
printForm (f1 :&: f2) = '(' : printForm f1 ++ " && " ++ printForm f2 ++ ")"

printClause :: Clause -> String
printClause [] = []
printClause (l:ls) = printLiteral l ++ printClause ls

printLiteral :: Literal -> String
printLiteral (Pos (Var name)) = name ++ " "
printLiteral (Neg (Var name)) = '-' : name ++ " "
printLiteral (Pos T) = "True "
printLiteral (Neg T) = ""

{-TTEW-}