{-# LANGUAGE TupleSections #-}
module Resolution where
import Data.Set (Set)
import qualified Data.Set as Set
import FOLSyntax
    ( Formula(Exists, Pred, Neg, Conj, Disj, Impl, All),
      Term(..),
      formulaVars,
      termFreeVars,
      termSubst,
      findFresh )
import Data.Maybe ( mapMaybe )
import Data.List
import Control.Monad (zipWithM)

type CNF = Set Clause
type Clause = Set Literal
type Literal = Formula
type Unifier = [(Term, Term)]


-- alpha equality on clauses, a clause [S(x), P(y)] is equal to [S(y), P(x)]
-- 1. get all permutations of both clauses
-- 2. build the crossproduct of permutations of c1 and c2
-- 3. unify each clause pair
-- 4. collect mgus where variables are only mapped to variables (meaning they are alpha eq)
-- 5. check if any such mgu exists
alphaEq :: Clause -> Clause -> Bool
alphaEq c1 c2 = any (uncurry unifyLiteralList) permutationPairs
    where
        permutationPairs = [(c1', c2') | c1' <- permutations $ Set.toList c1, c2' <- permutations $ Set.toList c2]

set1 :: Clause
set1 = Set.fromList [Pred "S" [Var "x"], Pred "P" [Var "y"]]
set2 :: Clause
set2 = Set.fromList [Pred "P" [Var "z"], Pred "S" [Var "y"]]

-- unification algorithm of martelli montanari
unify :: [(Term, Term)] -> Maybe Unifier
unify [] = Just []
-- (delete)
unify ((Var x, Var y) : rest) | x == y = unify rest
-- (decomp)
unify ((Fun f es, Fun g ds) : rest) | f == g && length es == length ds = unify $ zip es ds ++ rest
-- (conflict)
unify ((Fun _ _, Fun _ _) : _) = Nothing
-- (orient)
unify ((Fun f ts, Var x) : rest) = unify $ (Var x, Fun f ts) : rest
-- (occurs)
unify ((Var x, t) : _) | x `elem` termFreeVars t && Var x /= t = Nothing
-- (elim)
unify ((Var x, t) : rest) | notElem x (termFreeVars t) && x `elem` concatMap (\(t1, t2) -> termFreeVars t1 ++ termFreeVars t2) rest = unify $ (Var x, t) : map (\(t1, t2) -> (termSubst t1 x t, termSubst t2 x t)) rest
-- decent
unify ((t, s) : rest) = do
    rest' <- unify rest
    return $ (t, s) : rest'

unifyLiteralList :: [Literal] -> [Literal] -> Bool
unifyLiteralList ls1 ls2 = case zipWithM unifyLiterals ls1 ls2 of
    Nothing -> False
    Just mgus -> all checkSimpleMgu mgus
    where
        checkSimpleMgu :: Unifier -> Bool
        checkSimpleMgu [] = True
        checkSimpleMgu ((Var _, Var _) : _) = True
        checkSimpleMgu _ = False
        unifyLiterals :: Literal -> Literal -> Maybe Unifier
        unifyLiterals (Neg f) (Neg g) = unifyLiterals f g
        unifyLiterals (Pred p ts1) (Pred s ts2) | p == s && length ts1 == length ts2 = unify $ zip ts1 ts2
        unifyLiterals _ _ = Nothing

-- unifies predicates, e.g. P(x,y) == P(f(a), z)
unifyPredicates :: Literal -> Literal -> Maybe Unifier
unifyPredicates (Pred p1 ts1) (Neg (Pred p2 ts2)) | p1 == p2 && length p1 == length p2 = unify $ zip ts1 ts2
unifyPredicates _ _ = Nothing

setConcat :: Ord a => [Set a] -> Set a
setConcat = foldr Set.union Set.empty

-- applies an mgu to a given formula, asserts that the formula contains no quantifiers
applyMgu :: Formula -> [(Term, Term)] -> Formula
applyMgu (Pred p ts) mgu = Pred p $ map (`applyMguTerm` mgu) ts
applyMgu (Neg f) mgu = Neg $ applyMgu f mgu
applyMgu (Conj f1 f2) mgu = Conj (applyMgu f1 mgu) (applyMgu f2 mgu)
applyMgu (Disj f1 f2) mgu = Disj (applyMgu f1 mgu) (applyMgu f2 mgu)
applyMgu (Impl f1 f2) mgu = Impl (applyMgu f1 mgu) (applyMgu f2 mgu)
applyMgu f _ = f
-- applies mgu to term
applyMguTerm :: Term -> [(Term, Term)] -> Term
applyMguTerm (Var x) [] = Var x
applyMguTerm (Var x) ((Var y, t) : rest) = if x == y then t else applyMguTerm (Var x) rest
applyMguTerm (Fun f ts) mgu = Fun f $ map (`applyMguTerm` mgu) ts
applyMguTerm t _ = t

-- takes two clauses and makes the variables of these clauses disjunct
makeVariablesDisjunct :: Clause -> Clause -> (Clause, Clause)
makeVariablesDisjunct c1 c2 = (c1', c2')
    where
        (used, c1') = makeClauseDisjunct [] c1
        (_, c2') = makeClauseDisjunct used c2
-- takes a list of variable names and ensures that the clause does not contain these variables (by renaming), then returns all variables used in the clause + used before
makeClauseDisjunct :: [String] -> Clause -> ([String], Clause)
makeClauseDisjunct used clause = (concatMap formulaVars newClause ++ newUsed, newClause)
    where
        criticalVars = used `intersect` concatMap formulaVars clause
        (newUsed, newClause) = foldr foldFun (used, clause) criticalVars
            where
                foldFun oldVar (used', clause') = (v' : used', Set.map (\f -> renameFormula f oldVar v') clause')
                    where
                        v' = findFresh used'

-- renames all occurences of variable v with v' in a term
renameTerm :: Term -> String -> String -> Term
renameTerm t@(Var x) v v' = if x == v then Var v' else t
renameTerm (Fun g ts) v v' = Fun g (map (\t -> renameTerm t v v') ts)

-- renames all occurences of free variable v with v' in a formula
renameFormula :: Formula -> String -> String -> Formula
renameFormula (Pred p ts) v v' = Pred p (map (\t -> renameTerm t v v') ts)
renameFormula (Neg f') v v' = Neg $ renameFormula f' v v'
renameFormula (Conj f1 f2) v v' = Conj (renameFormula f1 v v') (renameFormula f2 v v')
renameFormula (Disj f1 f2) v v' = Disj (renameFormula f1 v v') (renameFormula f2 v v')
renameFormula (Impl f1 f2) v v' = Impl (renameFormula f1 v v') (renameFormula f2 v v')
renameFormula (All y f') v v' | y /= v = All y $ renameFormula f' v v'
renameFormula (Exists y f') v v' | y /= v = Exists y $ renameFormula f' v v'
renameFormula f' _ _ = f'

-- takes two clauses and tries to unify every literal in a crossproduct, returns the set of all clauses resulting from this resolution step
resolveClauses :: Clause -> Clause -> Set Clause
resolveClauses c1 c2 = newClauses
    where
        -- first make variables in both clauses disjunct
        (c1', c2') = makeVariablesDisjunct c1 c2
        zippedLiterals = [(lit1, lit2) | lit1 <- Set.toList c1', lit2 <- Set.toList c2']
        newClauses = Set.fromList $ mapMaybe (\(l1, l2) -> do
                                            -- first calculate mgu
                                            mgu <- unifyPredicates l1 l2
                                            -- now apply mgu to clauses without l1 and l2
                                            let c1'' = Set.map (`applyMgu` mgu) (c1' `Set.difference` Set.singleton l1)
                                            let c2'' = Set.map (`applyMgu` mgu) (c2' `Set.difference` Set.singleton l2)
                                            return $ c1'' `Set.union` c2''
                            ) zippedLiterals

-- a single resolution step as described in gloin
resolveStep :: CNF -> Either () CNF
resolveStep clauses = if Set.empty `Set.member` clauses then Left () else Right $ clauses `Set.union` newClauses
    where
        -- cross product of clauses
        zippedClauses = [(c1, c2) | c1 <- Set.toList clauses, c2 <- Set.toList clauses, c1 /= c2]
        -- after resolveStep add every clause that can be gained through resolution to clauseSet
        newClauses = setConcat $ map (uncurry resolveClauses) zippedClauses

resolveStepVerbose :: Set ([Clause], Clause) -> Either (Set ([Clause], Clause)) (Set ([Clause], Clause))
resolveStepVerbose clauses = if Set.empty `inRight` clauses then Left clauses else Right $ clauses `Set.union` newClauses
    where
        inRight clause set = clause `Set.member` Set.map snd set
        -- cross product of clauses
        zippedClauses = [(c1, c2) | c1 <- Set.toList (Set.map snd clauses), c2 <- Set.toList (Set.map snd clauses), c1 /= c2]
        -- after resolveStep add every clause that can be gained through resolution to clauseSet
        newClauses = setConcat' $ map (\(c1, c2) -> ([c1, c2], resolveClauses c1 c2)) zippedClauses

setConcat' :: [([Clause], Set Clause)] -> Set ([Clause], Clause)
setConcat' [] = Set.empty
setConcat' ((cs, clauses) : rest) = Set.map (cs,) clauses `Set.union` setConcat' rest

-- do resolution until we have proven unfulfillability of formula set
doResolution :: CNF -> Either () CNF
doResolution f = do
    f' <- resolveStep f
    doResolution f'

doResolutionIO :: CNF -> IO (Either () (Set ([Clause], Clause)))
doResolutionIO cnf = do
    let preparedCNF = Set.map ([], ) cnf
    resolutionHelper preparedCNF 1
    where
        resolutionHelper :: Set ([Clause], Clause) -> Int -> IO (Either () (Set ([Clause], Clause)))
        resolutionHelper set n = do
            let setE = resolveStepVerbose set
            case setE of
                Left _ -> do
                    putStrLn "Resolved to empty clause, so formula is proven!"
                    return $ Left ()
                Right set' -> do
                    putStrLn $ "Resolution iteration " ++ show n ++ ":\n"
                    printResolvents (Set.toList $ set' `Set.difference` set)
                    putStrLn ""
                    resolutionHelper set' (n + 1)

printResolvents :: [([Clause], Clause)] -> IO ()
printResolvents (([c1, c2], c) : rest) = do
    putStrLn $ "Resolving " ++ showClause c1 ++ " and " ++ showClause c2 ++ "\n-> " ++ showClause c
    printResolvents rest
printResolvents _ = return ()

showClause :: Clause -> String
showClause c = "[" ++ intercalate ", " (map show (Set.toList c)) ++ "]"