-- | A word search solver library
--
-- This solver is case sensitive; users should map data consistently to one case before using this library when such behavior is desired.

{-# LANGUAGE TemplateHaskell, DeriveDataTypeable #-}

module Text.WordSearchSolver (
                             -- * Types and containers
                               WordSearch, ws_grid, ws_search
                             , Grid
                             , Search
                             , Pos(..)
                             , PosIndex
                             , Match(..), m_dir, m_len, m_pos
                             , Dir(..)

                             -- * 'WordSearch' puzzles
                             , readWordSearch
                             , wordSearch
                             , solveWordSearch
                             , search
                             , tryMatch

                             -- * 'Grid' and 'Search' containers
                             , readGrid
                             , arrayToGrid
                             , setToSearch
                             , searchToSet

                             -- * Operations on solutions and rendering 'Grid's
                             , fillMatches
                             , showGridInsert

                             -- * 'Dir's
                             , dirs
                             , dirs'
                             , dirsPos
                             , dirsOpposite
                             , dirToOffset
                             , dirOpposite
                             , dirUpdatePos

                             -- * Helper functions
                             , inRangeOf
                             , posPlus
                             ) where

import Prelude hiding (id, (.), foldr, foldl, concat)
import Control.Category
import Control.Monad hiding (forM, forM_)
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Array
import Data.Array.MArray
import Data.Array.ST
import Data.Data (Data)
import Data.Foldable
import Data.List hiding (foldr, foldl, foldl', concat)
import Data.Record.Label
import Data.Typeable (Typeable)

--- Types and containers ---

-- | Abstract container of a word search puzzle
--
-- This can be created either from a 'Search' and a 'Grid' by the
-- 'wordSearch' function or from a properly formatted 'String' by the
-- 'readWordSearch' function.
data WordSearch a = WordSearch { _ws_grid   :: Grid a
                               , _ws_search :: Search a
                               } deriving (Data, Typeable, Eq, Ord, Show, Read)

-- | A grid in which to search
--
-- Constructors of this container usually assume that the grid is rectangular and properly sized; this precondition is \not\ checked.
data Grid a = Grid { _g_array :: Array Pos a
                   , _g_index :: GridIndex a
                   } deriving (Data, Typeable, Eq, Ord, Show, Read)

newtype GridIndex a = GridIndex { _gi_map :: M.Map a (S.Set Pos)
                                } deriving (Data, Typeable, Eq, Ord, Show, Read)

-- | A set of words or lists to search
newtype Search a = Search { _s_set :: S.Set [a]
                          } deriving (Data, Typeable, Eq, Ord, Show, Read)

-- | A position of a grid
newtype Pos = Pos (PosIndex, PosIndex) deriving (Data, Typeable, Eq, Ord, Ix, Show, Read)

-- | The integral type used for 'Pos'
type PosIndex = Integer

-- | An individual value describing a match
data Match = Match { _m_dir :: Dir
                   , _m_len :: Integer
                   , _m_pos :: Pos
                   } deriving (Data, Typeable, Eq, Ord, Show, Read)

data Dir = N
         | NW
         | W
         | SW
         | S
         | SE
         | E
         | NE
         deriving (Data, Typeable, Eq, Enum, Ord, Ix, Show, Read)

$(mkLabels [''WordSearch, ''Grid, ''GridIndex, ''Search, ''Match])


-- | Constructs a 'WordSearch' container from a properly formatted 'String'
--
-- The 'String' should contain two sections, separated by at least one empty
-- line.  The first section represents the grid, and thus is formatted as the
-- String that 'readGrid' expects.  The second section represents the search
-- words; it contains each word on its own separate line.  In the case that
-- is ill-formed, Nothing is returned.  The precondition that each grid row
-- has equal length is \not\ checked.

readWordSearch :: String -> Maybe (WordSearch Char)
readWordSearch xs = do
    let (former, latter) = span (not . null) . lines $ xs
        grid        = readGrid $ former
        searchTerms = Search . S.fromList . filter (not . null) $ dropWhile null latter
    guard $ (not . null $ former)
    return $ wordSearch grid searchTerms

-- | Constructs a 'WordSearch' container from a 'Grid' and a 'Search'
wordSearch :: Grid a -> Search a -> WordSearch a
wordSearch g s = WordSearch {_ws_grid = g, _ws_search = s}

-- | Constructs a 'Grid' from a formatted list of 'String's
--
-- The 'String' should be formatted as expected:
--
-- @[\"aoaoenxrcoedxncd\",
-- \"aoesnitdaoeusntd\",
-- \"itenohtneahuoteh\"]@
--
-- The precondition that each row has equal length is \not\ checked.
readGrid :: [String] -> Grid Char
readGrid xs = arrayToGrid . listArray (Pos (0, 0), Pos (pred . genericLength . head $ xs, pred . genericLength $ xs)) . concat . transpose $ xs

-- | Constructs a 'Grid' from an 'Array'
--
-- A grid array is indexed by (column, row) or (x, y), so users should be
-- careful that a list is in the proper order if listArray is used.
arrayToGrid :: (Ord a) => Array Pos a -> Grid a
arrayToGrid a = Grid { _g_array = a
                     , _g_index = gridIndex a
                     }
    where gridIndex arr = GridIndex $ foldr step M.empty $ assocs arr
          step ~(p, v) acc = M.insertWith (flip S.union) v (S.singleton p) acc

-- | Constructs a 'Search' from a set of lists
setToSearch :: S.Set [a] -> Search a
setToSearch = Search

-- | Returns the set of search terms from a 'Search' container
searchToSet :: Search a -> S.Set [a]
searchToSet = getL s_set

-- | Solves a 'WordSearch' and returns a set of matches together with a set of search terms for which a match was not found in a tuple
--
-- This algorithm solves word search puzzles by looking at the first cell of each search term, and looking for a match by checking each direction from each position whose cell contains the starting cell of the search term until a match is found.  The dictionary of individual cell values and sets of positions is part of the 'Grid' container; 'arrayToGrid' creates this dictionary automatically.
solveWordSearch :: (Eq a, Ord a) => WordSearch a -> (S.Set Match, Search a)
solveWordSearch ws = let (matches, nonmatches) = foldr step (S.empty, S.empty) . getL (s_set) $ getL ws_search ws
                     in  (matches, Search nonmatches)
    where step s (matches, nonmatches) =  -- look at algorithm
              let match     = foldr (flip mplus . search (getL ws_grid ws) s) Nothing $ positions
                  positions = case M.lookup (head s) $ getL (gi_map . g_index . ws_grid) ws of
                      (Just ps) -> ps
                      (Nothing) -> S.empty
              in  case match of
                      (Just m)  -> (m `S.insert` matches, nonmatches)
                      (Nothing) -> (matches, s `S.insert` nonmatches)

-- | Determines whether a given 'Search' term can be matched at a given position of a grid
--
-- This is done by trying each direction for a match from the given location.
search :: (Eq a) => Grid a -> [a] -> Pos -> Maybe Match
search _    [] _ = Nothing
search grid xs p = foldr mplus Nothing $ [tryMatch grid xs p d | d <- dirs]

-- | If the location and the direction matches the 'Search' term, returns the 'Match'; otherwise, returns Nothing
tryMatch :: (Eq a) => Grid a -> [a] -> Pos -> Dir -> Maybe Match
tryMatch grid = tryMatch' 0
              where arr = getL g_array grid
                    tryMatch' i []     p d = let d' = dirOpposite d
                                             in  Just $ Match {_m_dir = d', _m_len = i, _m_pos = dirUpdatePos d' p}
                    tryMatch' i (x:xs) p d
                        | not $ p `inRangeOf` arr = Nothing
                        | arr ! p == x            = tryMatch' (succ i) xs (dirUpdatePos d p) d
                        | otherwise               = Nothing

-- | Creates a 'Grid' in which every cell that does not match is set to a default value
fillMatches :: (Foldable t, Ord e) => e -> Grid e -> t Match -> Grid e
fillMatches deft g ms = arrayToGrid . runSTArray $ do
    let a = getL (g_array) g
    ma <- thaw $  listArray (bounds a) $ repeat deft
    let fill ~m@(Match dir len pos)
            | len == 0          = return ()
            | pos `inRangeOf` a = do
                writeArray ma pos $ a ! pos
                fill . modL m_pos (dirUpdatePos dir) . setL (m_len) (pred len) $ m
            | otherwise         = return ()
    forM_ ms $ fill
    return ma

-- | Renders a 'Grid', appending a cell, usually a newline character, after every row
showGridInsert :: a -> Grid a -> [a]
showGridInsert ins g = let a = getL (g_array) g
                           ~((Pos (0, 0)), Pos (w, h)) = bounds a
                           step ~p@(x, _) acc
                               | x == w    = (a ! Pos p):ins : acc
                               | otherwise = a ! Pos p       : acc
                       in  foldr step [] $ [(x, y) | y <- [0..h], x <- [0..w]]

--- Dir ---
-- | Complete set of possible 'Grid' 'Match' directions
dirs :: [Dir]
dirs =
    [ N
    , NW
    , W
    , SW
    , S
    , SE
    , E
    , NE
    ]

-- | More efficient (and real, unordered) 'Set' of 'dirs'
dirs' :: S.Set Dir
dirs' = S.fromList dirs

-- | 'Map' of directions and 'Pos' offsets
dirsPos :: M.Map Dir Pos
dirsPos = M.fromList . map (\ ~(a, b) -> (a, Pos b)) $
    [ (N,  ( 0, -1))
    , (NW, (-1, -1))
    , (W,  (-1,  0))
    , (SW, (-1,  1))
    , (S,  ( 0,  1))
    , (SE, ( 1,  1))
    , (E,  ( 1,  0))
    , (NE, ( 1, -1))
    ]

-- | Bidirectional 'Map' of opposite directions
dirsOpposite :: M.Map Dir Dir
dirsOpposite = M.fromList $
    [ (N,  S)
    , (NW, SE)
    , (W,  E)
    , (SW, NE)
    , (S,  N)
    , (SE, NW)
    , (E,  W)
    , (NE, SW)
    ]

-- | Returns the appropriate offset of a direction
dirToOffset :: Dir -> Pos
dirToOffset d = case M.lookup d dirsPos of
    (Just p)  -> p
    (Nothing) -> error "dirToOffset: unrecognized direction"  -- This shouldn't happen, since the lookup should always succeed

-- | Returns the opposite direction
dirOpposite :: Dir -> Dir
dirOpposite d = case M.lookup d dirsOpposite of
    (Just d') -> d'
    (Nothing) -> error "dirOpposite: unrecognized direction"  -- This shouldn't happen, since the lookup should always succeed

-- | Updates a position by one step in the given direction
dirUpdatePos :: Dir -> Pos -> Pos
dirUpdatePos d = (`posPlus` dirToOffset d)

--- Helper Functions ---
-- | Determines whether an index is within the range of the bounds of an array
inRangeOf :: (Ix a) => a -> Array a e -> Bool
i `inRangeOf` a = inRange (bounds a) i

-- | Adds two positions
posPlus :: Pos -> Pos -> Pos
posPlus (Pos (ax, ay)) (Pos (ba, by)) = Pos (ax + ba, ay + by)