content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE FlexibleContexts #-}
  -- | FIXME: Change the adornment/query building process such that
  -- conditional clauses are always processed last.  This is necessary
  -- so that all variables are bound.
  --
  -- FIXME: Add an assertion to say that ConditionalClauses cannot have
  -- Free variables.
  module Database.Datalog.Rules (
    Adornment(..),
    Term(..),
    Clause(..),
    AdornedClause(..),
    Rule(..),
    Literal(..),
    Query(..),
    QueryBuilder,
    PartialTuple(..),
    (|-),
    assertRule,
    relationPredicateFromName,
    inferencePredicate,
    ruleRelations,
    issueQuery,
    runQuery,
    queryToPartialTuple,
    queryPredicate,
    lit,
    negLit,
    cond1,
    cond2,
    cond3,
    cond4,
    cond5,
    bindQuery,
    partitionRules
    ) where
  
  import qualified Control.Monad.Catch as E
  import           Control.Monad.Trans.Class
  import           Control.Monad.Trans.State.Strict
  import           Data.Function ( on )
  import           Data.Hashable
  import           Data.HashMap.Strict ( HashMap )
  import qualified Data.HashMap.Strict as HM
  import           Data.List ( intercalate, groupBy, sortBy )
  import           Data.Maybe ( mapMaybe )
  import           Data.Monoid
  import           Data.Text ( Text )
  import qualified Data.Text as T
  import           Text.Printf ( printf )
  
  import           Prelude
  
  import           Database.Datalog.Adornment
  import           Database.Datalog.Relation
  import           Database.Datalog.Errors
  import           Database.Datalog.Database
  
  -- import Debug.Trace
  -- debug = flip trace
  
  data QueryState a = QueryState { intensionalDatabase :: Database a
                                 , conditionalIdSource :: Int
                                 , queryRules :: [(Clause a, [Literal Clause a])]
                                 }
  
  -- | The Monad in which queries are constructed and rules are declared
  type QueryBuilder m a = StateT (QueryState a) m
  
  nextConditionalId :: (E.MonadThrow m) => QueryBuilder m a Int
  nextConditionalId = do
    s <- get
    let cid = conditionalIdSource s
    put s { conditionalIdSource = cid + 1 }
    return cid
  
  data Term a = LogicVar !Text
                -- ^ A basic logic variable.  Equality is based on the
                -- variable name.
              | BindVar !Text
                -- ^ A special variable available in queries that can be
                -- bound at query execution time
              | Anything
                -- ^ A term that is allowed to take any value (this is
                -- sugar for a fresh logic variable)
              | Atom a
                -- ^ A user-provided literal from the domain a
              | FreshVar !Int
                -- ^ A fresh logic variable, generated internally for
                -- each Anything occurrence.  Not exposed to the user
  
  instance (Show a) => Show (Term a) where
    show (LogicVar t) = T.unpack t
    show (BindVar t) = "??" ++ T.unpack t
    show (Atom a) = show a
    show Anything = "*"
    show (FreshVar _) = "*"
  
  instance (Hashable a) => Hashable (Term a) where
    hashWithSalt s (LogicVar t) =
      s `hashWithSalt` t `hashWithSalt` (1 :: Int)
    hashWithSalt s (BindVar t) =
      s `hashWithSalt` t `hashWithSalt` (2 :: Int)
    hashWithSalt s (Atom a) = s `hashWithSalt` a
    hashWithSalt s Anything = s `hashWithSalt` (99 :: Int)
    hashWithSalt s (FreshVar i) =
      s `hashWithSalt` i `hashWithSalt` (22 :: Int)
  
  instance (Eq a) => Eq (Term a) where
    (LogicVar t1) == (LogicVar t2) = t1 == t2
    (BindVar t1) == (BindVar t2) = t1 == t2
    (Atom a1) == (Atom a2) = a1 == a2
    Anything == Anything = True
    FreshVar i1 == FreshVar i2 = i1 == i2
    _ == _ = False
  
  data Clause a = Clause { clauseRelation :: Relation
                         , clauseTerms :: [Term a]
                         }
  
  instance (Eq a) => Eq (Clause a) where
    (Clause r1 ts1) == (Clause r2 ts2) = r1 == r2 && ts1 == ts2
  
  instance (Show a) => Show (Clause a) where
    show (Clause p ts) =
      printf "%s(%s)" (show p) (intercalate ", " (map show ts))
  
  
  data AdornedClause a = AdornedClause { adornedClauseRelation :: Relation
                                       , adornedClauseTerms :: [(Term a, Adornment)]
                                       }
  
  instance (Eq a) => Eq (AdornedClause a) where
    (AdornedClause r1 cs1) == (AdornedClause r2 cs2) = r1 == r2 && cs1 == cs2
  
  instance (Hashable a) => Hashable (AdornedClause a) where
    hashWithSalt s (AdornedClause r ts) =
      s `hashWithSalt` r `hashWithSalt` ts
  
  instance (Show a) => Show (AdornedClause a) where
    show (AdornedClause p ats) =
      printf "%s(%s)" (show p) (intercalate ", " (map showAT ats))
      where
        showAT (t, a) = printf "%s[%s]" (show t) (show a)
  
  -- | Body clauses can be normal clauses, negated clauses, or
  -- conditionals.  Conditionals are arbitrary-arity (via a list)
  -- functions over literals and logic variables.
  data Literal ctype a = Literal (ctype a)
                       | NegatedLiteral (ctype a)
                       | ConditionalClause Int ([a] -> Bool) [Term a] (HashMap (Term a) Int)
  
  -- | This equality instance is complicated because conditional clauses
  -- contain functions.  We assign a unique id at conditional clause
  -- creation time so they have identity and we can compare on that.
  -- Rules from different queries cannot be compared safely, but that
  -- shouldn't be a problem because there isn't really a way to sneak a
  -- rule reference out of a query.  If there is a shady way to do so,
  -- don't because it will be bad.
  instance (Eq a, Eq (ctype a)) => Eq (Literal ctype a) where
    (Literal c1) == (Literal c2) = c1 == c2
    (NegatedLiteral c1) == (NegatedLiteral c2) = c1 == c2
    (ConditionalClause cid1 _ _ _) == (ConditionalClause cid2 _ _ _) = cid1 == cid2
    _ == _ = False
  
  instance (Hashable a, Hashable (ctype a)) => Hashable (Literal ctype a) where
    hashWithSalt s (Literal c) =
      s `hashWithSalt` c `hashWithSalt` (1 :: Int)
    hashWithSalt s (NegatedLiteral c) =
      s `hashWithSalt` c `hashWithSalt` (2 :: Int)
    hashWithSalt s (ConditionalClause cid _ ts vm) =
      s `hashWithSalt` cid `hashWithSalt` ts `hashWithSalt` HM.size vm
  
  lit :: (E.MonadThrow m) => Relation -> [Term a] -> QueryBuilder m a (Literal Clause a)
  lit p ts = return $ Literal $ Clause p ts
  
  negLit :: (E.MonadThrow m) => Relation -> [Term a] -> QueryBuilder m a (Literal Clause a)
  negLit p ts = return $ NegatedLiteral $ Clause p ts
  
  cond1 :: (E.MonadThrow m, Eq a, Hashable a)
           => (a -> Bool)
           -> Term a
           -> QueryBuilder m a (Literal Clause a)
  cond1 p t = do
    cid <- nextConditionalId
    return $ ConditionalClause cid (\[x] -> p x) [t] mempty
  
  cond2 :: (E.MonadThrow m, Eq a, Hashable a)
           => (a -> a -> Bool)
           -> (Term a, Term a)
           -> QueryBuilder m a (Literal Clause a)
  cond2 p (t1, t2) = do
    cid <- nextConditionalId
    return $ ConditionalClause cid (\[x1, x2] -> p x1 x2) [t1, t2] mempty
  
  
  cond3 :: (E.MonadThrow m, Eq a, Hashable a)
           => (a -> a -> a -> Bool)
           -> (Term a, Term a, Term a)
           -> QueryBuilder m a (Literal Clause a)
  cond3 p (t1, t2, t3) = do
    cid <- nextConditionalId
    return $ ConditionalClause cid (\[x1, x2, x3] -> p x1 x2 x3) [t1, t2, t3] mempty
  
  cond4 :: (E.MonadThrow m, Eq a, Hashable a)
           => (a -> a -> a -> a -> Bool)
           -> (Term a, Term a, Term a, Term a)
           -> QueryBuilder m a (Literal Clause a)
  cond4 p (t1, t2, t3, t4) = do
    cid <- nextConditionalId
    return $ ConditionalClause cid (\[x1, x2, x3, x4] -> p x1 x2 x3 x4) [t1, t2, t3, t4] mempty
  
  cond5 :: (E.MonadThrow m, Eq a, Hashable a)
           => (a -> a -> a -> a -> a -> Bool)
           -> (Term a, Term a, Term a, Term a, Term a)
           -> QueryBuilder m a (Literal Clause a)
  cond5 p (t1, t2, t3, t4, t5) = do
    cid <- nextConditionalId
    return $ ConditionalClause cid (\[x1, x2, x3, x4, x5] -> p x1 x2 x3 x4 x5) [t1, t2, t3, t4, t5] mempty
  
  instance (Show a, Show (ctype a)) => Show (Literal ctype a) where
    show (Literal c) = show c
    show (NegatedLiteral c) = '~' : show c
    show (ConditionalClause _ _ ts _) = printf "f(%s)" (show ts)
  
  -- | A rule has a head and body clauses.  Body clauses can be normal
  -- clauses, negated clauses, or conditionals.
  data Rule a = Rule { ruleHead :: AdornedClause a
                     , ruleBody :: [Literal AdornedClause a]
                     , ruleVariableMap :: HashMap (Term a) Int
                     }
  
  instance (Show a) => Show (Rule a) where
    show (Rule h b _) = printf "%s |- %s" (show h) (intercalate ", " (map show b))
  
  instance (Eq a) => Eq (Rule a) where
    (Rule h1 b1 vms1) == (Rule h2 b2 vms2) =
      h1 == h2 && b1 == b2 && vms1 == vms2
  
  instance (Hashable a) => Hashable (Rule a) where
    hashWithSalt s (Rule h b vms) =
      s `hashWithSalt` h `hashWithSalt` b `hashWithSalt` HM.size vms
  
  newtype Query a = Query { unQuery :: Clause a }
                  deriving (Show)
  infixr 0 |-
  
  -- | Assert a rule
  --
  -- FIXME: Check to make sure that clause arities match their declared
  -- schema.
  (|-), assertRule :: (E.MonadThrow m)
          => (Relation, [Term a]) -- ^ The rule head
          -> [QueryBuilder m a (Literal Clause a)] -- ^ Body literals
          -> QueryBuilder m a ()
  (|-) = assertRule
  assertRule (p, ts) b = do
    -- FIXME: Assert that Anything does not appear in the head terms
    -- (that is a range restriction violation).  Also check the range
    -- restriction here.
    b' <- sequence b
    let h = Clause p ts
        b'' = fst $ foldr freshenVars ([], [0..]) b'
    s <- get
    put s { queryRules = (h, b'') : queryRules s }
  
  -- | Replace all instances of Anything with a FreshVar with a unique
  -- (to the rule) index.  This lets later evaluation stages ignore
  -- these and just deal with clean FreshVars.
  freshenVars :: Literal Clause a
                 -> ([Literal Clause a], [Int])
                 -> ([Literal Clause a], [Int])
  freshenVars l (cs, ixSrc) =
    case l of
      ConditionalClause _ _ _ _ -> (l : cs, ixSrc)
      Literal (Clause h ts) ->
        let (ts', ixRest) = foldr freshen ([], ixSrc) ts
        in (Literal (Clause h ts') : cs, ixRest)
      NegatedLiteral (Clause h ts) ->
        let (ts', ixRest) = foldr freshen ([], ixSrc) ts
        in (NegatedLiteral (Clause h ts') : cs, ixRest)
    where
      freshen t (ts, src) =
        case t of
          Anything -> (FreshVar (head src) : ts, tail src)
          _ -> (t : ts, src)
  
  -- FIXME: Unify these and require inferred relations to be declared in
  -- the schema at db construction time.  That also gives an opportunity
  -- to name the columns
  
  -- | Retrieve a Relation handle from the IDB.  If the Relation does
  -- not exist, an error will be raised.
  relationPredicateFromName :: (E.MonadThrow m)
                               => Text
                               -> QueryBuilder m a Relation
  relationPredicateFromName name = do
    let rel = Relation name
    idb <- gets intensionalDatabase
    case rel `elem` databaseRelations idb of
      False -> lift $ E.throwM (NoRelationError rel)
      True -> return rel
  
  -- | Create a new predicate that will be referenced by an EDB rule
  inferencePredicate :: (E.MonadThrow m)
                        => Text
                        -> QueryBuilder m a Relation
  inferencePredicate = return . Relation
  
  -- | A partial tuple records the atoms in a tuple (with their indices
  -- in the tuple).  These are primarily used in database queries.
  newtype PartialTuple a = PartialTuple [Maybe a]
  
  instance (Show a) => Show (PartialTuple a) where
    show (PartialTuple vs) = show $ map show vs
  
  -- | Convert a 'Query' into a 'PartialTuple' that can be used in a
  -- 'select' of the IDB
  queryToPartialTuple :: Query a -> PartialTuple a
  queryToPartialTuple (Query c) =
    PartialTuple $! map takeAtom ts
    where
      ts = clauseTerms c
      takeAtom t =
        case t of
          Atom a -> Just a
          _ -> Nothing
  
  
  
  literalClauseRelation :: Literal AdornedClause a -> Maybe Relation
  literalClauseRelation bc =
    case bc of
      Literal c -> Just $ adornedClauseRelation c
      NegatedLiteral c -> Just $ adornedClauseRelation c
      _ -> Nothing
  
  ruleRelations :: Rule a -> [Relation]
  ruleRelations (Rule h bcs _) = adornedClauseRelation h : mapMaybe literalClauseRelation bcs
  
  -- | Turn a Clause into a Query.  This is meant to be the last
  -- statement in a QueryBuilder monad.
  issueQuery :: (E.MonadThrow m) => Relation -> [Term a] -> QueryBuilder m a (Query a)
  issueQuery r ts = return $ Query $ Clause r ts
  
  
  -- | Run the QueryBuilder action to build a query and initial rule
  -- database
  --
  -- Rules are adorned (marking each variable as Free or Bound as they
  -- appear) before being returned.
  runQuery :: (E.MonadThrow m, Eq a, Hashable a)
              => QueryBuilder m a (Query a) -> Database a -> m (Query a, [(Clause a, [Literal Clause a])])
  runQuery qm idb = do
    (q, QueryState _ _ rs) <- runStateT qm (QueryState idb 0 [])
    return (q, rs)
  
  -- | Group rules by their head relations.  This is needed to perform
  -- semi-naive evaluation easily.
  partitionRules :: [Rule a] -> [[Rule a]]
  partitionRules = groupBy gcmp . sortBy scmp
    where
      scmp = compare `on` (adornedClauseRelation . ruleHead)
      gcmp = (==) `on` (adornedClauseRelation . ruleHead)
  
  queryPredicate :: Query a -> Relation
  queryPredicate = clauseRelation . unQuery
  
  -- | Apply bindings to a query
  bindQuery :: Query a -> [(Text, a)] -> Query a
  bindQuery (Query (Clause r ts)) bs =
    Query $ Clause r $ foldr applyBinding [] ts
    where
      applyBinding t acc =
        case t of
          LogicVar _ -> t : acc
          BindVar name ->
            case lookup name bs of
              Nothing -> error ("No binding provided for BindVar " ++ show name)
              Just b -> Atom b : acc
          Anything -> t : acc
          Atom _ -> t : acc
          FreshVar _ -> error "Users cannot provide FreshVars"
  