module Main where

import Data.List
import Data.Maybe
import Debug.Trace

{-
First we define first order predicate logic
-}

data Term = Var String | Fun String [Term] deriving (Eq)
data Formula
    = Pred String [Term]
    | Neg Formula
    | Conj Formula Formula
    | Disj Formula Formula
    | Impl Formula Formula
    | All String Formula
    | Exists String Formula
    | T
    | F
    deriving (Eq)


instance Show Term where
    show (Var x) = x
    show (Fun f []) = f
    show (Fun f (x : xs)) = f ++ "(" ++ show x ++ foldr ((++) . (", "++) . show) ")" xs

instance Show Formula where
    show (Pred p []) = p
    show (Pred p (x : xs)) = p ++ "(" ++ show x ++ foldr ((++) . (", "++) . show) ")" xs
    show (Neg f) = "!(" ++ show f ++ ")"
    show (Conj f1 f2) = show f1 ++ " /\\ " ++ show f2
    show (Disj f1 f2) = show f1 ++ " \\/ " ++ show f2
    show (Impl f1 f2) = "(" ++ show f1 ++ " -> " ++ show f2 ++ ")"
    show (All x f) = "forall " ++ x ++ ". " ++ show f
    show (Exists x f) = "exists " ++ x ++ ". " ++ show f
    show T = "true"
    show F = "false"

-- free variables of a term
termFreeVars :: Term -> [String]
termFreeVars (Var x) = [x]
termFreeVars (Fun _ ts) = concatMap termFreeVars ts

-- substitution on terms
termSubst :: Term -> String -> Term -> Term
termSubst t@(Var x) y s = if x == y then s else t
termSubst (Fun f ts) y s = Fun f $ map (\t' -> termSubst t' y s) ts

-- unification algorithm of martelli montanari
unify :: [(Term, Term)] -> Maybe [(Term, Term)]
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'

{-
Now we define some functions to convert given terms to normalforms
i.e. negation normalform, prenex normal, skolemform, conjunctive normalform
-}

-- negation normalform
makeNNF :: Formula -> Formula
makeNNF form = if form == f' then form else makeNNF f'
    where
        f' = nnfStep form
        nnfStep (Conj f1 f2) = Conj (nnfStep f1) (nnfStep f2)
        nnfStep (Disj f1 f2) = Disj (nnfStep f1) (nnfStep f2)
        nnfStep (Impl f1 f2) = Disj (Neg f1) f2
        nnfStep (All x f) = All x (nnfStep f)
        nnfStep (Exists x f) = Exists x (nnfStep f)
        nnfStep (Neg (Conj f1 f2)) = Disj (Neg f1) (Neg f2)
        nnfStep (Neg (Disj f1 f2)) = Conj (Neg f1) (Neg f2)
        nnfStep (Neg (Impl f1 f2)) = Conj f1 (Neg f2)
        nnfStep (Neg (All x f)) = Exists x (Neg f)
        nnfStep (Neg (Exists x f)) = All x (Neg f)
        nnfStep (Neg (Neg f)) = f
        nnfStep (Neg T) = F
        nnfStep (Neg F) = T
        nnfStep f = f

-- infinite list of variable names
vars :: [String]
vars = ['v' : show n | n <- [(0 :: Int)..]]

-- return *all* vars in a formula, i.e. bound ones and free ones
formulaVars :: Formula -> [String]
formulaVars (Pred _ ts) = concatMap termFreeVars ts
formulaVars (Neg f) = formulaVars f
formulaVars (Conj f1 f2) = formulaVars f1 ++ formulaVars f2
formulaVars (Disj f1 f2) = formulaVars f1 ++ formulaVars f2
formulaVars (Impl f1 f2) = formulaVars f1 ++ formulaVars f2
formulaVars (All x f) = x : formulaVars f
formulaVars (Exists x f) = x : formulaVars f
formulaVars _ = []

-- 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'

-- finds a fresh variable i.e. a variable not occuring in vs
findFresh :: [String] -> String
findFresh vs = fromJust $ find (`notElem` vs) vars

