module KMPSeq (kmpMatch)

where

import qualified Data.Array.IArray as IArray ((!), Array, array)
import Data.List as List (map, filter, length, null, take, maximum, foldr, all, drop)
import qualified Data.ByteString.Lazy as B
import GHC.Int (Int64)

import Test.QuickCheck
import Data.ByteString as S (isSubstringOf, pack, ByteString) -- Only used during testing
import Debug.Trace (trace)
import System.IO.Unsafe
import System.IO

{-
 Returns the first index for which the pattern given is found
 in the string given, or -1 if the pattern does not appear.
 -}
kmpMatch :: B.ByteString -- ^ The pattern
  -> B.ByteString -- ^ The string to search
  -> Int64 -- ^ Result
kmpMatch pat search
  | B.null pat = 0 -- Empty pattern matches right away.
  | otherwise =
    let lookup = computeLookup pat
        patLen = fromIntegral $ B.length pat
        -- currPatIdx is our position in the pattern - sometimes we don't
        -- have to match the whole pattern again.
        --
        -- currIdx is our position in the string. We B.drop elements from str
        -- on each iteration so have to track our position separately.
        kmpMatch' :: B.ByteString -> Int64 -> Int -> Int64
        kmpMatch' str currIdx patShift
          | B.null str = -1
          | otherwise =
              -- Find length of prefix in our pattern which matches string (possibly zero)
              let matchLength :: Int
                  matchLength =
                    let matchLength' pat str cnt | {-# SCC "kmpMatch_nullPatStr" #-} B.null pat = cnt
                        matchLength' pat str cnt | {-# SCC "kmpMatch_nullStr" #-} B.null str = 0
                        matchLength' pat str cnt | {-# SCC "kmpMatch_eqHead" #-} (B.head pat == B.head str) = matchLength' (B.drop 1 pat) (B.drop 1 str) (cnt `seq` (cnt + 1))
                                                 |   otherwise = cnt
                    in
                    {-# SCC "kmpMatch_longestPrefix'" #-} matchLength' (B.drop (fromIntegral patShift) pat) str 0
                  shiftAmt = {-# SCC "kmpMatch_shiftAmt" #-} lookup matchLength
                  patShift64 = fromIntegral patShift
                  matchLen64 = fromIntegral matchLength
              in
                case matchLength of
                  _ | matchLength + patShift == patLen  -> currIdx - patShift64 -- found complete match
                    | matchLength == 0 -> kmpMatch' (B.drop 1 str) (currIdx `seq` currIdx + 1) 0
                    | shiftAmt > 0 -> kmpMatch' (B.drop matchLen64 str) (currIdx `seq` currIdx + matchLen64) shiftAmt
                    | otherwise -> kmpMatch' (B.drop matchLen64 str) (currIdx `seq` currIdx + matchLen64) patShift
    in
        kmpMatch' search 0 0

-- These tests have failed in the past
-- kmpMatch (toLazyBS "ccb") (toLazyBS "abcdcccb") == 5
-- kmpMatch (toLazyBS "bbaa") (toLazyBS "bdcdbbbbaa") == 6
-- kmpMatch (toLazyBS "cccadadc") (toLazyBS "adaccccadadc") == 4
-- This next one has caused an infinite loop previously.
-- kmpMatch (toLazyBS "bbbb") (toLazyBS "dddcaaaddaaabdacbcccabbada")

{-|
 
 Given our pattern, get all the prefixes of the pattern. For each of those
 prefixes, find the longest prefix from the original pattern that is also a
 suffix of the prefix segment being considered, and is not equal to it. The
 argument given to overlap is the length of the prefix matched so far, and the
 length of the longest prefix, which is a suffix and is not equal to it, is the
 value overlap returns.
 
 If a given prefix has no possible overlap, it is mapped to -1.
 
-}
overlap :: B.ByteString -> [(Int, Int)]
overlap pat =
  let patternLength = B.length pat
      -- Necessary to reverse prefixes since 'isSuffixOf' not implemented for lazy bytestrings.
      prefixes = List.map B.reverse $ List.take (fromIntegral patternLength) $ B.inits pat
      longestPreSuffix prefix =
        -- Find prefixes that are a prefix of this suffix, but
        -- are not equal to it. 
        let suffixes = List.filter (\p -> p `B.isPrefixOf` prefix && p /= prefix) prefixes
        in
          if B.null prefix || List.null suffixes
            then 0
            else List.maximum (List.map B.length suffixes)
  in
    List.map (\prefix -> (fromIntegral $ B.length prefix, fromIntegral $longestPreSuffix prefix)) prefixes

{- |
 Given a string representing a search pattern, this function
 returns a function which represents, for each prefix of that
 pattern, the maximally long prefix of the pattern which is a suffix
 of the indicated pattern segment.
 
 If there is no such prefix, 0 is returned.
 -}
computeLookup :: B.ByteString -> (Int -> Int)
computeLookup pat =
  let patLen = fromIntegral $ B.length pat
      table :: IArray.Array Int Int
      table = {-# SCC "computeLookup_table" #-} IArray.array (0, patLen - 1) (overlap pat)
  in
    (IArray.!) table

-- Types, instances and utility functions for testing purposes.
    
newtype PatternChar = PatternChar Char
  deriving Show
  
instance Arbitrary PatternChar where
  arbitrary = oneof (List.map (return . PatternChar) ['a', 'b', 'c', 'd'])

patternsToString :: [PatternChar] -> B.ByteString
patternsToString chars = B.pack $ List.foldr (\(PatternChar char) str -> (toEnum $ fromEnum char) : str) [] chars

patternsToStrictString :: [PatternChar] -> S.ByteString
patternsToStrictString chars = S.pack $ List.foldr (\(PatternChar char) str -> (toEnum $ fromEnum char) : str) [] chars

toLazyBS :: String -> B.ByteString
toLazyBS = B.pack . List.map (toEnum . fromEnum)

toStrictBS :: String -> S.ByteString
toStrictBS = S.pack . List.map (toEnum . fromEnum)

-- Test that 0 and 1 element always return 0, if present.
prop_testZero :: [PatternChar] -> Property
prop_testZero pat =
  let lookup = computeLookup (patternsToString pat)
  in
    not (List.null pat) ==>
      if List.length pat > 1
        then lookup 0 == 0 && lookup 1 == 0
        else lookup 0 == 0 
    
-- Test that all overlaps found are actually prefixes of the
-- pattern string
prop_testSubset :: [PatternChar] -> Property
prop_testSubset pat =
  let patStr = patternsToString pat
      lookup = computeLookup patStr
      prefix len = B.take (fromIntegral len)
      testPrefix len =
        if lookup len == 0
          then True
          else (prefix (lookup len) patStr) `B.isPrefixOf` (prefix len patStr)
  in
    not (List.null pat) ==>
      trivial (B.null patStr) $
        List.all testPrefix [0 .. List.length pat - 1]


-- Test that the prefix given is the maximal prefix. That is,
-- add one more character makes it either equal to the string
-- or not a prefix.
prop_testCorrectPrefix :: [PatternChar] -> Property
prop_testCorrectPrefix pat =
  let patStr = patternsToString pat
      lookup = computeLookup patStr
      isNeverSuffix len =
        let origPrefix = prefix len patStr
            -- Drop 1 to remove empty list
            allPrefixes = List.drop 1 $ B.inits origPrefix
        in
          List.all (\p -> B.null p || p == origPrefix || not ((B.reverse p) `B.isPrefixOf` (B.reverse origPrefix))) allPrefixes
      prefix len = B.take (fromIntegral len)
      -- True if the prefix returned from lookup for the length given produces
      -- a string which is a suffix of the original prefix.
      isRealSuffix len = (B.reverse (prefix (lookup len) patStr)) `B.isPrefixOf` (B.reverse $ prefix len patStr)
      isLongestSuffix len =
        let prefixPlus = prefix ((lookup len) + 1) patStr
            inputPrefix =  prefix len patStr
        in
          prefixPlus == inputPrefix ||
            not ((B.reverse prefixPlus) `B.isPrefixOf` (B.reverse inputPrefix))
      testTable len =
        if lookup len == 0
          then isNeverSuffix len
          else isRealSuffix len &&
                isLongestSuffix len
  in
    not (List.null pat) ==>
      List.all testTable [0 .. List.length pat - 1]
    
-- Verify that, if a match is found, it is where it's supposed to be in
-- the string and it can be independently found by other means.
prop_testMatch :: [PatternChar] -> [PatternChar] -> Property
prop_testMatch pat search =
  let patStr = patternsToString pat
      searchStr = patternsToString search
      strictStr = patternsToStrictString search
      strictPat = patternsToStrictString pat
      patLen = B.length patStr
      searchLen = B.length searchStr
      matchIdx = kmpMatch patStr searchStr
      verify =
        if matchIdx > -1
          then (B.take patLen $ B.drop matchIdx $ searchStr) == patStr && strictPat `S.isSubstringOf` strictStr
          else (B.null patStr && B.null searchStr) || not (strictPat `S.isSubstringOf` strictStr)
  in
    not (List.null pat) ==>
      trivial (B.null searchStr) $
        classify (patLen > searchLen) "Bigger pattern than search" $
        classify (patLen < searchLen) "Smaller pattern than search" $
        classify (patLen == searchLen) "Equal pattern and search" $
        classify (matchIdx > -1) "Match Exists" $
        classify (matchIdx == -1) "Match Doesn't Exist" $
          verify
        
-- Test a pattern that was known to fail.        
prop_testBadPat :: [PatternChar] -> Property
prop_testBadPat search =
  let patStr = toLazyBS "bbc"
      patLen = B.length patStr
      searchStr = patternsToString search
      matchIdx = kmpMatch patStr searchStr
      strictStr = patternsToStrictString search
      strictPat = toStrictBS "bbc"
      verify =
        if matchIdx > -1
          then (B.take patLen $ B.drop matchIdx $ searchStr) == patStr && strictPat `S.isSubstringOf` strictStr
          else (B.null patStr && B.null searchStr) || not (strictPat `S.isSubstringOf` strictStr)
  in
    trivial (List.null search) $
      verify

-- Test that a pattern on the end of the string is found OK.
prop_testEndPattern :: [PatternChar] -> [PatternChar] -> Property
prop_testEndPattern pat search =
  let patStr = patternsToString pat
      searchStr = patternsToString (search ++ pat)
      matchIdx = kmpMatch patStr searchStr
      strictStr = patternsToStrictString (search ++ pat)
      strictPat = patternsToStrictString pat
      patLen = B.length patStr
      verify =
        if matchIdx > -1
          then (B.take patLen $ B.drop matchIdx $ searchStr) == patStr && strictPat `S.isSubstringOf` strictStr
          else (B.null patStr && B.null searchStr) || not (strictPat `S.isSubstringOf` strictStr)
  in
    not (List.null pat) ==>
      trivial (B.null searchStr) $
        classify (matchIdx > -1) "Match Exists" $
        classify (matchIdx == -1) "Match Doesn't Exists" $
          verify