{-# LANGUAGE DeriveGeneric #-}
module OkapiBM25_utils where
import           Control.Parallel.Strategies
import           Data.Char
import qualified Data.Heap                     as Heap
import qualified Data.Map.Strict               as Map
import           Data.Maybe
import qualified Data.Set                      as Set
import           Data.Tuple
import           GHC.Generics                   ( Generic )


-- termFrequency: frequency of a word in this document
-- termSet: all the words in this Document (i.e. all keys of the termFrequency map)
data Document = Document
  { title         :: String
  , text          :: String
  , termFrequency :: Map.Map String Int
  , termSet       :: Set.Set String
  , docLen        :: Int
  }
  deriving Generic

scoreAllDocs
  :: Int
  -> String
  -> [Document]
  -> Map.Map String Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> [(Double, Document)]
scoreAllDocs chunks query docs dfMap corpusSize avgDocLen k b =
  map
      (\doc -> Data.Tuple.swap
        (getDocWithScore (tokenizeAndNormalize query)
                         doc
                         dfMap
                         corpusSize
                         avgDocLen
                         k
                         b
        )
      )
      docs
    `using` parListChunk chunks rseq

getDocWithScore
  :: [String]
  -> Document
  -> Map.Map String Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> (Document, Double)
getDocWithScore query doc dfMap corpusSize avgDocLen k b =
  (doc, sum [ individualScore q | q <- query ]) where
  individualScore :: String -> Double
  individualScore queryWord =
    calcIDF queryWord
      * fromIntegral (calcTF queryWord)
      * (k + 1.0)
      / ( fromIntegral (calcTF queryWord)
        + k
        * (1.0 - b + b * fromIntegral corpusSize / fromIntegral avgDocLen)
        )
  calcTF :: String -> Int
  calcTF queryWord = fromMaybe 0 (Map.lookup queryWord (termFrequency doc))
  calcIDF :: String -> Double
  calcIDF queryWord
    | Map.member queryWord dfMap = ln (fromIntegral corpusSize + 1)
    - ln (fromIntegral (fromMaybe 0 (Map.lookup queryWord dfMap)) + 0.5)
    | otherwise = ln (fromIntegral corpusSize * 2)
    where ln = logBase (exp 1)

readDocument :: String -> String -> Document
readDocument title' text' = Document title'
                                     text'
                                     tfMap
                                     (Map.keysSet tfMap)
                                     (sum (Map.elems tfMap))
  where tfMap = buildTFMap text'

-- Get the document count (DF) for each word
getAllDF :: [Document] -> Map.Map String Int
getAllDF docs = Map.fromListWith
  (+)
  [ tup | doc <- docs, tup <- zip (Set.toList (termSet doc)) [1 :: Int, 1 ..] ]

topKDocuments
  :: Int -> [(Double, Document)] -> Heap.MinPrioHeap Double Document
topKDocuments _ [] = Heap.fromList [] :: Heap.MinPrioHeap Double Document
topKDocuments topK (sd : sds) = maybePop
  (Heap.insert sd (topKDocuments topK sds)) where
  maybePop
    :: Heap.MinPrioHeap Double Document -> Heap.MinPrioHeap Double Document
  maybePop minHeap | Heap.size minHeap > topK = Heap.drop 1 minHeap
                   | otherwise                = minHeap

buildTFMap :: String -> Map.Map String Int
buildTFMap docText =
  Map.fromListWith (+) [ (word, 1) | word <- tokenizeAndNormalize docText ]

-- Return a list of normalized tokens
tokenizeAndNormalize :: String -> [String]
tokenizeAndNormalize query = map sanitizeWord (words query)
  where sanitizeWord word = map toLower (filter isLetter word)
