{-# LANGUAGE BangPatterns, NamedFieldPuns, RecordWildCards #-}

module Data.SearchEngine.Update (

    -- * Managing documents to be searched
    insertDoc,
    insertDocs,
    deleteDoc,

  ) where

import Data.SearchEngine.Types
import qualified Data.SearchEngine.SearchIndex as SI
import qualified Data.SearchEngine.DocTermIds as DocTermIds

import qualified Data.List as List
import Data.Ix
import Data.Array.Unboxed


insertDocs :: (Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>
              [doc] ->
              SearchEngine doc key field feature ->
              SearchEngine doc key field feature
insertDocs :: forall key field feature doc.
(Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>
[doc]
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
insertDocs [doc]
docs SearchEngine doc key field feature
se = (SearchEngine doc key field feature
 -> doc -> SearchEngine doc key field feature)
-> SearchEngine doc key field feature
-> [doc]
-> SearchEngine doc key field feature
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (\SearchEngine doc key field feature
se' doc
doc -> doc
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall key field feature doc.
(Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>
doc
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
insertDoc doc
doc SearchEngine doc key field feature
se') SearchEngine doc key field feature
se [doc]
docs


insertDoc :: (Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>
             doc ->
             SearchEngine doc key field feature ->
             SearchEngine doc key field feature
insertDoc :: forall key field feature doc.
(Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>
doc
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
insertDoc doc
doc se :: SearchEngine doc key field feature
se@SearchEngine{ searchConfig :: forall doc key field feature.
SearchEngine doc key field feature
-> SearchConfig doc key field feature
searchConfig = SearchConfig {
                                 doc -> key
documentKey :: doc -> key
documentKey :: forall doc key field feature.
SearchConfig doc key field feature -> doc -> key
documentKey,
                                 doc -> field -> [Term]
extractDocumentTerms :: doc -> field -> [Term]
extractDocumentTerms :: forall doc key field feature.
SearchConfig doc key field feature -> doc -> field -> [Term]
extractDocumentTerms,
                                 doc -> feature -> Float
documentFeatureValue :: doc -> feature -> Float
documentFeatureValue :: forall doc key field feature.
SearchConfig doc key field feature -> doc -> feature -> Float
documentFeatureValue
                               }
                             , SearchIndex key field feature
searchIndex :: SearchIndex key field feature
searchIndex :: forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex } =
    let key :: key
key = doc -> key
documentKey doc
doc
        searchIndex' :: SearchIndex key field feature
searchIndex' = key
-> (field -> [Term])
-> (feature -> Float)
-> SearchIndex key field feature
-> SearchIndex key field feature
forall key field feature.
(Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>
key
-> DocTerms field
-> DocFeatureValues feature
-> SearchIndex key field feature
-> SearchIndex key field feature
SI.insertDoc key
key (doc -> field -> [Term]
extractDocumentTerms doc
doc)
                                        (doc -> feature -> Float
documentFeatureValue doc
doc)
                                        SearchIndex key field feature
searchIndex
        oldDoc :: Maybe (DocTermIds field)
oldDoc       = SearchIndex key field feature -> key -> Maybe (DocTermIds field)
forall key field feature.
Ord key =>
SearchIndex key field feature -> key -> Maybe (DocTermIds field)
SI.lookupDocKey SearchIndex key field feature
searchIndex  key
key
        newDoc :: Maybe (DocTermIds field)
newDoc       = SearchIndex key field feature -> key -> Maybe (DocTermIds field)
forall key field feature.
Ord key =>
SearchIndex key field feature -> key -> Maybe (DocTermIds field)
SI.lookupDocKey SearchIndex key field feature
searchIndex' key
key

     in SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall field doc key feature.
Ix field =>
SearchEngine doc key field feature
-> SearchEngine doc key field feature
cacheBM25Context (SearchEngine doc key field feature
 -> SearchEngine doc key field feature)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall a b. (a -> b) -> a -> b
$
        Maybe (DocTermIds field)
-> Maybe (DocTermIds field)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall field doc key feature.
(Ix field, Bounded field) =>
Maybe (DocTermIds field)
-> Maybe (DocTermIds field)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
updateCachedFieldLengths Maybe (DocTermIds field)
oldDoc Maybe (DocTermIds field)
newDoc (SearchEngine doc key field feature
 -> SearchEngine doc key field feature)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall a b. (a -> b) -> a -> b
$
          SearchEngine doc key field feature
se { searchIndex = searchIndex' }


deleteDoc :: (Ord key, Ix field, Bounded field) =>
             key ->
             SearchEngine doc key field feature ->
             SearchEngine doc key field feature
deleteDoc :: forall key field doc feature.
(Ord key, Ix field, Bounded field) =>
key
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
deleteDoc key
key se :: SearchEngine doc key field feature
se@SearchEngine{SearchIndex key field feature
searchIndex :: forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex :: SearchIndex key field feature
searchIndex} =
    let searchIndex' :: SearchIndex key field feature
searchIndex' = key
-> SearchIndex key field feature -> SearchIndex key field feature
forall key field feature.
(Ord key, Ix field, Bounded field) =>
key
-> SearchIndex key field feature -> SearchIndex key field feature
SI.deleteDoc key
key SearchIndex key field feature
searchIndex
        oldDoc :: Maybe (DocTermIds field)
oldDoc       = SearchIndex key field feature -> key -> Maybe (DocTermIds field)
forall key field feature.
Ord key =>
SearchIndex key field feature -> key -> Maybe (DocTermIds field)
SI.lookupDocKey SearchIndex key field feature
searchIndex key
key

     in SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall field doc key feature.
Ix field =>
SearchEngine doc key field feature
-> SearchEngine doc key field feature
cacheBM25Context (SearchEngine doc key field feature
 -> SearchEngine doc key field feature)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall a b. (a -> b) -> a -> b
$
        Maybe (DocTermIds field)
-> Maybe (DocTermIds field)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall field doc key feature.
(Ix field, Bounded field) =>
Maybe (DocTermIds field)
-> Maybe (DocTermIds field)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
updateCachedFieldLengths Maybe (DocTermIds field)
oldDoc Maybe (DocTermIds field)
forall a. Maybe a
Nothing (SearchEngine doc key field feature
 -> SearchEngine doc key field feature)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
forall a b. (a -> b) -> a -> b
$
          SearchEngine doc key field feature
se { searchIndex = searchIndex' }


updateCachedFieldLengths :: (Ix field, Bounded field) =>
                            Maybe (DocTermIds field) -> Maybe (DocTermIds field) ->
                            SearchEngine doc key field feature ->
                            SearchEngine doc key field feature
updateCachedFieldLengths :: forall field doc key feature.
(Ix field, Bounded field) =>
Maybe (DocTermIds field)
-> Maybe (DocTermIds field)
-> SearchEngine doc key field feature
-> SearchEngine doc key field feature
updateCachedFieldLengths Maybe (DocTermIds field)
Nothing (Just DocTermIds field
newDoc) se :: SearchEngine doc key field feature
se@SearchEngine{UArray field Int
sumFieldLengths :: UArray field Int
sumFieldLengths :: forall doc key field feature.
SearchEngine doc key field feature -> UArray field Int
sumFieldLengths} =
    SearchEngine doc key field feature
se {
      sumFieldLengths =
        array (bounds sumFieldLengths)
              [ (i, n + DocTermIds.fieldLength newDoc i)
              | (i, n) <- assocs sumFieldLengths ]
    }
updateCachedFieldLengths (Just DocTermIds field
oldDoc) (Just DocTermIds field
newDoc) se :: SearchEngine doc key field feature
se@SearchEngine{UArray field Int
sumFieldLengths :: forall doc key field feature.
SearchEngine doc key field feature -> UArray field Int
sumFieldLengths :: UArray field Int
sumFieldLengths} =
    SearchEngine doc key field feature
se {
      sumFieldLengths =
        array (bounds sumFieldLengths)
              [ (i, n - DocTermIds.fieldLength oldDoc i
                      + DocTermIds.fieldLength newDoc i)
              | (i, n) <- assocs sumFieldLengths ]
    }
updateCachedFieldLengths (Just DocTermIds field
oldDoc) Maybe (DocTermIds field)
Nothing se :: SearchEngine doc key field feature
se@SearchEngine{UArray field Int
sumFieldLengths :: forall doc key field feature.
SearchEngine doc key field feature -> UArray field Int
sumFieldLengths :: UArray field Int
sumFieldLengths} =
    SearchEngine doc key field feature
se {
      sumFieldLengths =
        array (bounds sumFieldLengths)
              [ (i, n - DocTermIds.fieldLength oldDoc i)
              | (i, n) <- assocs sumFieldLengths ]
    }
updateCachedFieldLengths Maybe (DocTermIds field)
Nothing Maybe (DocTermIds field)
Nothing SearchEngine doc key field feature
se = SearchEngine doc key field feature
se