{-

usage: ./boggle /usr/share/dict/words random

benchmarks:
  compiled with            runtime  "dictionary" traced to stderr (line 62)
  ghc -O2                   29.5s   100x
  ghc -O2 -fno-state-hack    0.5s     1x

-}
import Control.Monad (forever, replicateM, replicateM_)
import Data.Array
import Data.Char (isLower, toLower, toUpper)
import Data.List (foldl')
import qualified Data.Map.Strict as M
import Data.Maybe (mapMaybe, catMaybes)
import Data.Map.Strict (Map)
import qualified Data.Set as S
import System.Environment (getArgs)
import System.Random (randomRIO)

import Debug.Trace (trace)

main = do
  dictFile:args <- getArgs
  dict <- parseDict `fmap` readFile dictFile
  case args of
    [] -> interact (main' dict)
    _ -> replicateM_ 100 $ putStr . main' dict =<< randomBoard

randomBoard = do
  b <- replicateM 16 (randomRIO ('A', 'Z'))
  putStr . unlines . map expandQu' . chunk 4 $ b
  return (expandQu b)

chunk n [] = []
chunk n xs = let (ys, zs) = splitAt n xs in ys : chunk n zs

main' dict = unlines . mapMaybe (fmap (unwords . longestFirst . boggle dict) . parseBoard) . lines

boggle dict0 board = concatMap (go dict0 emptyUsed "") . indices $ board
  where
    go dict used word ix
      | used ! ix = []
      | otherwise = case lookupTrie (board ! ix) dict of
          Nothing -> []
          Just dict' ->
            let word' = (board ! ix) : word
                used' = used // [(ix, True)]
            in  (if isWord dict' then (reverse word' :) else id)
                  (concatMap (go dict' used' word') (neighbours ! ix))

bogBounds = ((0,0),(3,3))
emptyUsed = array bogBounds [(ix,False) | ix <- range bogBounds ]
neighbours = array bogBounds
  [ ((i, j), [ (i',j') | di <- [-1,0,1], let i' = i + di, dj <- [-1,0,1], let j' = j + dj
    , inRange bogBounds (i', j')
    ])
  | (i, j) <- range bogBounds
  ]

parseDict = trace "dictionary" . buildTrie . mapMaybe canonical . filter (all isLower) . lines
parseBoard = fmap (listArray bogBounds) . (checkLength =<<) . canonical
  where checkLength l | length l == 16 = Just l
                      | otherwise = Nothing

canonical = compressQu . filter isLetter . map toUpper
isLetter c = 'A' <= c && c <= 'Z'

expandQu = concatMap qu
  where qu 'Q' = "QU"
        qu c = [c]
expandQu' = concatMap qu
  where qu 'Q' = "Qu"
        qu c = [c, ' ']
compressQu ('Q':'U':xs) = ('Q' :) `fmap` compressQu xs
compressQu ('Q':_) = Nothing
compressQu (x:xs) = (x :) `fmap` compressQu xs
compressQu [] = Just []

data Trie = Trie{ isWord :: !Bool, nextLetters :: !(Map Char Trie) }
emptyTrie = Trie{ isWord = False, nextLetters = M.empty }
buildTrie = foldl' (flip insertTrie) emptyTrie
insertTrie [] t = t{ isWord = True }
insertTrie (c:cs) t = case lookupTrie c t of
  Nothing -> t{ nextLetters = M.insert c (insertTrie cs emptyTrie) (nextLetters t) }
  Just t' -> t{ nextLetters = M.insert c (insertTrie cs t'       ) (nextLetters t) }
lookupTrie c t = M.lookup c (nextLetters t)

longestFirst = map snd . S.toAscList . S.fromList . mapMaybe (short . expandQu)
short l = case length l of
  m | m < 3 -> Nothing
    | otherwise -> Just (-m, map toLower l)