{-# LANGUAGE NamedFieldPuns, RecordWildCards #-}

module Data.SearchEngine.Types (
    -- * Search engine types and helper functions
    SearchEngine(..),
    SearchConfig(..),
    SearchRankParameters(..),
    BM25F.FeatureFunction(..),
    initSearchEngine,
    cacheBM25Context,

    -- ** Helper type for non-term features
    NoFeatures,
    noFeatures,

    -- * Re-export SearchIndex and other types
    SearchIndex, Term, TermId,
    DocIdSet, DocId,
    DocTermIds, DocFeatVals,

    -- * Internal sanity check
    invariant,
  ) where

import Data.SearchEngine.SearchIndex (SearchIndex, Term, TermId)
import qualified Data.SearchEngine.SearchIndex as SI
import Data.SearchEngine.DocIdSet (DocIdSet, DocId)
import qualified Data.SearchEngine.DocIdSet as DocIdSet
import Data.SearchEngine.DocFeatVals (DocFeatVals)
import Data.SearchEngine.DocTermIds  (DocTermIds)
import qualified Data.SearchEngine.BM25F as BM25F

import Data.Ix
import Data.Array.Unboxed



data SearchConfig doc key field feature = SearchConfig {
       forall doc key field feature.
SearchConfig doc key field feature -> doc -> key
documentKey          :: doc -> key,
       forall doc key field feature.
SearchConfig doc key field feature -> doc -> field -> [Term]
extractDocumentTerms :: doc -> field -> [Term],
       forall doc key field feature.
SearchConfig doc key field feature -> Term -> field -> Term
transformQueryTerm   :: Term -> field -> Term,
       forall doc key field feature.
SearchConfig doc key field feature -> doc -> feature -> Float
documentFeatureValue :: doc -> feature -> Float
     }

data SearchRankParameters field feature = SearchRankParameters {
       forall field feature. SearchRankParameters field feature -> Float
paramK1                 :: !Float,
       forall field feature.
SearchRankParameters field feature -> field -> Float
paramB                  :: field -> Float,
       forall field feature.
SearchRankParameters field feature -> field -> Float
paramFieldWeights       :: field -> Float,
       forall field feature.
SearchRankParameters field feature -> feature -> Float
paramFeatureWeights     :: feature -> Float,
       forall field feature.
SearchRankParameters field feature -> feature -> FeatureFunction
paramFeatureFunctions   :: feature -> BM25F.FeatureFunction,

       forall field feature. SearchRankParameters field feature -> Int
paramResultsetSoftLimit   :: !Int,
       forall field feature. SearchRankParameters field feature -> Int
paramResultsetHardLimit   :: !Int,
       forall field feature. SearchRankParameters field feature -> Int
paramAutosuggestPrefilterLimit  :: !Int,
       forall field feature. SearchRankParameters field feature -> Int
paramAutosuggestPostfilterLimit :: !Int
     }

data SearchEngine doc key field feature = SearchEngine {
       forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex      :: !(SearchIndex      key field feature),
       forall doc key field feature.
SearchEngine doc key field feature
-> SearchConfig doc key field feature
searchConfig     :: !(SearchConfig doc key field feature),
       forall doc key field feature.
SearchEngine doc key field feature
-> SearchRankParameters field feature
searchRankParams :: !(SearchRankParameters field feature),

       -- cached info
       forall doc key field feature.
SearchEngine doc key field feature -> UArray field Int
sumFieldLengths :: !(UArray field Int),
       forall doc key field feature.
SearchEngine doc key field feature -> Context TermId field feature
bm25Context     :: BM25F.Context TermId field feature
     }

invariant :: (Ord key, Ix field, Bounded field) =>
             SearchEngine doc key field feature -> Bool
invariant :: forall key field doc feature.
(Ord key, Ix field, Bounded field) =>
SearchEngine doc key field feature -> Bool
invariant 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} =
    forall key field feature.
(Ord key, Ix field, Bounded field) =>
SearchIndex key field feature -> Bool
SI.invariant SearchIndex key field feature
searchIndex
-- && check caches

initSearchEngine :: (Ix field, Bounded field, Ix feature, Bounded feature) =>
                    SearchConfig doc key field feature ->
                    SearchRankParameters field feature ->
                    SearchEngine doc key field feature
initSearchEngine :: forall field feature doc key.
(Ix field, Bounded field, Ix feature, Bounded feature) =>
SearchConfig doc key field feature
-> SearchRankParameters field feature
-> SearchEngine doc key field feature
initSearchEngine SearchConfig doc key field feature
config SearchRankParameters field feature
params =
    forall field doc key feature.
Ix field =>
SearchEngine doc key field feature
-> SearchEngine doc key field feature
cacheBM25Context
      SearchEngine {
        searchIndex :: SearchIndex key field feature
searchIndex      = forall key field feature. SearchIndex key field feature
SI.emptySearchIndex,
        searchConfig :: SearchConfig doc key field feature
searchConfig     = SearchConfig doc key field feature
config,
        searchRankParams :: SearchRankParameters field feature
searchRankParams = SearchRankParameters field feature
params,
        sumFieldLengths :: UArray field Int
sumFieldLengths  = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
listArray (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound) (forall a. a -> [a]
repeat Int
0),
        bm25Context :: Context TermId field feature
bm25Context      = forall a. HasCallStack => a
undefined
      }

cacheBM25Context :: Ix field =>
                    SearchEngine doc key field feature ->
                    SearchEngine doc key field feature
cacheBM25Context :: forall field doc key feature.
Ix field =>
SearchEngine doc key field feature
-> SearchEngine doc key field feature
cacheBM25Context
    se :: SearchEngine doc key field feature
