import           Control.Monad.Reader
import           Control.Parallel.Strategies
import           Data.Char
import qualified Data.Heap                     as Heap
import           Data.Time
import           OkapiBM25_utils
import           Prelude
import           System.Environment             ( getArgs
                                                , getProgName
                                                )
import           System.Exit                    ( die )

main :: IO ()
main = do
  (file, query, topK, chunks) <- getArgs >>= checkArgs
  output                      <- runOkapiBM25 file
                                              (filterLetters query)
                                              (read topK :: Int)
                                              (read chunks :: Int)
  mapM_ print output
 where
  checkArgs [file, query, topK, chunks] = return (file, query, topK, chunks)
  checkArgs _                           = do
    progName <- getProgName
    die
      $  "Usage: ./"
      ++ progName
      ++ " <filename> <query> <K (no. of top documents)> <no. of chunks>"
  filterLetters = filter (\c -> isLetter c || isSpace c || c == '\'')


runOkapiBM25 :: String -> String -> Int -> Int -> IO [(String, Double)]
runOkapiBM25 file query topK chunks = do
  documents' <- getAllDocuments file chunks
  documents  <- sequence documents'
  let corpusSize = length documents -- N
  let (avgDocLen, dfMap) = runEval $ do
        avgDocLen' <-
          rpar $ (sum [ docLen doc | doc <- documents ]) `div` corpusSize
        dfMap' <- rpar $ getAllDF documents
        _      <- rseq avgDocLen'
        _      <- rseq dfMap'
        return (avgDocLen', dfMap')

  let (k, b) = (1.2, 0.75)
  startTime <- documents `seq` getCurrentTime
  let scoredDocs =
        scoreAllDocs chunks query documents dfMap corpusSize avgDocLen k b
  let result = parseOutput (Heap.toAscList (topKDocuments topK scoredDocs))
  endTime <- result `seq` getCurrentTime
  print $ diffUTCTime endTime startTime
  return result

getAllDocuments :: String -> Int -> IO [IO Document]
getAllDocuments file chunks = do
  fileText <- readFile file -- titles.txt
  let documents =
        [ fmap (\(text', title'') -> readDocument title'' text')
               (readFile path >>= return . \fileTxt -> (fileTxt, title'))
        | line <- lines fileText
        , let (path, title') = splitAtFirst ',' line
        ]
        `using` parListChunk chunks rseq
  return documents
  where splitAtFirst x = fmap (drop 1) . break (x ==)

parseOutput :: [(Double, Document)] -> [(String, Double)]
parseOutput inputs = reverse [ (title doc, score) | (score, doc) <- inputs ]
