{-# LANGUAGE OverloadedStrings, StandaloneDeriving #-}

import Chiou.FirstOrderLogic (Sentence(..), Term(..), Quantifier(..), Connective(..))
import Chiou.Monad (ProverT, runProverT)
import Chiou.NormalForm (ImplicativeNormalForm(..), NormalSentence(..), ConjunctiveNormalForm(..), distributeAndOverOr)
import Chiou.KnowledgeBase (askKB, loadKB, showKB, theoremKB)
import Chiou.Resolution (SetOfSupport)
import Control.Monad.Identity (runIdentity)
import Control.Monad.Trans (MonadIO, liftIO)
import Data.Either (rights)
import Data.List (intercalate)
import Data.String (IsString(..))
import Test.HUnit
import System.Exit

newtype V = V String deriving (Eq, Show)

instance IsString V where
    fromString = V

-- |A newtype for the Primitive Predicate parameter.
newtype Pr = Pr String deriving (Eq, Show)

instance IsString Pr where
    fromString = Pr

data AtomicFunction
    = Fn String
    | Skolem Int
    deriving (Eq, Show)

instance Enum AtomicFunction where
    toEnum n = Skolem n
    fromEnum (Skolem n) = n
    fromEnum (Fn _) = error "Enum Chiou.AtomicFunction"

instance IsString AtomicFunction where
    fromString = Fn

main :: IO ()
main = runTestTT (TestList [loadTest, distributeTest, proofTest1, proofTest2]) >>= \ counts ->
       exitWith (if errors counts /= 0 || failures counts /= 0 then ExitFailure 1 else ExitSuccess)

loadTest :: Test
loadTest =
    TestCase (assertEqual "loadKB test" expected (runIdentity (runProverT (loadKB sentences >>= return . map snd))))
    where
      expected :: [Maybe [ImplicativeNormalForm V Pr AtomicFunction]]
      expected = [Just [INF [] [NFPredicate (Pr "Dog") [Constant (Skolem 1)]],INF [] [NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Constant (Skolem 1)]]],
                  Just [INF [NFPredicate (Pr "Dog") [Variable (V "y")],NFPredicate (Pr "Owns") [Variable (V "x"),Variable (V "y")]] [NFPredicate (Pr "AnimalLover") [Variable (V "x")]]],
                  Just [INF [NFPredicate (Pr "AnimalLover") [Variable (V "x")],NFPredicate (Pr "Animal") [Variable (V "y")],NFPredicate (Pr "Kills") [Variable (V "x"),Variable (V "y")]] []],
                  Just [INF [] [NFPredicate (Pr "Kills") [Constant (Fn "Jack"),Constant (Fn "Tuna")],NFPredicate (Pr "Kills") [Constant (Fn "Curiosity"),Constant (Fn "Tuna")]]],
                  Just [INF [] [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")]]],
                  Just [INF [NFPredicate (Pr "Cat") [Variable (V "x")]] [NFPredicate (Pr "Animal") [Variable (V "x")]]]]
      
{-
testLoad :: Monad m => [(Sentence V Pr AtomicFunction, Maybe [ImplicativeNormalForm V Pr AtomicFunction])] -> ProverT V Pr AtomicFunction m ()
testLoad ss =
    if actual /= expected
    then error ("Expected:\n " ++ show expected ++ "\nActual:\n " ++ show actual)
    else return () -- liftIO (putStrLn "testLoad ok")
    where
      actual :: [Maybe [ImplicativeNormalForm V Pr AtomicFunction]]
      actual = map snd ss
-}

distributeTest :: Test
distributeTest =
    TestCase
    (assertEqual "distribute test" 

     (Connective
      (Connective 
       (Connective (Not (Predicate (Pr "q") [Variable (V "x"),Variable (V "y")])) Or (Connective (Not (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")])) Or (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")])))
       And
       (Connective (Not (Predicate (Pr "q") [Variable (V "x"),Variable (V "y")])) Or (Connective (Not (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")])) Or (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")]))))
      And 
      (Connective
       (Connective (Connective (Connective (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")]) Or (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")])) Or (Predicate (Pr "q") [Variable (V "x"),Variable (V "y")]))
        And
        (Connective (Connective (Not (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")])) Or (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")])) Or (Predicate (Pr "q") [Variable (V "x"),Variable (V "y")])))
       And
       (Connective
        (Connective (Connective (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")]) Or (Not (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")]))) Or (Predicate (Pr "q") [Variable (V "x"),Variable (V "y")]))
        And
        (Connective (Connective (Not (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")])) Or (Not (Predicate (Pr "f") [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")]))) Or (Predicate (Pr "q") [Variable (V "x"),Variable (V "y")])))))

     (distributeAndOverOr
      (Connective
       (Connective
        (Not (Predicate "q" [Variable (V "x"),Variable (V "y")]))
        Or
        (Connective
         (Connective
          (Not (Predicate "f" [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")]))
          Or
          (Predicate "f" [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")]))
         And
         (Connective
          (Not (Predicate "f" [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")]))
          Or
          (Predicate "f" [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")]))))
       And
       (Connective
        (Connective
         (Connective
          (Predicate "f" [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")])
          And
          (Not (Predicate "f" [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")])))
         Or
         (Connective
          (Predicate "f" [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "y")])
          And
          (Not (Predicate "f" [SkolemFunction 1 [Variable (V "x"),Variable (V "x"),Variable (V "y"),Variable (V "z")],Variable (V "x")]))))
        Or
        (Predicate "q" [Variable (V "x"),Variable (V "y")])) :: Sentence V Pr AtomicFunction)))

proofTest1 :: Test
proofTest1 = TestCase (assertEqual "proof test 1" proof1 (runIdentity (runProverT (loadKB sentences >> theoremKB (Predicate "Kills" [Constant (Fn "Jack"), Constant (Fn "Tuna")])))))

proof1 :: (Bool, SetOfSupport V Pr AtomicFunction {-[ImplicativeNormalForm V Pr AtomicFunction]-})
proof1 = ( False
         , [(INF [NFPredicate (Pr "Kills") [Constant (Fn "Jack"),Constant (Fn "Tuna")]] [],[]),(INF [] [NFPredicate (Pr "Kills") [Constant (Fn "Curiosity"),Constant (Fn "Tuna")]],[]),(INF [NFPredicate (Pr "Animal") [Constant (Fn "Tuna")],NFPredicate (Pr "AnimalLover") [Constant (Fn "Curiosity")]] [],[]),(INF [NFPredicate (Pr "Dog") [Variable (V "y")],NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Variable (V "y")],NFPredicate (Pr "Animal") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")],NFPredicate (Pr "AnimalLover") [Constant (Fn "Curiosity")]] [],[]),(INF [NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Constant (Skolem 1)],NFPredicate (Pr "Animal") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")],NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Variable (V "y")],NFPredicate (Pr "Dog") [Variable (V "y")]] [],[]),(INF [NFPredicate (Pr "Dog") [Variable (V "y")],NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Variable (V "y")],NFPredicate (Pr "Cat") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "AnimalLover") [Constant (Fn "Curiosity")]] [],[]),(INF [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")],NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Constant (Skolem 1)]] [],[]),(INF [NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Constant (Skolem 1)],NFPredicate (Pr "Cat") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Variable (V "y")],NFPredicate (Pr "Dog") [Variable (V "y")]] [],[]),(INF [NFPredicate (Pr "Dog") [Variable (V "y")],NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Variable (V "y")]] [],[]),(INF [NFPredicate (Pr "Owns") [Constant (Fn "Curiosity"),Constant (Skolem 1)]] [],[])] )

proofTest2 :: Test
proofTest2 = TestCase (assertEqual "proof test 2" proof2 (runIdentity (runProverT (loadKB sentences >> theoremKB (Predicate "Kills" [Constant (Fn "Curiosity"), Constant (Fn "Tuna")])))))

proof2 :: (Bool, SetOfSupport V Pr AtomicFunction)
proof2 = ( True
         , [(INF [NFPredicate (Pr "Kills") [Constant (Fn "Curiosity"),Constant (Fn "Tuna")]] [],[]),(INF [] [NFPredicate (Pr "Kills") [Constant (Fn "Jack"),Constant (Fn "Tuna")]],[]),(INF [NFPredicate (Pr "Animal") [Constant (Fn "Tuna")],NFPredicate (Pr "AnimalLover") [Constant (Fn "Jack")]] [],[]),(INF [NFPredicate (Pr "Dog") [Variable (V "y")],NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Variable (V "y")],NFPredicate (Pr "Animal") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")],NFPredicate (Pr "AnimalLover") [Constant (Fn "Jack")]] [],[]),(INF [NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Constant (Skolem 1)],NFPredicate (Pr "Animal") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Dog") [Constant (Skolem 1)],NFPredicate (Pr "Animal") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")],NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Variable (V "y")],NFPredicate (Pr "Dog") [Variable (V "y")]] [],[]),(INF [NFPredicate (Pr "Dog") [Variable (V "y")],NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Variable (V "y")],NFPredicate (Pr "Cat") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "AnimalLover") [Constant (Fn "Jack")]] [],[]),(INF [NFPredicate (Pr "Animal") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")],NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Constant (Skolem 1)]] [],[]),(INF [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")],NFPredicate (Pr "Dog") [Constant (Skolem 1)]] [],[]),(INF [NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Constant (Skolem 1)],NFPredicate (Pr "Cat") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Variable (V "y")],NFPredicate (Pr "Dog") [Variable (V "y")]] [],[]),(INF [NFPredicate (Pr "Dog") [Constant (Skolem 1)],NFPredicate (Pr "Cat") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Dog") [Variable (V "y")],NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Variable (V "y")]] [],[]),(INF [NFPredicate (Pr "Cat") [Constant (Fn "Tuna")]] [],[]),(INF [NFPredicate (Pr "Owns") [Constant (Fn "Jack"),Constant (Skolem 1)]] [],[]),(INF [NFPredicate (Pr "Dog") [Constant (Skolem 1)]] [],[]),(INF [] [],[])] )

testProof :: MonadIO m => String -> (Sentence V Pr AtomicFunction, Bool, [ImplicativeNormalForm V Pr AtomicFunction]) -> ProverT V Pr AtomicFunction m ()
testProof label (question, expectedAnswer, expectedProof) =
    theoremKB question >>= \ (actualFlag, actualProof) ->
    let actual' = (actualFlag, map fst actualProof) in
    if actual' /= (expectedAnswer, expectedProof)
    then error ("\n Expected:\n  " ++ show (expectedAnswer, expectedProof) ++
                "\n Actual:\n  " ++ show actual')
    else liftIO (putStrLn (label ++ " ok"))
    where
      format (flag, proof) = "(" ++ show flag ++ ",\n   [" ++ intercalate "\n    " (map show proof) ++ "])"

loadCmd :: Monad m => ProverT V Pr AtomicFunction m [(Sentence V Pr AtomicFunction, Maybe [ImplicativeNormalForm V Pr AtomicFunction])]
loadCmd = loadKB sentences

sentences :: [Sentence V Pr AtomicFunction]
sentences = [Quantifier Exists ["x"] (Connective (Predicate "Dog" [Variable "x"]) And (Predicate "Owns" [Constant "Jack", Variable "x"])),
             Quantifier ForAll ["x"] (Connective (Connective (Quantifier Exists ["y"] (Predicate "Dog" [Variable "y"])) And
                                                             (Predicate "Owns" [Variable "x", Variable "y"])) Imply
                                                 (Predicate "AnimalLover" [Variable "x"])),

             Quantifier ForAll ["x"] (Connective (Predicate "AnimalLover" [Variable "x"]) Imply
                                                 (Quantifier ForAll ["y"] (Connective (Predicate "Animal" [Variable "y"]) Imply
                                                                                      (Not (Predicate "Kills" [Variable "x", Variable "y"]))))),

             Connective (Predicate "Kills" [Constant "Jack", Constant "Tuna"]) Or
                        (Predicate "Kills" [Constant "Curiosity", Constant "Tuna"]),
             Predicate "Cat" [Constant "Tuna"],
             Quantifier ForAll ["x"] (Connective (Predicate "Cat" [Variable "x"]) Imply (Predicate "Animal" [Variable "x"]))]

{-
From Prover.hs:

tellCmd :: (Monad m, Eq v, Eq p, Eq f, IsString p, Enum f) => Sentence v p f -> ProverT v p f m [ImplicativeNormalForm v p f]
tellCmd = tellKB


askCmd :: (MonadIO m, Eq v, Eq p, Eq f, IsString p, Show v, Show p, Show f) =>
          Sentence v p f -> ProverT v p f m (Bool, SetOfSupport v p f)
askCmd sentence = 
    liftIO (putStrLn ("askCmd " ++ show sentence)) >>
    askKB sentence >>= \ result ->
    liftIO (putStrLn (format result)) >>
    return result
    where format (flag, proof) = "(" ++ show flag ++
                                 ",\n   [" ++ intercalate "\n    " (map show proof) ++ "])"

deleteCmd :: Monad m => Int -> ProverT v p f m String
deleteCmd i = deleteKB i

unloadCmd :: Monad m => ProverT v p f m ()
unloadCmd = unloadKB

resetCmd :: Monad m => ProverT v p f m ()
resetCmd = emptyKB

showCmd :: (MonadIO m, Show v, Show p, Show f) => ProverT v p f m ()
showCmd = showKB >>= liftIO . putStrLn

debugCmd :: (Eq v, IsString p, Show v, Show p, Show f) => Sentence v p f -> String
debugCmd sentence = showINFDerivation sentence
-}

{---------------------}
{-- Pretty Printing --}
{---------------------}

{-
flattenListString :: String -> String
flattenListString s = filter (\c -> c /= '[' && c /= ']' && c /= '\"') s

instance (Show v, Show p, Show f) => Show (Sentence v p f) where
    show (Connective s1 c s2) = "(" ++ show s1 ++ show c ++ show s2 ++ ")"
    show (Quantifier q [] s) = show q ++ " " ++ show s
    show (Quantifier q vs s) = show q ++ " " ++ (flattenListString (show vs)) ++
			       " " ++ show s
    show (Not s) = "NOT " ++ "(" ++ show s ++ ")"
    show (Predicate p []) = show p
    show (Predicate p ts) = show p ++ "(" ++ (flattenListString (show ts)) ++ ")"
    show (Equal t1 t2) = show t1 ++ "=" ++ show t2

instance (Show v, Show f) => Show (Term v f) where
    show (Function f ts) = show f ++ "(" ++ (flattenListString (show ts)) ++ ")"
    show (Constant c) = show c
    show (Variable v) = show v
    show (SkolemFunction i ts) = "SKOLEM_FUNCTION_" ++ show i ++ "(" ++
				 (flattenListString (show ts)) ++ ")"
    show (SkolemConstant i) = "SKOLEM_CONSTANT_" ++ show i

instance Show Connective where
    show (Imply) = " => "
    show (Equiv) = " <=> "
    show (And)   = " AND "
    show (Or)    = " OR "

instance Show Quantifier where
    show (ForAll) = "ForAll"
    show (Exists) = "Exists"

{--------------------}
{-- Pretty Printing -}
{--------------------}

connectNormalSentence :: (Show v, Show p, Show f) => String -> [NormalSentence v p f] -> String
connectNormalSentence s [] = case s of
			       " AND " -> "True"
			       " OR "  -> "False"
                               _ -> error ("unknown connective: " ++ show s)
connectNormalSentence _s (x:[]) = show x
connectNormalSentence s (x:xs) = show x ++ s ++ connectNormalSentence s xs

instance (Show v, Show p, Show f) => Show (ConjunctiveNormalForm v p f) where
    show (CNF xs) = connectNormalSentence "\n" xs

instance (Show v, Show p, Show f) => Show (ImplicativeNormalForm v p f) where
    show (INF ls rs) = connectNormalSentence " AND " ls ++ " => " ++
		       connectNormalSentence " OR " rs

instance (Show v, Show p, Show f) => Show (NormalSentence v p f) where
    show (NFNot s) = "NOT " ++ "(" ++ show s ++ ")"
    show (NFPredicate p []) = show p
    show (NFPredicate p ts) = show p ++ "(" ++ (flattenListString (show ts)) ++ ")"
    show (NFEqual t1 t2) = show t1 ++ "=" ++ show t2

--reportSetOfSupport :: SetOfSupport v p f -> String
--reportSetOfSupport [] = ""
--reportSetOfSupport ((inf, _theta):xs) = show inf ++ ('\n':reportSetOfSupport xs)
-}

-- Ugly Printing

deriving instance (Show v, Show p, Show f) => Show (Sentence v p f)
deriving instance (Show v, Show p, Show f) => Show (ImplicativeNormalForm v p f)
deriving instance (Show v, Show p, Show f) => Show (NormalSentence v p f)
deriving instance (Show v, Show f) => Show (Term v f)
deriving instance Show Quantifier
deriving instance Show Connective