import Data.Ix
import Data.List
import Data.List.Split hiding (chunk, split)
import Control.Parallel.Strategies(using, parList, rdeepseq) --rseq, Eval, runEval, rpar
import System.Environment (getArgs)

main :: IO ()
main = do
    args <- getArgs
    if length args /= 2
    then putStrLn "Usage: ./nQueens <option: 1 - show the number of solutions; 2 - print all the solutions> <length of the board>"
    else do
      let (o:s:_) = args
      let size = rInt s
      let option = rInt o
      if option == 1
        then nQueensSolutionNum size
        else do 
          printNQueensSolutions size

-- nQueensNumSoln calculates the number of solutions for a given size 
nQueensSolutionNum :: Int -> IO ()
nQueensSolutionNum n 
  | n > 3 = print $ length $ genSolutions n
  | otherwise = print (0 :: Int)

-- printNQueensSolutions prints all the solutions in board for a given size 
printNQueensSolutions :: Int -> IO ()
printNQueensSolutions n
  | n > 3 = mapM_ (printBoard n) $ genSolutions n
  | otherwise = putStrLn "No solutions"

-- genSolutions generates all the valid solutions for a given size
genSolutions :: Int -> [[(Int, Int)]]
genSolutions n = removeDiagsInChunks $ getCoords $ getAllPermutations n

-- removeDiagsInChunks removes invalid solutions violating diagnoal rule and returns all valid solutions
removeDiagsInChunks :: [[(Int, Int)]] -> [[(Int, Int)]]
removeDiagsInChunks xs = foldr1 merge $ (map (filter (not . invalidSolution)) (split 32 xs) `using` parList rdeepseq)

-- getAllPermutations returns a list of all permutations given an int
-- the output is every possible permutation of queen's position on board
getAllPermutations :: Int -> [[Int]]
getAllPermutations 0 = [[]]
getAllPermutations n = permutations [1..n]

-- getCoords converts position of queens to coordinates
getCoords :: [[Int]] -> [[(Int, Int)]]
getCoords xs = map (zip [1..]) xs

-- invalidSolution checks if a solution has any two queens on the same diagonal
invalidSolution :: [(Int, Int)] -> Bool
invalidSolution [] = False
invalidSolution (x:xs) = (True `elem` map (checkDiagonal x) xs) || invalidSolution xs

-- helper function to merge two list of solutions 
merge :: [[(Int, Int)]] -> [[(Int, Int)]] -> [[(Int, Int)]]
merge x y = x ++ y

-- checkDiagonal returns true if two coordinates are on the same diagonal 
checkDiagonal :: (Int, Int) -> (Int, Int) -> Bool
checkDiagonal (x1, y1) (x2, y2) = abs(x1 - x2) == abs(y1 - y2)

split :: Int -> [[(Int, Int)]] -> [[[(Int, Int)]]]
split numChunks xs = chunk (length xs `quot` numChunks) xs

chunk :: Int -> [[(Int, Int)]] -> [[[(Int, Int)]]]
chunk _ [] = []
chunk n xs = let (as,bs) = splitAt n xs in as : chunk n bs

-- helper function that convert string to integer
rInt :: String -> Int
rInt = read

-- printBoard prints a solution for a given solution and its size
printBoard :: Int -> [(Int, Int)] -> IO ()
printBoard n xs = do
  putStrLn ""
  mapM_ putStrLn $ map unwords $ chunksOf n $ map (printCell xs) $ range ((1,1),(n,n))
  putStrLn ""

-- printCell prints Q to represent queen and X represent empty cell on the board
printCell :: [(Int, Int)] -> (Int, Int) -> String
printCell xs x
  | x `elem` xs = "Q "
  | otherwise = "X "

{-
-- Deprecated:
-- Too many possible solutions and not enough work per solution for per-solution filter parallelism;
-- Overhead was dominated and Sparks were overflowed. Change to split work into larger chunks.

-- removeDiags removes any solutions containing diagonally aligned queens
removeDiags :: [[(Int, Int)]] -> [[(Int,Int)]]
removeDiags xs = runEval $ parFilter (not.invalidSolution) xs

-- parFilter filters invalid solutions and outputs solutions with no queens on the same diagonal 
parFilter :: ([(Int, Int)] -> Bool) -> [[(Int, Int)]] -> Eval [[(Int, Int)]]
parFilter _ [] = return []
parFilter f (a: as) = do
  b <- rpar (f a)
  bs <- parFilter f as
  if b
    then return (a: bs)
  else return bs
-}
