{-
 Compile with:

 $ stack ghc -- --make -Wall -O generate_grid.hs

-}

module GenerateGrid where

import Data.List (sortBy)
import Data.Map.Strict (Map, (!))
import qualified Data.Map.Strict as Map
import Data.Set (Set)
import qualified Data.Set as Set
import System.Random
import SquareDefs

type Index = (Int, Int)
type Constraint = (Int, Int, Int, Int)
type Constraints = Map Index Constraint

{-
test_seed :: StdGen
test_seed = mkStdGen 42
-}

probs :: RandomGen g => Int -> g -> [Float]
probs 0 _ = []
probs n s = r : probs (n-1) s'
           where
               (r, s') = randomR (0.0::Float, 1.0::Float) s

indices :: Int -> Int -> Set Index
indices h w = Set.fromList $ [(i,j) | i <- [0..(h - 1)], j <- [0..(w - 1)]]

initConstraints :: Int -> Int -> Constraints
initConstraints h w = Map.fromSet (\_ -> (0, w - 1, 0, h - 1)) (indices h w)

updateConst :: Index -> Index -> Constraint -> Constraint
updateConst (bi, bj) (i, j) (l,r,t,b) = (nl, nr, nt, nb)
                            where
                                (nl, nr) | not sameRow = (l, r)
                                         | left = (max l bj + 1, r)
                                         | right = (l, min r bj - 1)
                                         | otherwise = (i, j)
                                (nt, nb) | not sameCol = (t,b)
                                         | above = (max t bi + 1, b)
                                         | below = (t, min b bi - 1)
                                         | otherwise = (i, j)
                                sameRow = bi == i
                                (left, right) = (bj < j, bj > j)
                                sameCol = bj == j
                                (above, below) = (bi < i, bi > i)

adjustWithKeyList :: Ord k => (k -> a -> a) -> [k] -> Map k a -> Map k a
adjustWithKeyList _ [] m = m
adjustWithKeyList f (x:xs) m = Map.adjustWithKey f x $ adjustWithKeyList f xs m

addBS :: Index -> Constraints -> Constraints
addBS (i,j) con = updateRow $ updateCol con
                       where
                         updateCol c = adjustWithKeyList f col c
                         updateRow c = adjustWithKeyList f row c
                         f = updateConst (i,j)
                         col = filter (\(_,b) -> b == j) keyList
                         row = filter (\(a,_) -> a == i) keyList
                         keyList = Map.keys con

checkSquare :: Index -> Constraints -> Int -> Bool
checkSquare (i,j) c m | i - top < m && i - top /= 0 = False
                      | bot - i < m && bot - i /= 0 = False
                      | j - left < m && j - left /= 0 = False
                      | right - j < m && right - j /= 0 = False
                      | otherwise = True
                         where
                          (left, right, top, bot) = c ! (i,j)

spiralIndices :: Int -> Int -> [Index]
spiralIndices 0 _ = []
spiralIndices 1 w = [ (0, i) | i <- [0..(w - 1)]]
spiralIndices h w = (++) (f 0 0) $ adj $ spiralIndices (h - 2) (w - 2)
                        where
                          f i j = (i,j) : l i j
                          l i j | j < w - 1 && i == 0 = f i (j + 1)
                                | j == w - 1 && i < h - 1 = f (i + 1) j
                                | j > 0 && i == h - 1 = f i (j - 1)
                                | j == 0 && i > 1 = f (i - 1) j
                                | otherwise = []
                          adj spiral = map (\(a, b) -> (a + 1, b + 1)) spiral

fillIn :: [Index] -> [Bool] -> Constraints -> Int -> [Square]
fillIn [] _ _ _ = []
fillIn _ [] _ _ = []
fillIn (i:is) (b:bs) c m | not b = white
                         | checkSquare i c m = black
                         | otherwise = white
                             where
                               white = (White '_') : fillIn is bs c m
                               black = Black : fillIn is bs (addBS i c) m

unravel :: Int -> [(Index, Square)] -> [[Square]]
unravel h s = map getRow [0..(h - 1)]
                where
                  getRow i = map snd $ sortBy my_compare $ filter (isRow i) s
                  my_compare ((_,y1),_) ((_,y2),_) = compare y1 y2
                  isRow i ((x, _), _) = x == i

generateGrid :: RandomGen g => Int -> Int -> Float -> g -> Maybe [[Square]]
generateGrid dim m thresh seed | even dim = Just even_grid
                               | otherwise = Nothing
                                where
                                  even_grid = half ++ rev_half
                                  rev_half = reverse $ map reverse half
                                  half = unravel half_h $ zip s $ squares
                                  squares = fillIn s b c m
                                  s = spiralIndices half_h dim
                                  c = initConstraints half_h dim
                                  b = map (\x -> x <= thresh) p
                                  p = probs (half_h * dim) seed
                                  half_h = div dim 2

printGrid :: [[Square]] -> IO ()
printGrid grid = do mapM_ (putStrLn) $ map show grid
                    putStrLn ""

{-
main :: IO ()
main = do args <- getArgs
          case args of
              [dimS, mS, threshS, dictPath, deterministic] -> do
                  seed <- case deterministic of
                                "y" -> return $ mkStdGen 15
                                _   -> newStdGen
                  let dim = read dimS :: Int
                  let m = read mS :: Int
                  let thresh = read threshS :: Float
                  let result = generateGrid dim m thresh seed
                  let peeled = maybe [[]] (\x -> x) $ result
                  -- mapM_ (putStrLn) $ map show peeled
                  dictStream <- readFile dictPath
                  let dict = build_dict dictStream
                  putStrLn $ show peeled
                  let filled = fill_and_check dict peeled
                  -- mapM_ (putStrLn) $ map show filled
                  mapM_ (printGrid) filled
                  putStrLn $ "found " ++ show (length filled) ++ " ways to fill the given grid"
              _ -> do pn <- getProgName
                      die $ "Usage: "++pn++" <dimension> <min_word_length>\
                                              \ <probability_threshold> <dict_filepath> <deterministic>"
-}
