module FTrie where

import Types

import qualified Data.IntSet as IS
import qualified Data.IntMap.Lazy as IM
import Data.List (groupBy, sortOn)
import Control.Arrow (second)
import Data.Function (on)

-- this is a path-compressed version of feature vector indexing as described in https://link.springer.com/chapter/10.1007/978-3-642-36675-8_3
-- we implement perfect subsumption checking by using all literals as features;
-- a Node stores whether a literal does not occur (left) or does occur (right) in a set of clause keys
data FTrie = Leaf IS.IntSet | Node FTrie Name IS.IntSet FTrie
  deriving (Eq, Show)

emptyFTrie :: FTrie
emptyFTrie = Leaf IS.empty

buildFTrie :: KeyConjForm -> FTrie
buildFTrie = run [] . group . sortOn snd . map (second IS.toAscList) . IM.toList
  where
    group = groupBy (groupHead `on` snd)
    groupHead _ [] = True
    groupHead [] _ = False
    groupHead (k1:_) (k2:_) = k1 == k2
    run :: [Key] -> [[(Key, [Literal])]] -> FTrie
    run lks [] = Leaf $ IS.fromList lks
    -- note: as we call run with groupBy, ther will always be at least one element in the list
    run lks gkcs@(gkc@((_, []):_):gkct) = run nLks gkct
      where nLks = lks ++ map fst gkc
    run lks gkcs@(gkc@((_, l:_):_):gkct) = Node left l (IS.fromList ks) right
      where
        ks = lks ++ concatMap (map fst) gkcs 
        left = run lks gkct
        right = run [] $ group $ map (second tail) gkc
    run _ _ = error "unreachable code"

insertFTrie :: Key -> Clause -> FTrie -> FTrie
insertFTrie k c fTrie = run fTrie $ IS.toAscList c
  where
    run (Leaf ks) [] = Leaf $ IS.insert k ks
    run (Node left l ks right) [] = Node (run left []) l (IS.insert k ks) right
    run t@(Leaf ks) (l:ls) = Node t l (IS.insert k ks) (run emptyFTrie ls)
    run t@(Node left l ks right) ls@(l':lt)
      | l < l' = Node (run left ls) l nks right
      | l == l' = Node left l nks (run right lt)
      | otherwise = Node t l' nks (run emptyFTrie lt)
      where nks = IS.insert k ks

fTrieSubsumes :: FTrie -> Clause -> Bool
fTrieSubsumes fTrie = find fTrie . IS.toAscList
  where
    find :: FTrie -> [Literal] -> Bool
    find (Leaf ks) _ = not $ IS.null ks
    find (Node left _ _ _) [] = find left []
    find t@(Node left l _ right) ls@(l':lt)
      | l < l' = find left ls
      | l == l' = find left lt || find right lt
      | otherwise = find t lt

filterSubsumedFTrie :: Clause -> FTrie -> (IS.IntSet, FTrie)
filterSubsumedFTrie c fTrie = filterAux fTrie $ IS.toAscList c
  where
    filterAux (Leaf ks) [] = (ks, emptyFTrie)
    filterAux (Leaf ks) _ = (IS.empty, Leaf ks)
    filterAux (Node _ _ ks _) [] = (ks, emptyFTrie)
    filterAux ftrie@(Node left l ks right) ls@(l':lt)
      | l < l' = let (ksl, fTrieL) = filterAux left ls
                     (ksr, fTrieR) = filterAux right ls in
        (ksl `IS.union` ksr, compress fTrieL fTrieR l $ ks IS.\\ (ksl `IS.union` ksr))
      | l == l' = let (ksr, fTrieR) = filterAux right lt in
        (ksr, compress left fTrieR l $ ks IS.\\ ksr)
      | otherwise = (IS.empty, ftrie)
    compress left@(Leaf ksl) right@(Leaf ksr) l ks
      | IS.null ksl && IS.null ksr = Leaf ks
      | otherwise = Node left l ks right
    compress left right l ks = Node left l ks right

filterSubsumedDataKeyConjFormFTrie :: Clause -> KeyDataConjForm -> FTrie -> (KeyDataConjForm, FTrie)
filterSubsumedDataKeyConjFormFTrie c kcs fTrie =
  let (ks, fTrie') = c `filterSubsumedFTrie` fTrie
      kcs' = kcs `IM.withoutKeys` ks in
  (kcs', fTrie')