{-# LANGUAGE ParallelListComp #-}
module Exercise_4 where
import Data.List hiding (stripPrefix)
import Data.Maybe
import Control.Arrow
import Control.Monad
import Control.Parallel.Strategies

-- Eliminates the longest common prefix from two words.
stripPrefix :: (String, String) -> (String, String)
stripPrefix (x : xs, y : ys) | x == y = stripPrefix (xs, ys)
stripPrefix z = z

-- Computes the edit distance between two words.
editDistance :: String -> String -> Int
editDistance = editDistanceBound Nothing

-- Computes the edit distance, optionally subject to a given bound. If the real edit distance is strictly
-- greater than the bound ub, the function may return ub + 1.
editDistanceBound :: Maybe Int -> String -> String -> Int
editDistanceBound ub xs ys = editDistanceBoundAux ub xs' ys'
  where -- Eliminate common prefixes and suffixes
        (xs', ys') = (reverse *** reverse) . stripPrefix . (reverse *** reverse) . stripPrefix $ (xs, ys)

outsideBound :: Int -> Maybe Int -> Bool
outsideBound _ Nothing  = False
outsideBound y (Just x) = y > x

editDistanceBoundAux :: Maybe Int -> String -> String -> Int
editDistanceBoundAux _  []  ys  = length ys
editDistanceBoundAux _  xs  []  = length xs
editDistanceBoundAux _  [x] ys  = length ys - if x `elem` ys then 1 else 0
editDistanceBoundAux _  xs  [y] = length xs - if y `elem` xs then 1 else 0
editDistanceBoundAux ub xs  ys
  | abs (m - n) `outsideBound` ub = fromJust ub + 1
  | otherwise = go (2, xs, row0, row1)
  where m = length xs
        n = length ys
        row0 = [0..n]
        row1 = 1 : let k = length (takeWhile (/= head xs) ys) in [1..k] ++ [k..n-1]
        go (i, xTwoLeft : xLeft : xs', butLastRow, lastRow)
          | minimum (butLastRow ++ lastRow) `outsideBound` ub = fromJust ub + 1
          | otherwise =
              let k = let k' = lastRow !! 1 in k' + if head ys == xLeft && k' == i - 1 then 0 else 1
                  row = i : k : [minimum $ [aLeft + 1,
                                            aBelow + 1,
                                            aBelowLeft + if xLeft == yLeft then 0 else 1] ++
                                            [aTwoBelowLeft + 1 | xLeft == yTwoLeft && xTwoLeft == yLeft]
                           | aLeft <- tail row | aBelow <- drop 2 lastRow | aBelowLeft <- tail (lastRow) 
                           | aTwoBelowLeft <- butLastRow | yLeft <- tail ys | yTwoLeft <- ys]
              in  go (i + 1, xLeft : xs', lastRow, row)
        go (_, _, _, lastRow) = last lastRow
{-WETT-}
-- Spell-checks a list of words against the given dictionary in parallel
spellCorrect :: [String] -> [String] -> [[String]]
spellCorrect d xs = map go xs `using` parListChunk 500 rdeepseq
  where go x = case d of
                 [] -> []
                 (y : ys) -> reverse $ fst $ foldl go' ([y], editDistance x y) ys
          where go' (acc, ub) y =
                  let dst = editDistanceBound (Just ub) x y
                  in  case compare dst ub of
                        LT -> ([y], dst)
                        GT -> (acc, ub)
                        EQ -> (y : acc, dst)
{-TTEW-}