{-# LANGUAGE DeriveDataTypeable #-}

module Main where

import Control.Monad (when)
import Data.Array (Array, array, bounds, (!), (//))
import Data.Functor ((<$>))
import Data.Maybe
import Data.List (sort)
import Data.Ord (compare)
import Data.Char (toUpper)
import System.Console.CmdArgs (Data, Typeable, cmdArgs)
import System.Random (RandomGen, getStdGen, split, randomR, randomRs)
import System.Random.Shuffle

-- Vectors to represent position and orientation.
data Vec = Vec { x::Int, y::Int } deriving (Show, Eq)

origin = Vec 0 0
left = Vec (-1) 0
right = Vec 1 0
up = Vec 0 (-1)
down = Vec 0 1

add :: Vec -> Vec -> Vec
add a b = Vec (x a + x b) (y a + y b)

index :: Vec -> (Int, Int)
index p = (x p, y p)

-- Pose is a tuple of position and orientation.
data Pose = Pose { position::Vec, orientation::Vec } deriving (Show, Eq)

advance :: Pose -> Pose
advance p = p {position = add (position p) (orientation p)}

-- Move is a pose combined with a cost.
data Move = Move { pose::Pose, cost::Int } deriving (Show, Eq)

instance Ord Move where
    (Move _ a) `compare` (Move _ b) = a `compare` b

-- We just use Char for each field.
empty = '.'

-- Board is an immutable array.
type Board = Array (Int, Int) Char

showBoard :: Board -> String
showBoard arr =
  unlines $ map (map (arr !)) indices
  where indices = [[(x, y) | x <- [sX..eX]] | y <- [sY..eY]]
        ((sX, sY), (eX, eY)) = bounds arr

mkBoard :: (Int, Int) -> Array (Int, Int) Char
mkBoard (d, e) = array ((0, 0), (m, n)) [((x, y), empty) | x<-[0..m], y<-[0..n]]
  where m = d-1
        n = e-1

charAt :: Board -> Vec -> Maybe Char
charAt b p =
    if x p >= 0 && y p >= 0 && x p <= eX && y p <= eY then
        Just (b ! index p)
    else
        Nothing
  where ((sX, sY), (eX, eY)) = bounds b

allPoses :: Board -> [Pose]
allPoses b =
    [Pose (Vec x y) o | x <- [sX..eX], y <- [sY..eY], o <- [right, down]]
  where
    ((sX, sY), (eX, eY)) = bounds b

-- Test whether it's possible to insert a string at the given pose.
-- Return the associated cost or Nothing, if the string didn't fit.
computeCost :: String -> Pose -> Board -> Maybe Int
computeCost [] _ _ = Just 0
computeCost (w:ws) p b =
  case charAt b (position p) of
    Nothing -> Nothing
    Just c ->
        if c /= empty && c /= w then
            Nothing
        else do
          t <- computeCost ws (advance p) b
          return ((if c == w then 0 else 1) + t)

tryInsert :: String -> Pose -> Board -> Maybe Move
tryInsert w p b = do
  cost <- computeCost w p b
  return $ Move p cost

-- Given a string and a pose, return a list of updates for a board to
-- insert the string.
mkPatch :: String -> Pose -> [((Int, Int), Char)]
mkPatch [] _ = []
mkPatch (w:ws) p = (index $ position p, w) : mkPatch ws (advance p)

-- Actually insert the string returning a new board.
insert :: String -> Pose -> Board -> Maybe Board
insert [] _ b = Just b
insert w p b =
  if isNothing (computeCost w p b) then
    Nothing
  else
    Just (b // mkPatch w p)

-- Insert all given words into the board.
insertAll :: RandomGen gen => gen -> [String] -> Board -> Maybe Board
insertAll _ [] b = Just b
insertAll rnd (w:ws) b = do
  when (null moves) Nothing
  b' <- insert w (pose $ sort moves !! nth) b
  insertAll rnd' ws b'
  where
    (nth, rnd') = randomR (0, min 10 (length moves)) rnd
    moves = mapMaybe (\x -> tryInsert w x b) shuffled
    shuffled = shuffle' all (length all) rnd
    all = allPoses b

-- Fill all empty fields with a random character.
fill :: RandomGen gen => gen -> Board -> Board
fill rnd b = b // zip [(x, y) | x <- [sX..eX], y <- [sY..eY], (b ! (x, y)) == empty] (randomRs ('A', 'Z') rnd)
     where
       ((sX, sY), (eX, eY)) = bounds b

data Puzzle = Puzzle { blanks :: Bool
                     , dimension :: (Int, Int)
                     } deriving (Data, Typeable, Show)

main :: IO ()
main = do
  args <- cmdArgs $ Puzzle False (30, 30)
  rnd <- getStdGen
  c <- getContents
  let guests = lines (map toUpper c) in
    case (showBoard . (if blanks args then id else fill rnd)) <$> insertAll rnd guests (mkBoard $ dimension args) of
      Just s -> putStr s
      Nothing -> putStr "goesnt"