content: -- ----------------------------------------------------------------------------
  
  {- |
    Module     : Holumbus.Index.Inverted.PrefixMem
    Copyright  : Copyright (C) 2007 - 2009 Sebastian M. Schlatt, Timo B. Huebel, Uwe Schmidt
    License    : MIT
    
    Maintainer : Uwe Schmidt (uwe@fh-wedel.de)
    Stability  : experimental
    Portability: portable
    Version    : 0.3
    
    A variant of the Inverted.Memory index with an optimized prefix tree
    instead of a trie as central data structure. This version should be
    more space efficient as the trie and more runtime efficient when combining
    whole indexes.
  
    For switching from Memory to this module, only the import has to be modified
  
  -}
  
  -- ----------------------------------------------------------------------------
  
  {-# OPTIONS -XTypeSynonymInstances -XFlexibleInstances -XMultiParamTypeClasses #-}
  
  module Holumbus.Index.Inverted.PrefixMem 
      (
       -- * Inverted index types
       Inverted (..)
      , Parts
      , Part
    
      -- * Construction
      , singleton
      , emptyInverted
      )
  where
  
  import 		 Control.DeepSeq
  
  import 		 Data.Binary 			hiding (Word)
  import 		 Data.Function
  import qualified Data.IntMap 			as IM
  import qualified Data.IntSet 			as IS
  import 		 Data.List
  import 		 Data.Map 			(Map)
  import qualified Data.Map 			as M
  import 		 Data.Maybe
  
  import 		 Holumbus.Control.MapReduce.MapReducible
  import qualified Holumbus.Data.PrefixTree	as PT
  import 		 Holumbus.Index.Common
  import 		 Holumbus.Index.Compression
  
  import 		 Text.XML.HXT.Core
  
  -- ----------------------------------------------------------------------------
  
  -- | The index consists of a table which maps documents to ids and a number of index parts.
  
  newtype Inverted 	= Inverted { indexParts :: Parts  -- ^ The parts of the index, each representing one context.
  				   } 
                            deriving (Show, Eq)
  
  -- | The index parts are identified by a name, which should denote the context of the words.
  type Parts       	= Map Context Part
  
  -- | The index part is the real inverted index. Words are mapped to their occurrences.
  type Part        	= PT.PrefixTree CompressedOccurrences
  
  -- ----------------------------------------------------------------------------
  
  instance MapReducible Inverted (Context, Word) Occurrences where
    mergeMR i1 i2       = return $ mergeIndexes i1 i2
    reduceMR _ (c,w) os = do
                          let idx = singleton c w (IM.unionsWith IS.union os)
                              _   = rnf idx
                          return $ Just $ idx 
  
  -- ----------------------------------------------------------------------------
  
  instance HolIndex Inverted where
    sizeWords 			= M.fold ((+) . PT.size) 0 . indexParts
    contexts 			= map fst . M.toList . indexParts
  
    allWords i c 			= map (\(w, o) -> (w, inflateOcc o)) $ PT.toList $ getPart c i
    prefixCase i c q 		= map (\(w, o) -> (w, inflateOcc o)) $ PT.prefixFindWithKey q $ getPart c i
    prefixNoCase i c q 		= map (\(w, o) -> (w, inflateOcc o)) $ PT.prefixFindNoCaseWithKey q $ getPart c i
    lookupCase i c q 		= map (\    o  -> (q, inflateOcc o)) $ maybeToList (PT.lookup q $ getPart c i)
    lookupNoCase i c q 		= map (\(w, o) -> (w, inflateOcc o)) $ PT.lookupNoCase q $ getPart c i
  
    mergeIndexes i1 i2 		= Inverted (mergeParts (indexParts i1) (indexParts i2))
    substractIndexes i1 i2 	= Inverted (substractParts (indexParts i1) (indexParts i2))
  
    insertOccurrences c w o i 	= mergeIndexes (singleton c w o) i
    deleteOccurrences c w o i 	= substractIndexes i (singleton c w o)
  
    splitByContexts (Inverted ps) = splitInternal (map (uncurry annotate) . M.toList $ ps)
      where
      annotate c p 		= let
   				  i = Inverted (M.singleton c p)
                                    in
                                    (sizeWords i, i)
  
    splitByDocuments i 		= splitInternal ( map convert $
  						  IM.toList $
  						  IM.unionsWith unionDocs' docResults
  						)
      where
      unionDocs' 			= M.unionWith (M.unionWith IS.union)
      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) = (IS.size p + rs , insertOccurrences c w (IM.singleton d p) ri)
  
    splitByWords i 		= splitInternal indexes
      where
      indexes 			= map 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 (Inverted parts)
  				= Inverted (M.mapWithKey updatePart parts)
      where
      updatePart c p 		= PT.mapWithKey (\w o -> IM.foldWithKey (updateDocument c w) IM.empty o) p
      updateDocument c w d p r 	= IM.insertWith mergePositions (f c w d) p r
        where
        mergePositions p1 p2 	= deflatePos $ IS.union (inflatePos p1) (inflatePos p2)
  
    updateDocIds' f
  				= Inverted . M.map updatePart . indexParts
      where
      updatePart	 		= PT.map updateOcc
      updateOcc			= IM.foldWithKey updateId IM.empty
      updateId			= IM.insert . f
  
    toList i 			= concat $ map convertPart $ M.toList (indexParts i) 
      where convertPart (c,p) = map (\(w, o) -> (c, w, inflateOcc o)) $ PT.toList $ p
  
  -- ----------------------------------------------------------------------------
  
  instance NFData Inverted where
      rnf (Inverted parts) 	= rnf parts
  
  -- ----------------------------------------------------------------------------
  
  instance XmlPickler Inverted where
    xpickle 			=  xpElem "indexes" $
  				   xpWrap (\p -> Inverted p, \(Inverted p) -> p) xpParts
  
  -- | The XML pickler for the index parts.
  xpParts 			:: PU Parts
  xpParts 			= xpWrap (M.fromList, M.toList) (xpList xpContext)
    where
    xpContext 			= xpElem "part" (xpPair (xpAttr "id" xpText) xpPart)
  
  -- | The XML pickler for a single part.
  xpPart 				:: PU Part
  xpPart 				= xpElem "index" (xpWrap (PT.fromList, PT.toList) (xpList xpWord))
    where
    xpWord 			= xpElem "word" $
  				  xpPair (xpAttr "w" xpText)
  					 (xpWrap (deflateOcc, inflateOcc) xpOccurrences)
  
  -- ----------------------------------------------------------------------------
  
  instance Binary Inverted where
      put (Inverted parts) 	= put parts
      get 			= do
  				  parts <- get
  				  return (Inverted parts)
  
  -- ----------------------------------------------------------------------------
  
  -- | Create an empty index.
  emptyInverted 			:: Inverted
  emptyInverted 			= Inverted M.empty
                    
  -- | Create an index with just one word in one context.
  singleton 			:: Context -> String -> Occurrences -> Inverted
  singleton c w o 		= Inverted (M.singleton c (PT.singleton w (deflateOcc o)))
  
  -- | Merge two sets of index parts.
  mergeParts 			:: Parts -> Parts -> Parts
  mergeParts 			= M.unionWith mergePart
  
  -- | Merge two index parts.
  mergePart 			:: Part -> Part -> Part
  mergePart 			= PT.unionWith mergeDiffLists
    where
    mergeDiffLists o1 o2 		= deflateOcc $
  				  mergeOccurrences (inflateOcc o1) (inflateOcc o2)
  
  -- | Substract a set of index parts from another.
  substractParts 			:: Parts -> Parts -> Parts
  substractParts 			= M.differenceWith substractPart
  
  -- | Substract one index part from another.
  substractPart 			:: Part -> Part -> Maybe Part
  substractPart p1 p2 		= if PT.null diffPart then Nothing else Just diffPart
    where
    diffPart 			= PT.differenceWith substractDiffLists p1 p2
      where
      substractDiffLists o1 o2 	= if diffOcc == emptyOccurrences then Nothing else Just (deflateOcc diffOcc)
        where
        diffOcc 			= substractOccurrences (inflateOcc o1) (inflateOcc o2)
  
  -- | Internal split function used by the split functions from the HolIndex interface (above).
  splitInternal 			:: [(Int, Inverted)] -> Int -> [Inverted]
  splitInternal 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)]
  createBuckets n 		= (replicate n (0, emptyInverted))
    
  -- | Return a part of the index for a given context.
  getPart 			:: Context -> Inverted -> Part
  getPart c i 			= fromMaybe PT.empty (M.lookup c $ indexParts i)
  
  -- ----------------------------------------------------------------------------
  