{-
  Module containing the Rabin-Karp algorithm.
-}
module Search
    ( rabinKarpMain
    ) where

import Data.Array ( (!), Array )
import qualified Data.Word8 as W

-- No of alphabets in input
d :: Int
d = 26

-- Prime number used in hash function
q :: Int
q = 13

{-
  Calculate the hash value of the given word.
-}
calcHashPatt :: Int -> [W.Word8] -> Int
calcHashPatt = foldl (\ h x -> ((h * d) + fromEnum x) `mod` q)

{-
  Calculate the hash value of a string present in the grid.
  If 'isRow' is True, calculate the hash value of a string at arr(fixedIdx,f) of size (l-f+1)
  If 'isRow' is False, calculate the hash value of a string at arr(f, fixedIdx) of size (l-f+1)
-}
calcHash :: Int -> Int -> Int -> Array (Int, Int) W.Word8 -> Int -> Bool -> Int
calcHash fixedIdx f l text h isRow
  | f == l = h
  | isRow = calcHash fixedIdx (f+1) l text (((h * d) + fromEnum (text ! (fixedIdx, f))) `mod` q) isRow
  | otherwise = calcHash fixedIdx (f+1) l text (((h * d) + fromEnum (text ! (f, fixedIdx))) `mod` q) isRow

{-
  Check if a substring present in grid matches the pattern.
  If 'isRow' is True, compare pattern with string starting at arr(fixedIdx,baseIdx) of size (l-f+1)
  If 'isRow' is False, compare pattern with string starting at arr(baseIdx,fixedIdx) of size (l-f+1)
  Returns a boolean.
-}
areEqual :: Int -> Int -> Int -> [W.Word8] -> Array (Int, Int) W.Word8 -> Int -> Bool -> Bool
areEqual baseIdx f l patt text fixedIdx isRow
  | f == l = True
  | isRow = ((patt !! f) == (text ! (fixedIdx, baseIdx+f))) &&
      areEqual baseIdx (f+1) l patt text fixedIdx isRow
  | otherwise = ((patt !! f) == (text ! (baseIdx+f, fixedIdx))) &&
      areEqual baseIdx (f+1) l patt text fixedIdx isRow

{-
  Calculate the new hash value of a string in grid from existing hash by sliding by 1 character.
-}
calcSlidingHash :: Int -> Int -> Int -> Int -> Int -> Array (Int, Int) W.Word8 -> Int -> Bool -> Int
calcSlidingHash tHash baseIdx l m h text fixedIdx isRow
  | baseIdx < l = let
              nHash = if isRow then
                        (d * (tHash - fromEnum ( text ! (fixedIdx, baseIdx) ) * h) +
                         fromEnum(text ! (fixedIdx, baseIdx+m))) `mod` q
                      else
                        (d * (tHash - fromEnum ( text ! (baseIdx, fixedIdx) ) * h) +
                         fromEnum(text ! (baseIdx+m, fixedIdx))) `mod` q
            in
              if nHash < 0 then
                nHash + q
              else
                nHash
  | otherwise = tHash

{-
  Update the result with matched index if the word matches with the pattern.
  Returns updated result list.
-}
findMatch :: [W.Word8] -> Array (Int, Int) W.Word8 -> Int -> Int -> Int -> Bool -> [Int] -> Int -> [Int]
findMatch patt text i fixedIdx m isRow res matched_idx
 = let isMatch = areEqual i 0 (m-1) patt text fixedIdx isRow
       new_res = if isMatch then
                  matched_idx : res
                 else
                  res
   in new_res

{-
  The main Rabin-Karp algorithm that iterates through a row/column
  and compares the hash value of words with the pattern and it's reverse.
  It returns a list of indices of matches.
-}
rabinKarp :: Int -> Int -> Int -> Array (Int, Int) W.Word8 -> [W.Word8] -> [W.Word8] -> Int
            -> Int -> [Int] -> [Int] -> Int -> Int -> Int -> Bool -> [Int]
rabinKarp tHash pHash pRevHash text patt revPatt i l res revRes m h fixedIdx isRow
    | i == l = res ++ revRes
    | tHash == pHash = let
                          new_res = findMatch patt text i fixedIdx m isRow res i
                          new_tHash = calcSlidingHash tHash i l m h text fixedIdx isRow
                       in
                          rabinKarp new_tHash pHash pRevHash text patt revPatt (i+1) l new_res revRes m h fixedIdx 
                          isRow
    | tHash == pRevHash = let
                           new_res = findMatch revPatt text i fixedIdx m isRow revRes (i+m-1)
                           new_tHash = calcSlidingHash tHash i l m h text fixedIdx isRow
                          in
                            rabinKarp new_tHash pHash pRevHash text patt revPatt (i+1) l res new_res m h fixedIdx 
                            isRow
    | otherwise = let new_tHash = calcSlidingHash tHash i l m h text fixedIdx isRow
                  in
                    rabinKarp new_tHash pHash pRevHash text patt revPatt (i+1) l res revRes m h fixedIdx isRow

{-
  Calculates the hash of pattern, reversed pattern and first word on the row/column.
  It then calls the main Rabin-Karp algorithm which returns matches.
  Returns a list of (x,y) coordinates of matches.
-}
rabinKarpMain :: [W.Word8] -> [W.Word8] -> Array (Int, Int) W.Word8 -> Int -> Int -> Bool -> Int -> [(Int, Int)]
rabinKarpMain patt revPatt text m n isRow fixedIdx =
                                                      let tHash = calcHash fixedIdx 0 m text 0 isRow
                                                          pHash = calcHashPatt 0 patt
                                                          pRevHash = calcHashPatt 0 revPatt
                                                          h = foldl (\acc _ -> (acc*d) `mod` q) 1 $ replicate 1 (m-1)
                                                          res = rabinKarp tHash pHash pRevHash text patt revPatt 
                                                                0 (n-m+1) [] [] m h fixedIdx isRow
                                                      in
                                                        if isRow then
                                                          zip (repeat fixedIdx) res
                                                        else
                                                          zip res (repeat fixedIdx)
