module Main where

import Data.List
import Data.Maybe
import qualified Data.Set as Set
import FOLSyntax
import Parser

fromRight :: Either a b -> b
fromRight (Left _) = error "fromRight called on left value!"
fromRight (Right b) = b

-- 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 form
    where
        go f = let f' = pnfStep f in
            if f == f' then f else go f'
        -- swapping rules
        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 (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 (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 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 form
    where
        go f = let f' = cnfStep f in
            if f == f' then f else go f'
        cnfStep (Conj phi psi) = Conj (cnfStep phi) (cnfStep psi)
        cnfStep (Disj (Conj phi psi) xi) = Conj (Disj phi xi) (Disj psi xi)
        cnfStep (Disj xi (Conj phi psi)) = Conj (Disj xi phi) (Disj xi psi)
        cnfStep (All _ f) = cnfStep f
        cnfStep (Exists _ f) = cnfStep 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 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

-- replace every fresh variable in a formula with a new name, this is only correct in our resolution context!!
makeVariablesDisjunct :: [[Formula]] -> [[Formula]]
makeVariablesDisjunct form = snd $ foldr (\clause (used', clauses') -> let (used'', clause') = makeClauseDisjunct used' clause in (used'', clause' : clauses')) ([], []) form
makeClauseDisjunct :: [String] -> [Formula] -> ([String], [Formula])
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', map (\f -> renameFormula f oldVar v') clause') 
                    where
                        v' = findFresh used'

-- a single resolution step as described in gloin
-- TODO before each step, rename every clause so that variable names are disjunct!
rifStep :: [[Formula]] -> Either () [[Formula]]
rifStep clauses = if [] `elem` clauses then Left () else Right newClauses
    where
        -- rename all variables
        clauses' = makeVariablesDisjunct clauses
        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 = Set.toList . Set.fromList $ clauses' ++ map (\((f1, f2), mgu) -> map (`applyMgu` mgu) f1 ++ map (`applyMgu` mgu) f2) clausesWithMgus

-- do resolution until we have proven unfulfillability of formula set
doResolution :: [[Formula]] -> Either () [[Formula]]
doResolution f = do
    f' <- rifStep f
    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 . makeCNF . makeSkolem . makePNF . renameBinders . makeNNF $ Neg form

proveFormulaIO :: Formula -> IO ()
proveFormulaIO f = do
    let f' = makeNNF (Neg f)
    print f'
    putStrLn ""
    let f'' = renameBinders f'
    print f''
    putStrLn ""

    let f''' = makePNF f''
    print f'''
    putStrLn ""

    let f'''' = makeSkolem f'''
    print f''''
    putStrLn ""

    let g = makeCNF f''''
    print g
    putStrLn ""

    let g' = makeCNFList g
    print g'
    case doResolution g' of
        Left _ -> putStrLn "Success!!"
        _ -> return ()

-- 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 :: String
formula2 = "forall x. (forall y. L(y,x)) -> exists y.M(x,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 (already in cnf)
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, fifth clause is wrong, arguments to S need to be swapped!!
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"]]
    ]

formula8 :: [[Formula]]
formula8 = [
    [Neg $ Pred "E" [Var "x2"], Pred "I" [Var "x2"], Pred "Z" [Fun "f" [Var "x21", Var "x22", Var "x2"]]],
    [Neg $ Pred "E" [Var "x3"], Pred "I" [Var "x3"], Pred "S" [Var "x3", Fun "f" [Var "x31", Var "x32", Var "x3"]]],
    [Pred "D" [Fun "c" [Var "x5", Var "x4"]]],
    [Pred "E" [Fun "c" [Var "x5'", Var "x4'"]]],
    [Neg $ Pred "S" [Fun "c" [Var "x5''", Var "x4''"], Var "y"], Pred "D" [Var "y"]],
    [Neg $ Pred "I" [Var "x44"], Neg $ Pred "D" [Var "x44"]],
    [Neg $ Pred "Z" [Var "x55"], Neg $ Pred "D" [Var "x55"]]
    ]

-- aerzte und quacksalber v2
formula9 :: Formula 
formula9 = Impl (Conj psi1 psi2) psi3
    where
        psi1 = All "x" $ Impl (Pred "D" [Var "x"]) (Exists "y" $ Conj (Pred "P" [Var "y"]) (Pred "L" [Var "y", Var "x"]))
        psi2 = All "x" $ Impl (Pred "P" [Var "x"]) (All "y" $ Impl (Pred "Q" [Var "y"]) (Neg $ Pred "L" [Var "x", Var "y"]))
        psi3 = All "x" $ Impl (Pred "D" [Var "x"]) (Neg $ Pred "Q" [Var "x"])

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 (fromRight $ parseFormula formula2)
    putStrLn $ "Now making Skolemform of formula: " ++ show formula2
    print $ makeSkolem (fromRight $ parseFormula 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 ()
    putStrLn $ "Now Proving formula by resolution: " ++ show (Neg formula5)
    case proveFormula (Neg formula5) of
        Left _ -> putStrLn "Success!"
        Right _ -> return ()
    putStrLn $ "Now proving formula by resolution: " ++ show formula7
    case doResolution formula7 of
        Left _ -> putStrLn "Success!"
        Right _ -> return ()