{-# 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 Data.Ix
import Data.Array.Unboxed
import Data.List


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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\SearchEngine doc key field feature
se' doc
doc -> 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 :: forall doc key field feature.
SearchConfig doc key field feature -> doc -> key
documentKey :: doc -> key
documentKey,
                                 doc -> field -> [Term]
extractDocumentTerms :: forall doc key field feature.
SearchConfig doc key field feature -> doc -> field -> [Term]
extractDocumentTerms :: doc -> field -> [Term]
extractDocumentTerms,
                                 doc -> feature -> Float
documentFeatureValue :: forall doc key field feature.
SearchConfig doc key field feature -> doc -> feature -> Float
documentFeatureValue :: doc -> feature -> Float
documentFeatureValue
                               }
                             , 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 key :: key
key = doc -> key
documentKey doc
doc
        searchIndex' :: SearchIndex key field feature
searchIndex' = 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       = 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       = forall key field feature.
Ord key =>
SearchIndex key field feature -> key -> Maybe (DocTermIds field)
SI.lookupDocKey SearchIndex key field feature
searchIndex' key
key

     in forall field doc key feature.
Ix field =>
SearchEngine doc key field feature
-> SearchEngine doc key field feature
cacheBM25Context forall a b. (a -> b) -> a -> b
$
        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 forall a b. (a -> b) -> a -> b
$
          SearchEngine doc key field feature
se { searchIndex :: SearchIndex key field feature
searchIndex = SearchIndex key field feature
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 :: SearchIndex key field feature
searchIndex :: forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex} =
    let searchIndex' :: SearchIndex key field feature
searchIndex' = 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       = forall key field feature.
Ord key =>
SearchIndex key field feature -> key -> Maybe (DocTermIds field)
SI.lookupDocKey SearchIndex key field feature
searchIndex key
key

     in forall field doc key feature.
Ix field =>
SearchEngine doc key field feature
-> SearchEngine doc key field feature
cacheBM25Context forall a b. (a -> b) -> a -> b
$
        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 forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$
          SearchEngine doc key field feature
se { searchIndex :: SearchIndex key field feature
searchIndex = SearchIndex key field feature
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 :: 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 :: UArray field Int
sumFieldLengths =
        forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray field Int
sumFieldLengths)
              [ (field
i, Int
n forall a. Num a => a -> a -> a
+ forall field.
(Ix field, Bounded field) =>
DocTermIds field -> field -> Int
DocTermIds.fieldLength DocTermIds field
newDoc field
i)
              | (field
i, Int
n) <- forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray field Int
sumFieldLengths ]
    }
updateCachedFieldLengths (Just DocTermIds field
oldDoc) (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 :: UArray field Int
sumFieldLengths =
        forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray field Int
sumFieldLengths)
              [ (field
i, Int
n forall a. Num a => a -> a -> a
- forall field.
(Ix field, Bounded field) =>
DocTermIds field -> field -> Int
DocTermIds.fieldLength DocTermIds field
oldDoc field
i
                      forall a. Num a => a -> a -> a
+ forall field.
(Ix field, Bounded field) =>
DocTermIds field -> field -> Int
DocTermIds.fieldLength DocTermIds field
newDoc field
i)
              | (field
i, Int
n) <- forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray field Int
sumFieldLengths ]
    }
updateCachedFieldLengths (Just DocTermIds field
oldDoc) Maybe (DocTermIds field)
Nothing 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 :: UArray field Int
sumFieldLengths =
        forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray field Int
sumFieldLengths)
              [ (field
i, Int
n forall a. Num a => a -> a -> a
- forall field.
(Ix field, Bounded field) =>
DocTermIds field -> field -> Int
DocTermIds.fieldLength DocTermIds field
oldDoc field
i)
              | (field
i, Int
n) <- forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray field Int
sumFieldLengths ]
    }
updateCachedFieldLengths Maybe (DocTermIds field)
Nothing Maybe (DocTermIds field)
Nothing SearchEngine doc key field feature
se = SearchEngine doc key field feature
se