se@SearchEngine {
      searchRankParams :: forall doc key field feature.
SearchEngine doc key field feature
-> SearchRankParameters field feature
searchRankParams = SearchRankParameters{Float
Int
field -> Float
feature -> Float
feature -> FeatureFunction
paramAutosuggestPostfilterLimit :: Int
paramAutosuggestPrefilterLimit :: Int
paramResultsetHardLimit :: Int
paramResultsetSoftLimit :: Int
paramFeatureFunctions :: feature -> FeatureFunction
paramFeatureWeights :: feature -> Float
paramFieldWeights :: field -> Float
paramB :: field -> Float
paramK1 :: Float
paramAutosuggestPostfilterLimit :: forall field feature. SearchRankParameters field feature -> Int
paramAutosuggestPrefilterLimit :: forall field feature. SearchRankParameters field feature -> Int
paramResultsetHardLimit :: forall field feature. SearchRankParameters field feature -> Int
paramResultsetSoftLimit :: forall field feature. SearchRankParameters field feature -> Int
paramFeatureFunctions :: forall field feature.
SearchRankParameters field feature -> feature -> FeatureFunction
paramFeatureWeights :: forall field feature.
SearchRankParameters field feature -> feature -> Float
paramFieldWeights :: forall field feature.
SearchRankParameters field feature -> field -> Float
paramB :: forall field feature.
SearchRankParameters field feature -> field -> Float
paramK1 :: forall field feature. SearchRankParameters field feature -> Float
..},
      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,
      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 { bm25Context :: Context TermId field feature
bm25Context = Context TermId field feature
bm25Context' }
  where
    bm25Context' :: Context TermId field feature
bm25Context' = BM25F.Context {
      numDocsTotal :: Int
BM25F.numDocsTotal    = forall key field feature. SearchIndex key field feature -> Int
SI.docCount SearchIndex key field feature
searchIndex,
      avgFieldLength :: field -> Float
BM25F.avgFieldLength  = \field
f -> forall a b. (Integral a, Num b) => a -> b
fromIntegral (UArray field Int
sumFieldLengths forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! field
f)
                                  forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall key field feature. SearchIndex key field feature -> Int
SI.docCount SearchIndex key field feature
searchIndex),
      numDocsWithTerm :: TermId -> Int
BM25F.numDocsWithTerm = DocIdSet -> Int
DocIdSet.size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall key field feature.
SearchIndex key field feature -> TermId -> DocIdSet
SI.lookupTermId SearchIndex key field feature
searchIndex,
      paramK1 :: Float
BM25F.paramK1         = Float
paramK1,
      paramB :: field -> Float
BM25F.paramB          = field -> Float
paramB,
      fieldWeight :: field -> Float
BM25F.fieldWeight     = field -> Float
paramFieldWeights,
      featureWeight :: feature -> Float
BM25F.featureWeight   = feature -> Float
paramFeatureWeights,
      featureFunction :: feature -> FeatureFunction
BM25F.featureFunction = feature -> FeatureFunction
paramFeatureFunctions
    }


-----------------------------

data NoFeatures = NoFeatures
  deriving (NoFeatures -> NoFeatures -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NoFeatures -> NoFeatures -> Bool
$c/= :: NoFeatures -> NoFeatures -> Bool
== :: NoFeatures -> NoFeatures -> Bool
$c== :: NoFeatures -> NoFeatures -> Bool
Eq, Eq NoFeatures
NoFeatures -> NoFeatures -> Bool
NoFeatures -> NoFeatures -> Ordering
NoFeatures -> NoFeatures -> NoFeatures
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: NoFeatures -> NoFeatures -> NoFeatures
$cmin :: NoFeatures -> NoFeatures -> NoFeatures
max :: NoFeatures -> NoFeatures -> NoFeatures
$cmax :: NoFeatures -> NoFeatures -> NoFeatures
>= :: NoFeatures -> NoFeatures -> Bool
$c>= :: NoFeatures -> NoFeatures -> Bool
> :: NoFeatures -> NoFeatures -> Bool
$c> :: NoFeatures -> NoFeatures -> Bool
<= :: NoFeatures -> NoFeatures -> Bool
$c<= :: NoFeatures -> NoFeatures -> Bool
< :: NoFeatures -> NoFeatures -> Bool
$c< :: NoFeatures -> NoFeatures -> Bool
compare :: NoFeatures -> NoFeatures -> Ordering
$ccompare :: NoFeatures -> NoFeatures -> Ordering
Ord, NoFeatures
forall a. a -> a -> Bounded a
maxBound :: NoFeatures
$cmaxBound :: NoFeatures
minBound :: NoFeatures
$cminBound :: NoFeatures
Bounded, Int -> NoFeatures -> ShowS
[NoFeatures] -> ShowS
NoFeatures -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [NoFeatures] -> ShowS
$cshowList :: [NoFeatures] -> ShowS
show :: NoFeatures -> String
$cshow :: NoFeatures -> String
showsPrec :: Int -> NoFeatures -> ShowS
$cshowsPrec :: Int -> NoFeatures -> ShowS
Show)

instance Ix NoFeatures where
  range :: (NoFeatures, NoFeatures) -> [NoFeatures]
range   (NoFeatures, NoFeatures)
_   = []
  inRange :: (NoFeatures, NoFeatures) -> NoFeatures -> Bool
inRange (NoFeatures, NoFeatures)
_ NoFeatures
_ = Bool
False
  index :: (NoFeatures, NoFeatures) -> NoFeatures -> Int
index   (NoFeatures, NoFeatures)
_ NoFeatures
_ = -Int
1

noFeatures :: NoFeatures -> a
noFeatures :: forall a. NoFeatures -> a
noFeatures NoFeatures
_ = forall a. HasCallStack => String -> a
error String
"noFeatures"