content: {-# LANGUAGE CPP                        #-}
  {-# LANGUAGE DeriveDataTypeable         #-}
  {-# LANGUAGE FlexibleInstances          #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE MultiParamTypeClasses      #-}
  
  module Holumbus.Index.Inverted.CompressedPrefixMem
      ( Inverted(..)
      , Parts
      , Part
      , Inverted0
      , InvertedCompressed
      , InvertedSerialized
      , InvertedCSerialized
      , InvertedOSerialized
  
      , ComprOccurrences(..)
  
      , emptyInverted
      , emptyInverted0
      , emptyInvertedCompressed
      , emptyInvertedSerialized
      , emptyInvertedCSerialized
      , emptyInvertedOSerialized
  
      , mapOcc
      , zipOcc
      , emptyOcc
      , theOcc
      , nullOcc
      , unionOcc
      , diffOcc
      , insertPosOcc
      , deletePosOcc
      , updateDocIdOcc
      , deleteDocIds
  
      , removeDocIdsInverted
      )
  where
  
  import           Control.Arrow                  (second)
  import           Control.DeepSeq
  
  import qualified Data.Binary                    as B
  import qualified Data.ByteString                as BS
  import qualified Data.ByteString.Lazy           as BL
  import qualified Data.ByteString.Short          as SS
  
  import           Data.Function                  (on)
  
  import           Data.List                      (foldl', sortBy)
  import qualified Data.Map.Strict                as M
  import           Data.Maybe
  import qualified Data.StringMap.Strict          as PT
  import           Data.Typeable
  
  import qualified Holumbus.ByteStringCompression as BZ
  import           Holumbus.Index.Common
  import           Holumbus.Index.Compression
  
  import           Text.XML.HXT.Core              (PU, XmlPickler, xpAttr, xpElem,
                                                   xpList, xpPair, xpText, xpWrap,
                                                   xpickle)
  
  #if sizeable == 1
  import           Data.Size
  #endif
  -- ----------------------------------------------------------------------------
  
  compress   :: BL.ByteString -> BL.ByteString
  compress   = BZ.compressBZipSmart
  
  decompress :: BL.ByteString -> BL.ByteString
  decompress = BZ.decompressBZipSmart
  
  traceFrom :: Occurrences -> Occurrences
  traceFrom = id
  -- traceFrom = traceOccPos
  
  -- ----------------------------------------------------------------------------
  
  class ComprOccurrences s where
    toOccurrences         :: s -> Occurrences
    fromOccurrences       :: Occurrences -> s
  
  mapOcc                  :: (ComprOccurrences s) => (Occurrences -> Occurrences) -> s -> s
  mapOcc f                = fromOccurrences . f . toOccurrences
  
  zipOcc                  :: (ComprOccurrences s) => (Occurrences -> Occurrences -> Occurrences) -> s -> s -> s
  zipOcc op x y           = fromOccurrences $ op (toOccurrences x)(toOccurrences y)
  
  emptyOcc                :: (ComprOccurrences s) => s
  emptyOcc                = fromOccurrences  $ emptyOccurrences
  
  theOcc                  :: (ComprOccurrences s) => s -> Occurrences
  theOcc                  = toOccurrences
  
  nullOcc                 :: (ComprOccurrences s) => s -> Bool
  nullOcc                 = (== emptyOccurrences) . toOccurrences         -- better nullOccurrences, but not yet there
  
  unionOcc                :: (ComprOccurrences s) => Occurrences -> s -> s
  unionOcc os             = mapOcc $ mergeOccurrences os
  
  diffOcc                 :: (ComprOccurrences s) => Occurrences -> s -> s
  diffOcc os              = mapOcc $ substractOccurrences os
  
  insertPosOcc            :: (ComprOccurrences s) => DocId -> Position -> s -> s
  insertPosOcc d p        = mapOcc $ insertOccurrence d p
  
  deletePosOcc            :: (ComprOccurrences s) => DocId -> Position -> s -> s
  deletePosOcc d p        = mapOcc $ deleteOccurrence d p
  
  updateDocIdOcc          :: (ComprOccurrences s) => (DocId -> DocId) -> s -> s
  updateDocIdOcc f        = mapOcc $ updateOccurrences f
  
  deleteDocIds            :: (ComprOccurrences s) => Occurrences -> s -> s
  deleteDocIds ids        = mapOcc $ flip diffOccurrences ids
  
  -- ----------------------------------------------------------------------------
  --
  -- auxiliary type for strict Bytestrings
  -- the use of the smart constructor assures that the Bytestring is
  -- fully evaluated before constructing the value
  --
  -- This type is used in the variants working with compressed and/or serialised data
  -- to represent occurences 'OccCompressed', 'OccSerialized',  'OccCserialized', OccOSerialized'
  --
  -- The constructor must not be called directly, only via mkBs
  --
  -- The SByteString is a candidate for a BS.ShortByteString available with bytestring 0.10.4,
  -- then 5 machine words can be saved per value
  
  newtype SByteString     = Bs { unSs :: SS.ShortByteString }
                            deriving (Eq, Show, Typeable)
  
  mkBs                    :: BS.ByteString -> SByteString
  mkBs s                  = Bs $!! SS.toShort s
  
  unBs                    :: SByteString -> BS.ByteString
  unBs                    = SS.fromShort . unSs
  
  -- in mkBs the ByteString is physically copied into a ShortByteString
  -- to avoid sharing data with other possibly very large strings
  -- this can occur e.g in the Binary instance with get
  -- Working with ShortByteString prevents fragmentation of the heap store,
  -- ByteStrings can't be move in the heap (are PINNED), ShortByteString can.
  
  instance NFData SByteString where
  -- use default implementation: eval to WHNF, and that's sufficient
  
  instance B.Binary SByteString where
    put                   = B.put . unBs
    get                   = B.get >>= return . mkBs
  
  -- to avoid sharing the data with with the input the ByteString is physically copied
  -- before return. This should be the single place where sharing is introduced,
  -- else the copy must be moved to mSBs
  
  -- ----------------------------------------------------------------------------
  --
  -- the pure occurrence type, just wrapped in a newtype for instance declarations
  -- and for forcing evaluation
  
  newtype Occ0            = Occ0 { unOcc0 :: Occurrences }
                            deriving (Show, Typeable)
  
  mkOcc0                  :: Occurrences -> Occ0
  mkOcc0 os               = Occ0 $!! os
  
  instance ComprOccurrences Occ0 where
    fromOccurrences       = mkOcc0 . traceFrom
    toOccurrences         = unOcc0
  
  instance NFData Occ0 where
  -- use default implementation: eval to WHNF, and that's sufficient
  
  instance B.Binary Occ0 where
    put                   = B.put . unOcc0
    get                   = B.get >>= return . mkOcc0
  
  -- ----------------------------------------------------------------------------
  --
  -- the simple-9 compressed occurrence type, just wrapped in a newtype for instance declarations
  
  newtype OccCompressed   = OccCp { unOccCp :: CompressedOccurrences }
                            deriving (Eq, Show, Typeable)
  
  mkOccCp                  :: Occurrences -> OccCompressed
  mkOccCp os               = OccCp $!! deflateOcc os
  
  instance ComprOccurrences OccCompressed where
    fromOccurrences       = mkOccCp . traceFrom
    toOccurrences         = inflateOcc . unOccCp
  
  instance NFData OccCompressed where
  -- use default implementation: eval to WHNF, and that's sufficient
  
  instance B.Binary OccCompressed where
    put                   = B.put . unOccCp
    get                   = B.get >>= ((return . OccCp) $!!)
  
  -- ----------------------------------------------------------------------------
  --
  -- the simpe-9 compresses occurrences serialized into a byte strings
  
  newtype OccSerialized   = OccBs { unOccBs :: SByteString }
                            deriving (Eq, Show, NFData, Typeable)
  
  instance ComprOccurrences OccSerialized where
    fromOccurrences       = OccBs . mkBs . BL.toStrict . B.encode . deflateOcc . traceFrom
    toOccurrences         = inflateOcc . B.decode . BL.fromStrict . unBs . unOccBs
  
  instance B.Binary OccSerialized where
    put                   = B.put . unOccBs
    get                   = B.get >>= return . OccBs
  
  -- ----------------------------------------------------------------------------
  --
  -- the simple-9 compressed occurrences serialized and bzipped into a byte string
  
  newtype OccCSerialized  = OccCBs { unOccCBs :: SByteString }
                            deriving (Eq, Show, NFData, Typeable)
  
  instance ComprOccurrences OccCSerialized where
    fromOccurrences       = OccCBs . mkBs . BL.toStrict . compress . B.encode . deflateOcc . traceFrom
    toOccurrences         = inflateOcc . B.decode  . decompress . BL.fromStrict . unBs . unOccCBs
  
  instance B.Binary OccCSerialized where
    put                   = B.put . unOccCBs
    get                   = B.get >>= return . OccCBs
  
  -- ----------------------------------------------------------------------------
  --
  -- the pure occurrences serialized and bzipped into a byte string
  -- this seems to be the best choice: compression is about 4% better then simple-9 & bzip
  -- and lookup is about 8% better
  
  newtype OccOSerialized  = OccOBs { unOccOBs :: SByteString }
                            deriving (Eq, Show, NFData, Typeable)
  
  instance ComprOccurrences OccOSerialized where
    fromOccurrences       = OccOBs . mkBs . BL.toStrict . compress . B.encode .traceFrom
    toOccurrences         = B.decode . decompress . BL.fromStrict . unBs . unOccOBs
  
  instance B.Binary OccOSerialized where
    put                   = B.put . unOccOBs
    get                   = B.get >>= return . OccOBs
  
  -- ----------------------------------------------------------------------------
  
  -- | The index consists of a table which maps documents to ids and a number of index parts.
  
  newtype Inverted occ    = Inverted
                            { unInverted :: Parts  occ    -- ^ The parts of the index, each representing one context.
                            }
                            deriving (Show, Eq, NFData, Typeable)
  
  -- | The index parts are identified by a name, which should denote the context of the words.
  
  type Parts occ          = M.Map Context (Part occ)
  
  -- | The index part is the real inverted index. Words are mapped to their occurrences.
  -- The part is implemented as a prefix tree
  
  type Part occ           = PT.StringMap occ
  
  -- ----------------------------------------------------------------------------
  
  instance (ComprOccurrences occ) => XmlPickler (Inverted occ) where
    xpickle               =  xpElem "indexes" $
                             xpWrap (Inverted, unInverted) xpParts
  
  xpParts                 :: (ComprOccurrences occ) => PU (Parts occ)
  xpParts                 = xpWrap (M.fromList, M.toList) (xpList xpContext)
    where
    xpContext             = xpElem "part" (xpPair (xpAttr "id" xpText) xpPart)
  
  xpPart                  :: (ComprOccurrences occ) => PU (Part occ)
  xpPart                  = xpElem "index" (xpWrap (PT.fromList, PT.toList) (xpList xpWord))
    where
    xpWord                = xpElem "word" $
                            xpPair (xpAttr "w" xpText)
                                   (xpWrap (fromOccurrences, toOccurrences) xpOccurrences)
  
  -- ----------------------------------------------------------------------------
  
  instance (B.Binary occ) => B.Binary (Inverted occ) where
    put                   = B.put . unInverted
    get                   = B.get >>= return . Inverted
  
  -- ----------------------------------------------------------------------------
  
  instance (B.Binary occ, ComprOccurrences occ) => HolIndex (Inverted occ) where
    sizeWords                     = M.foldl' (\ x y -> x + PT.size y) 0 . unInverted
    contexts                      = fmap fst . M.toList . unInverted
  
    allWords     i c              = fmap (second  toOccurrences) .
                                    PT.toList .
                                    getPart c $ i
  
    prefixCase   i c q            = fmap (second  toOccurrences) .
                                    PT.toListShortestFirst .
                                    PT.prefixFilter q .
                                    getPart c $ i
  
    prefixNoCase i c q            = fmap (second  toOccurrences) .
                                    PT.toListShortestFirst .
                                    PT.prefixFilterNoCase q .
                                    getPart c $ i
  
    lookupCase   i c q            = fmap ((,) q . toOccurrences) .
                                    maybeToList .
                                    PT.lookup q .
                                    getPart c $ i
  
    lookupNoCase i c q            = fmap (second  toOccurrences) .
                                    PT.toList .
                                    PT.lookupNoCase q .
                                    getPart c $ i
  
    mergeIndexes                  = zipInverted $ M.unionWith      $ PT.unionWith (zipOcc mergeOccurrences)
    substractIndexes              = zipInverted $ M.differenceWith $ substractPart
  
    insertOccurrences c w o i     = mergeIndexes     i (singletonInverted c w o)          -- see "http://holumbus.fh-wedel.de/hayoo/hayoo.html#0:unionWith%20module%3AData.Map"
    deleteOccurrences c w o i     = substractIndexes i (singletonInverted c w o)
  
    toList                        = concatMap (uncurry convertPart) . toListInverted
                                    where
                                    convertPart c  = map (\(w, o) -> (c, w, toOccurrences o)) . PT.toList
  
    splitByContexts               = splitInverted
                                    . map ( (\ i -> (sizeWords i, i))
                                            . Inverted
                                            . uncurry M.singleton
                                          )
                                    . toListInverted
  
    splitByDocuments i            = splitInverted ( map (uncurry convert) $
                                                    toListDocIdMap $
                                                    unionsWithDocIdMap (M.unionWith (M.unionWith unionPos)) docResults
                                                  )
      where
      docResults                  = map (\c -> resultByDocument c (allWords i c)) (contexts i)
      convert d cs                = foldl' makeIndex (0, emptyInverted) (M.toList cs)
        where
        makeIndex r (c, ws)       = foldl' makeOcc r (M.toList ws)
          where
          makeOcc (rs, ri) (w, p) = (sizePos p + rs, insertOccurrences c w (singletonDocIdMap d p) ri)
  
    splitByWords i                = splitInverted ( map (uncurry convert) .
                                                    M.toList .
                                                    M.unionsWith (M.unionWith mergeOccurrences) $ wordResults
                                                  )
      where
      wordResults                 = map (\c -> resultByWord c (allWords i c)) (contexts i)
      convert w cs                = foldl' makeIndex (0, emptyInverted) (M.toList cs)
        where
        makeIndex (rs, ri) (c, o) = (rs + sizeOccurrences o, insertOccurrences c w o ri)
  
    updateDocIds f                = mapInverted (M.mapWithKey updatePart)
      where
      updatePart c                = PT.mapWithKey (updateOcc (f c))
      updateOcc f' w              = mapOcc $ updateOccurrences (f' w)
  
  -- ----------------------------------------------------------------------------
  
  -- | Return a part of the index for a given context.
  
  getPart                         :: Context -> Inverted i -> Part i
  getPart c                       = fromMaybe PT.empty . M.lookup c . unInverted
  
  -- | Substract one index part from another.
  substractPart                   :: (ComprOccurrences i) => Part i -> Part i -> Maybe (Part i)
  substractPart p1 p2
      | PT.null res               = Nothing
      | otherwise                 = Just res
      where
      res                         = diffPart p1 p2
  
  diffPart                        :: (ComprOccurrences i) => Part i -> Part i -> Part i
  diffPart                        = PT.differenceWith subtractDiffLists
      where
      subtractDiffLists o1 o2
          | nullOcc res           = Nothing
          | otherwise             = Just res
          where
          res                     = zipOcc substractOccurrences o1 o2
  
  removeDocIdsPart                :: (ComprOccurrences i) => Occurrences -> Part i -> Part i
  removeDocIdsPart ids            = PT.foldrWithKey removeDocIds PT.empty
      where
      removeDocIds k occ acc
          | nullOcc occ'          = acc
          | otherwise             = PT.insert k occ' acc
          where
          occ'                    = deleteDocIds ids occ
  
  -- ----------------------------------------------------------------------------
  
  mapInverted                     :: (Parts i -> Parts i) -> Inverted i -> Inverted i
  mapInverted f                   = Inverted . f . unInverted
  
  zipInverted                     :: (Parts i -> Parts i -> Parts i) -> Inverted i -> Inverted i -> Inverted i
  zipInverted op i1 i2            = Inverted $ op (unInverted i1) (unInverted i2)
  
  emptyInverted                   :: Inverted i
  emptyInverted                   = Inverted M.empty
  
  toListInverted                  :: Inverted i -> [(Context, Part i)]
  toListInverted                  = M.toList . unInverted
  
  -- | Create an index with just one word in one context.
  
  singletonInverted               :: (ComprOccurrences i) => Context -> String -> Occurrences -> Inverted i
  singletonInverted c w o         = Inverted . M.singleton c . PT.singleton w . fromOccurrences $ o
  
  -- | Remove DocIds from index
  
  removeDocIdsInverted            :: (ComprOccurrences i) => Occurrences -> Inverted i -> Inverted i
  removeDocIdsInverted ids        = mapInverted $ M.map (removeDocIdsPart ids)
  
  -- ----------------------------------------------------------------------------
  --
  -- copied from Holumbus.Index.Inverted.Memory
  
  splitInverted                   :: (B.Binary i, ComprOccurrences i) => [(Int, Inverted i)] -> Int -> [Inverted i]
  splitInverted inp n             = allocate mergeIndexes stack buckets
    where
    buckets                       = zipWith const (createBuckets n) stack
    stack                         = reverse (sortBy (compare `on` fst) inp)
  
  -- | Allocates values from the first list to the buckets in the second list.
  allocate                        :: (a -> a -> a) -> [(Int, a)] -> [(Int, a)] -> [a]
  allocate _ _ []                 = []
  allocate _ [] ys                = map snd ys
  allocate f (x:xs) (y:ys)        = allocate f xs (sortBy (compare `on` fst) ((combine x y):ys))
    where
    combine (s1, v1) (s2, v2)     = (s1 + s2, f v1 v2)
  
  -- | Create empty buckets for allocating indexes.
  createBuckets                   :: Int -> [(Int, Inverted i)]
  createBuckets n                 = (replicate n (0, emptyInverted))
  
  -- ----------------------------------------------------------------------------
  --
  -- the 5 variants for the inverted index as prefix tree,
  
  -- | The pure inverted index implemented as a prefix tree without any space optimizations.
  -- This may be taken as a reference for space and time measurements for the other index structures
  
  type Inverted0                  = Inverted Occ0
  
  -- | The inverted index with simple-9 encoding of the occurence sets
  
  type InvertedCompressed         = Inverted OccCompressed
  
  -- | The inverted index with serialized occurence maps with simple-9 encoded sets
  
  type InvertedSerialized         = Inverted OccSerialized
  
  -- | The inverted index with serialized occurence maps with simple-9 encoded sets
  -- and with the serialized bytestrings compressed with bzip2
  
  type InvertedCSerialized        = Inverted OccCSerialized
  
  -- | The pure inverted index with serialized occurence maps
  -- and with the serialized bytestrings compressed with bzip2, no simple-9 encoding.
  -- This is the most space efficient index of the 5 variants, even a few percent smaller
  -- then InvertedCSerialized, and a few percent faster in lookup
  
  type InvertedOSerialized        = Inverted OccOSerialized
  
  -- ----------------------------------------------------------------------------
  
  emptyInverted0                  :: Inverted0
  emptyInverted0                  = emptyInverted
  
  emptyInvertedCompressed         :: InvertedCompressed
  emptyInvertedCompressed         = emptyInverted
  
  emptyInvertedSerialized         :: InvertedSerialized
  emptyInvertedSerialized         = emptyInverted
  
  emptyInvertedCSerialized        :: InvertedCSerialized
  emptyInvertedCSerialized        = emptyInverted
  
  emptyInvertedOSerialized        :: InvertedOSerialized
  emptyInvertedOSerialized        = emptyInverted
  
  
  -- ----------------------------------------------------------------------------
  #if sizeable == 1
  
  instance Sizeable SByteString where
      dataOf                      = dataOf  . unSs
      bytesOf                     = bytesOf . unSs
      statsOf                     = statsOf . unSs
  
  instance Sizeable OccSerialized where
      dataOf              = dataOf  . unOccBs
      bytesOf             = bytesOf . unOccBs
      statsOf             = statsOf . unOccBs
  
  instance Sizeable OccCSerialized where
      dataOf              = dataOf  . unOccCBs
      bytesOf             = bytesOf . unOccCBs
      statsOf             = statsOf . unOccCBs
  
  instance Sizeable Occ0 where
      dataOf              = dataOf  . unOcc0
      bytesOf             = bytesOf . unOcc0
      statsOf             = statsOf . unOcc0
  
  instance Sizeable OccCompressed where
      dataOf              = dataOf  . unOccCp
      bytesOf             = bytesOf . unOccCp
      statsOf             = statsOf . unOccCp
  
  instance Sizeable OccOSerialized where
      dataOf              = dataOf  . unOccOBs
      bytesOf             = bytesOf . unOccOBs
      statsOf             = statsOf . unOccOBs
  
  instance (Typeable occ, Sizeable occ) => Sizeable (Inverted occ) where
      dataOf              = dataOf  . unInverted
      bytesOf             = bytesOf . unInverted
      statsOf             = statsOf . unInverted
  
  #endif
  -- ----------------------------------------------------------------------------
  
  
  