-- rename every binder in formula to be disjunct
renameBinders :: Formula -> Formula
renameBinders f = fst $ go f []
    where
        go :: Formula -> [String] -> (Formula, [String])
        go (All x f') vs = (All x' f'', vs')
            where
                (f'', vs') = go (renameFormula f' x x') (x' : vs)
                x' = findFresh (formulaVars f ++ vs)
        go (Exists x f') vs = (Exists x' f'', vs')
            where
                (f'', vs') = go (renameFormula f' x x') (x' : vs)
                x' = findFresh (formulaVars f ++ vs)
        go (Neg f') vs = (Neg f'', vs')
            where
                (f'', vs') = go f' vs
        go (Conj f1 f2) vs = (Conj f1' f2', vs'')
            where
                (f1', vs') = go f1 vs
                (f2', vs'') = go f2 vs'
        go (Disj f1 f2) vs = (Disj f1' f2', vs'')
            where
                (f1', vs') = go f1 vs
                (f2', vs'') = go f2 vs'
        go (Impl f1 f2) vs = (Impl f1' f2', vs'')
            where
                (f1', vs') = go f1 vs
                (f2', vs'') = go f2 vs'
        go f' vs = (f', vs)

-- prenex normalform
-- TODO make it so that forall has higher priority to be moved left (makes skolem functions smaller)
makePNF :: Formula -> Formula
makePNF form = go $ renameBinders . makeNNF $ form
    where
        go f = let f' = pnfStep f in
            if f == f' then f else go f'
        -- swapping rules
        pnfStep (Conj phi (All x psi)) = All x (Conj phi psi)
        pnfStep (Disj phi (All x psi)) = All x (Disj phi psi)
        pnfStep (Conj (All x psi) phi) = All x (Conj psi phi)
        pnfStep (Disj (All x psi) phi) = All x (Disj psi phi)
        pnfStep (Conj phi (Exists x psi)) = Exists x (Conj phi psi)
        pnfStep (Disj phi (Exists x psi)) = Exists x (Disj phi psi)
        pnfStep (Conj (Exists x psi) phi) = Exists x (Conj psi phi)
        pnfStep (Disj (Exists x psi) phi) = Exists x (Disj psi phi)
        {-
        pnfStep (Conj phi (Exists x psi)) = Exists x (Conj phi psi)
        pnfStep (Conj phi (All x psi)) = All x (Conj phi psi)
        pnfStep (Disj phi (Exists x psi)) = Exists x (Disj phi psi)
        pnfStep (Disj phi (All x psi)) = All x (Disj phi psi)
        pnfStep (Conj (Exists x psi) phi) = Exists x (Conj psi phi)
        pnfStep (Conj (All x psi) phi) = All x (Conj psi phi)
        pnfStep (Disj (Exists x psi) phi) = Exists x (Disj psi phi)
        pnfStep (Disj (All x psi) phi) = All x (Disj psi phi)
        -}
        -- descent rules
        pnfStep (All x f) = All x (pnfStep f)
        pnfStep (Exists x f) = Exists x (pnfStep f)
        pnfStep (Conj f1 f2) = Conj (pnfStep f1) (pnfStep f2)
        pnfStep (Disj f1 f2) = Disj (pnfStep f1) (pnfStep f2)
        pnfStep (Impl f1 f2) = Impl (pnfStep f1) (pnfStep f2) -- can't happen, gets removed in makeNNF
        pnfStep (Neg f) = Neg (pnfStep f)
        pnfStep f = f

-- infinite function names for skolemizing
skolemFuns :: [String]
skolemFuns = ["sk" ++ show n | n <- [(0 :: Int)..]]

-- returns all function names in a formula
usedFunctions :: Formula -> [String]
usedFunctions (Pred _ ts) = concatMap termFunctions ts
    where
        termFunctions (Var _) = []
        termFunctions (Fun f ts') = f : concatMap termFunctions ts'
usedFunctions (Neg f) = usedFunctions f
usedFunctions (Conj f1 f2) = usedFunctions f1 ++ usedFunctions f2
usedFunctions (Disj f1 f2) = usedFunctions f1 ++ usedFunctions f2
usedFunctions (Impl f1 f2) = usedFunctions f1 ++ usedFunctions f2
usedFunctions (All _ f) = usedFunctions f
usedFunctions (Exists _ f) = usedFunctions f
usedFunctions _ = []

-- skolem form of a formula
makeSkolem :: Formula -> Formula
makeSkolem form = go (makePNF form) [] []
    where
        substTermInFormula :: Formula -> String -> Term -> Formula
        substTermInFormula (Pred p ts) x s = Pred p $ map (\t -> termSubst t x s) ts
        substTermInFormula (Neg f) x s = Neg $ substTermInFormula f x s
        substTermInFormula (Conj f1 f2) x s = Conj (substTermInFormula f1 x s) (substTermInFormula f2 x s)
        substTermInFormula (Disj f1 f2) x s = Disj (substTermInFormula f1 x s) (substTermInFormula f2 x s)
        substTermInFormula (Impl f1 f2) x s = Impl (substTermInFormula f1 x s) (substTermInFormula f2 x s)
        substTermInFormula f@(All y f') x s = if x == y then f else All y (substTermInFormula f' x s)
        substTermInFormula f@(Exists y f') x s = if x == y then f else Exists y (substTermInFormula f' x s)
        substTermInFormula f _ _ = f
        go (All x f) vs skolems = All x (go f (vs ++ [Var x]) skolems)
        go (Exists x f) vs skolems = go (substTermInFormula f x (Fun newSkolem vs)) vs (newSkolem : skolems)
            where newSkolem = fromJust $ find (`notElem` skolems) skolemFuns
        go f _ _ = f

-- conjunctive normalform, removes all quantors, so its ready for resolution
makeCNF :: Formula -> Formula
makeCNF form = go $ makeSkolem form
    where
        go f = let f' = cnfStep f in
            if f == f' then f else go f'
        cnfStep (Conj phi psi) = Conj (makeCNF phi) (makeCNF psi)
        cnfStep (Disj (Conj phi psi) xi) = Conj (makeCNF $ Disj phi xi) (makeCNF $ Disj psi xi)
        cnfStep (Disj xi (Conj phi psi)) = Conj (makeCNF $ Disj xi phi) (makeCNF $ Disj xi psi)
        cnfStep (All _ f) = makeCNF f
        cnfStep (Exists _ f) = makeCNF f
        cnfStep f = f

-- create the list of clauses from a formula
-- TODO not working correctly for `makeCNFList formula5`
makeCNFList :: Formula -> [[Formula]]
makeCNFList form = go $ makeCNF form
    where
        go (Conj f1 f2) = go f1 ++ go f2
        go (Disj f1 f2) = [collectDisjs f1 ++ collectDisjs f2]
            where
                collectDisjs (Disj f1' f2') = collectDisjs f1' ++ collectDisjs f2'
                collectDisjs f' = [f']
        go f = [[f]]

-- unifies predicates, e.g. P(x,y) == P(f(a), z)
unifyPredicates :: (Formula, [Formula]) -> (Formula, [Formula]) -> Maybe (([Formula], [Formula]), [(Term, Term)])
unifyPredicates (e1@(Pred p1 ts1), c1) (e2@(Neg (Pred p2 ts2)), c2) | p1 == p2 && length p1 == length p2 = do
        mgu <- unify $ zip ts1 ts2
        return ((c1', c2'), mgu)
    where
        c1' = filter (/= e1) c1
        c2' = filter (/= e2) c2
unifyPredicates _ _ = Nothing

-- 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

-- a single resolution step as described in gloin
rifStep :: [[Formula]] -> Either () [[Formula]]
rifStep clauses | trace (show clauses) True = if [] `elem` clauses then Left () else Right newClauses
    where
        resolveClauses :: [Formula] -> [Formula] -> [(([Formula], [Formula]), [(Term, Term)])]
        resolveClauses c1 c2 = let zippedElems = [((e1, c1), (e2, c2)) | e1 <- c1, e2 <- c2] in mapMaybe (uncurry unifyPredicates) zippedElems
        zippedClauses = [(c1, c2) | c1 <- clauses, c2 <- clauses, c1 /= c2]
        clausesWithMgus = concatMap (uncurry resolveClauses) zippedClauses
        newClauses = clauses ++ map (\((f1, f2), mgu) -> map (`applyMgu` mgu) f1 ++ map (`applyMgu` mgu) f2) clausesWithMgus
rifStep _ = undefined

-- do resolution until we have proven unfulfillability of formula set
doResolution :: [[Formula]] -> Either () [[Formula]]
doResolution f= do
    f' <- rifStep f
    -- TODO after every resolution step make variable names of clauses disjunct
    doResolution f'

{-
To prove a formula we:
1. Construct `Neg phi`
2. Transform `Neg phi` to a clause list
3. doResolution on clause list
-}
proveFormula :: Formula -> Either () [[Formula]]
proveFormula form = doResolution $ makeCNFList (Neg form)

-- unification examples
terma1 :: Term
terma1 = Fun "f" [Var "x", Fun "g" [Var "y"]]
termb1 :: Term
termb1 = Fun "f" [Fun "g" [Var "z"], Var "z"]
terma2 :: Term
terma2 = Fun "f" [Var "x", Fun "g" [Var "x"], Fun "h" [Var "y"]]
termb2 :: Term
termb2 = Fun "f" [Fun "k" [Var "y"], Fun "g" [Var "z"], Var "z"]
terma3 :: Term
terma3 = Fun "f" [Var "x", Fun "g" [Var "x"]]
termb3 :: Term
termb3 = Fun "f" [Var "z", Var "z"]

-- NNF example from gloin
formula1 :: Formula
formula1 = Neg (Conj (Disj (Pred "A" []) (Neg $ Pred "B" [])) (Pred "C" []))

-- PNF and skolem example from gloin
formula2 :: Formula
formula2 = All "x" $ Impl (All "y" $ Pred "L" [Var "y", Var "x"]) (Exists "y" $ Pred "M" [Var "x", Var "y"])

-- Resolution example from gloin script
formula3 :: Formula
formula3 = Impl (Conj (Pred "P" [Fun "a" []]) (All "x" $ Impl (Pred "P" [Var "x"]) (Pred "P" [Fun "f" [Var "x"]]))) (Exists "x" $ Pred "P" [Fun "f" [Fun "f" [Var "x"]]])

-- Resolution example from gloin exercises
formula4 :: Formula
formula4 = Conj (Disj (Pred "P" [Fun "f" [Var "x"], Var "y"]) (Disj (Pred "S" [Var "y", Var "z"]) (Pred "P" [Var "y"]))) (Conj (Neg $ Pred "S" [Fun "f" [Fun "f" [Var "x"]], Var "x"]) (Neg $ Pred "P" [Fun "f" [Var "z"]]))

-- now a big example, sheet 11, exercise 6: Drogenschmuggel, this doesn't work yet but I'm sure its just the exercise thats wrong...
phi1 :: Formula
phi1 = All "x" $ Impl (Conj (Pred "E" [Var "x"]) (Neg $ Pred "I" [Var "x"])) (Exists "y" $ Conj (Pred "Z" [Var "y"]) (Pred "S" [Var "y", Var "x"]))
phi2 :: Formula
phi2 = Exists "x" $ Conj (Conj (Pred "D" [Var "x"]) (Pred "E" [Var "x"])) (All "y" $ Impl (Pred "S" [Var "y", Var "x"]) (Pred "D" [Var "y"]))
phi3 :: Formula
phi3 = All "x" $ Impl (Pred "I" [Var "x"]) (Neg $ Pred "D" [Var "x"])
psi' :: Formula
psi' = Exists "x" $ Conj (Pred "Z" [Var "x"]) (Pred "D" [Var "x"])

formula5 :: Formula
formula5 = Impl (Conj (Conj phi1 phi2) phi3) psi'

-- exercise 2: Ärzte und Quacksalber
formula6 :: [[Formula]]
formula6 = [
    [Neg $ Pred "D" [Var "x1"], Pred "L" [Fun "f" [Var "x1"], Var "x1"]], 
    [Pred "P" [Fun "f" [Var "x2"]], Neg $ Pred "L" [Fun "f" [Var "x2"], Var "x2"]],
    [Neg $ Pred "P" [Var "x3"], Neg $ Pred "Q" [Var "y3"], Neg $ Pred "L" [Var "x3", Var "y3"]],
    [Pred "D" [Fun "a" []]],
    [Pred "Q" [Fun "a" []]]]

-- Drogenschmuggel but already as clauses
formula7 :: [[Formula]]
formula7 = [
    [Neg $ Pred "E" [Var "x2"], Pred "I" [Var "x2"], Pred "Z" [Fun "f" [Var "x2"]]],
    [Neg $ Pred "E" [Var "x3"], Pred "I" [Var "x3"], Pred "S" [Var "x3", Fun "f" [Var "x3"]]],
    [Pred "D" [Fun "c" []]],
    [Pred "E" [Fun "c" []]],
    [Neg $ Pred "S" [Fun "c" [], Var "y"], Pred "D" [Var "y"]],
    [Neg $ Pred "I" [Var "x4"], Neg $ Pred "D" [Var "x4"]],
    [Neg $ Pred "Z" [Var "x5"], Neg $ Pred "D" [Var "x5"]]
    ]

main :: IO ()
main = do
    putStrLn $ "Now making NNF of formula: " ++ show formula1
    print $ makeNNF formula1
    putStrLn $ "Now making PNF of formula: " ++ show formula2
    print $ makePNF formula2
    putStrLn $ "Now making Skolemform of formula: " ++ show formula2
    print $ makeSkolem formula2
    putStrLn $ "Now proving formula by resolution: " ++ show formula3
    case proveFormula formula3 of
        Left _ -> putStrLn "Success!"
        Right _ -> return ()
    putStrLn $ "Now Proving formula by resolution: " ++ show formula4
    case doResolution $ makeCNFList formula4 of
        Left _ -> putStrLn "Success!"
        Right _ -> return ()
    putStrLn $ "Now Proving formula by resolution: " ++ show formula6
    case doResolution formula6 of
        Left _ -> putStrLn "Success!"
        Right _ -> return ()