module AutoComplete
  ( autoComplete,
  )
where


import qualified Data.Map.Strict as M
import qualified Data.Trie as T
import qualified Data.ByteString.Char8 as B
import Data.Ord (comparing)
import Data.Maybe(fromJust)
import Data.List (maximumBy)

import BuildKnowledgeBase (processSentence)

chunkSize :: Int
chunkSize = 2000

iOChunkSize :: Int
iOChunkSize = 5000

readNGramFreqs :: B.ByteString -> M.Map [B.ByteString] Int
readNGramFreqs contents = M.fromList $ map getListTuple (B.lines contents)
  where
    getListTuple x = case B.words x of
      [freq, a, b, c] -> ([a, b, c], fst $ fromJust $ B.readInt freq)
      [freq, a, b] -> ([a, b], fst $ fromJust $ B.readInt freq)
      _ -> ([], 0)

buildTrie :: B.ByteString -> T.Trie Int
buildTrie contents = T.fromList $ map getTuple (B.lines contents)
  where
    getTuple s = case B.words s of
      [freq, a] -> (a, fst $ fromJust $ B.readInt freq)
      _ -> (B.empty, 0)

autoComplete :: String -> String -> IO ()
autoComplete knowledgeBaseDir inputFile = do
  -- INPUT
  ifile <- readFile inputFile
  let input = processSentence ifile
  let prefix = B.pack $ last input
  let g2 = [B.pack $ last $ init input]
  let g3 = map B.pack $ drop (length input - 3) $ init input

  -- TRIE
  contents <- B.readFile $ fname 1
  let trie = buildTrie contents

  -- WORDS
  let words = T.toList $ T.submap prefix trie

  -- N-GRAMS
  f2 <- B.readFile $ fname 2
  f3 <- B.readFile $ fname 3
  let m2 = readNGramFreqs f2
  let m3 = readNGramFreqs f3

  let scores = zipWith3 comb words (calcF m2 g2 words) (calcF m3 g3 words)
  print $ fst $ maximumBy (comparing snd) scores

  where
    fname x = knowledgeBaseDir ++ "/" ++ show x ++ ".txt"
    makeNGrams l x = l ++ [fst x]
    getFreq m t = case M.lookup t m of
      Just f -> (t, f)
      Nothing -> (t, 0)
    comb (a1,a2) (b1,b2) (c1,c2) = (a1, b2+c2)
    makeGrams words g = map (makeNGrams g) words
    calcF m g w = map (getFreq m) (makeGrams w g)