module Main where

import Data.Char
import Control.Monad
import qualified Data.Map as Map
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Text.ParserCombinators.Parsec as P

-- --------------------------------------------------------------------

type Var = Int

data Term
  = Var !Int
  | Str String [Term]
  deriving (Eq, Show)

type Subst = IM.IntMap Term

nullSubst :: Subst
nullSubst = IM.empty

class Substitutable t where
  apply :: Subst -> t -> t
  vars  :: t -> IS.IntSet

instance Substitutable Term where
  apply s (Var u)  =
      case IM.lookup u s of
        Just t  -> t
        Nothing -> Var u
  apply s (Str f xs) = Str f (apply s xs)
  vars (Var u) = IS.singleton u
  vars (Str f xs) = vars xs

instance Substitutable a => Substitutable [a] where
  apply s = map (apply s)
  vars = IS.unions . map vars

(@@) :: Subst -> Subst -> Subst
s1 @@ s2 = s1 `IM.union` fmap (apply s1) s2

mgu :: MonadPlus m => Term -> Term -> m Subst
mgu (Str f xs) (Str g ys) | f==g = mguList xs ys
mgu (Var u) t = bind u t
mgu t (Var u) = bind u t
mgu t1 t2 = mzero

mguList :: MonadPlus m => [Term] -> [Term] -> m Subst
mguList [] [] = return nullSubst
mguList (x:xs) (y:ys) = do
  s1 <- mgu x y
  s2 <- mguList (apply s1 xs) (apply s1 ys)
  return (s2 @@ s1)
mguList _ _ = mzero

bind :: MonadPlus m => Var -> Term -> m Subst
bind u t
  | t == Var u = return nullSubst
  | u `IS.member` vars t = mzero
  | otherwise = return (IM.singleton u t)

-- --------------------------------------------------------------------

newtype PL a = PL{ runPL :: Subst -> Int -> [(a, Subst, Int)] }

instance Monad PL where
    return x = PL $ \s n -> [(x, s, n)]
    m >>= f = PL $ \s n -> concat [ runPL (f x) s n' | (x, s, n') <- runPL m s n ]
    fail _ = mzero
instance MonadPlus PL where
    mzero = PL $ \_ _ -> []
    mplus m1 m2 = PL $ \s n -> runPL m1 s n ++ runPL m2 s n

fresh :: PL Var
fresh = PL $ \s n -> [(n, s, n+1)]

getSubst :: PL Subst
getSubst = PL $ \s n -> [(s, s, n)]

extSubst :: Subst -> PL ()
extSubst s2 = PL $ \s1 n -> [((), s2 @@ s1, n)]

unify :: Term -> Term -> PL ()
unify t1 t2 = do
  s1 <- getSubst
  s2 <- mgu (apply s1 t1) (apply s1 t2)
  extSubst s2

instantiate :: Substitutable t => t -> PL t
instantiate t = do
  s <- liftM IM.fromList $ forM (IS.toList (vars t)) $ \v -> do
    v2 <- fresh
    return (v, Var v2)
  return $ apply s t

data Rule = Term :- [Term] deriving (Show)
type Query = [Term]

instance Substitutable Rule where
  apply s (head :- body) = apply s head :- apply s body
  vars (head :- body) = vars head `IS.union` vars body

run :: [Rule] -> Query -> [Subst]
run rules q = 
    [ IM.filterWithKey (\v _  -> v `IS.member` vs) s
    | let vs = vars q
    , (_, s, _) <- runPL (mapM_ call q) nullSubst 0 ] 
  where    
    call :: Term -> PL ()
    call q = msum [ do (head :- body) <- instantiate rule
                       unify q head
                       mapM_ call body
                  | rule <- rules ]

-- --------------------------------------------------------------------

type Parser = P.GenParser Char PState
type PState = (Map.Map String Var, Int)

emptyPState :: PState
emptyPState = (Map.empty, 0)

program :: Parser [Rule]
program = optspaces >> P.many rule

rule :: Parser Rule
rule = do
  P.setState emptyPState
  head <- term
  body <- P.option [] (P.string ":-" >> optspaces >> terms)
  P.string "." >> optspaces
  return (head :- body)

query :: Parser [Term]
query = do
  body <- terms 
  P.string "." >> optspaces
  return body

term :: Parser Term
term = msum
  [ do c <- P.satisfy isLower
       cs <- P.many (P.satisfy isAlphaNum)
       let f = c:cs
       optspaces
       xs <- P.option [] $
             P.between (P.char '(' >> optspaces) (P.char ')' >> optspaces) $
             terms
       return (Str f xs)
  , do c <- P.satisfy isUpper
       cs <- P.many (P.satisfy isAlphaNum)
       optspaces
       v <- internVar (c:cs)
       return (Var v)
  ]
  where
    internVar :: String -> Parser Var
    internVar s = do
      (m, n) <- P.getState
      case Map.lookup s m of
        Just v -> return v
        Nothing -> do
          P.setState (Map.insert s n m, n+1)
          return n

terms :: Parser [Term]
terms = term `P.sepBy` (P.char ',' >> optspaces)

optspaces :: Parser ()
optspaces = P.optional P.spaces

-- --------------------------------------------------------------------

test = run prog q
  where
    Right q = P.runParser query emptyPState "" "mortal(X)."
    Right prog = P.runParser program emptyPState "" "mortal(X) :- human(X).\nhuman(socrat)."