module Exercise_4 where
import Data.List hiding (insert)
import Data.Array
import Data.Ord (comparing)
import qualified Data.IntMap as M
import Data.Maybe
import Control.Monad

{-WETT-}
arrFromList xs = listArray (0,length xs - 1) xs 

editDistanceArr a b = ds ! (n,m)
  where n = (snd $ bounds a) + 1
        m = (snd $ bounds b) + 1 
        d i 0 = i
        d 0 j = j
        d i j
          | a ! (i-1) ==  b ! (j-1) = ds ! (i-1,j-1)
          | otherwise = 1 + minimum [ ds ! (i-1,j)
                                    , ds ! (i,j-1)
                                    , ds ! (i-1,j-1)
                                    , if i > 1 && j > 1 && a ! (i-1) == b ! (j-2) && a ! (i-2) == b ! (j-1) then ds ! (i-2,j-2) else n+m
                                    ]

        ds = array ((0,0),(n,m)) [ ((i,j),d i j) | i <- [0..n], j <- [0..m]]


editDistance :: Eq a => [a] -> [a] -> Int
editDistance a b = editDistanceArr (arrFromList a) (arrFromList b)

{- Implementation of a BKTree -}
data Tree = Empty | Tree (Array Int Char) (M.IntMap Tree)
  deriving Show

insert :: Tree -> Array Int Char -> Tree
insert Empty w = Tree w M.empty
insert (Tree r ts) w =
  let d = editDistanceArr r w
  in case M.lookup d ts of
          Nothing -> Tree r (M.insert d (Tree w M.empty) ts)
          Just c -> Tree r (M.adjust (flip insert $ w) d ts)

query :: Int -> Array Int Char -> Tree -> [Array Int Char]
query _ _ Empty = []
query md w (Tree r ts) =
  let dist = editDistanceArr r w
      ts' = mapMaybe (`M.lookup` ts) [(dist-md)..(dist+md)] 
      ms = concatMap (query md w) ts'
  in if dist <= md then r:ms else ms 

findClosest :: Array Int Char -> Tree -> [Array Int Char]
findClosest _ Empty = []
findClosest w (Tree r ts) = fst $ findClosest' w ([r], editDistanceArr w r) (Tree r ts) 
    
findClosest' _ res Empty = res 
findClosest' w (cands,dist) (Tree r ts) = foldl' (findClosest' w) newCands children
  where
    newDist = editDistanceArr w r
    newCands = case compare dist newDist of
                 LT -> (cands,dist)
                 EQ -> (r:cands,dist)
                 GT -> ([r],newDist)
    children = mapMaybe (`M.lookup` ts) [dist-newDist..dist+newDist]
    
findClosest'' w t = aux 0
  where aux i = case query i w t of
                  [] -> aux (i+1)
                  res -> res

spellCorrect :: [String] -> [String] -> [[String]]
spellCorrect dict xs = 
  let dict' = map arrFromList dict
      xs' = map arrFromList xs
      bk = foldl' insert Empty dict'
  in map (\x -> (nub . map elems) $ findClosest'' x bk) xs'
{-TTEW-}