-- Copyright 2022 University of Freiburg
-- Janek Spaderna <janek.spaderna@pluto.uni-freiburg.de>
{-# LANGUAGE QualifiedDo #-}

module Solutions.NQueens where

import qualified Data.List as List
import SimplePrelude as SP
import Prelude ()

-------------------------------------------------------------------------------
-- `Monad []` instance.

{-
instance Monad [] where
  return a = [a]
  xs >>= f = concatMap f xs
-}

-------------------------------------------------------------------------------
-- Search algorithm

-- | A list of columns. The integer in each column specifies at which 0-based
-- index the queen is positioned.
type Solution = [Integer]

nqueens :: Integer -> [Solution]
nqueens n = nqueens' [] n n

nqueens' :: [Integer] -> Integer -> Integer -> [Solution]
nqueens' placed _rows _cols@0 = return placed
nqueens' placed rows cols = SP.do
  -- Draw (or "guess") the next row in which we try to position the next queen.
  r <- [0 .. rows - 1]
  -- Check that there are no other queens in this row.
  guard $ r `notElem` placed
  -- Check that there are no queens diagonally above.
  guard $ r `notElem` zipWith (+) placed [1 ..]
  -- Check that there are no queens diagonally below.
  guard $ r `notElem` zipWith (-) placed [1 ..]
  -- No conflicts for `r`. Solve the sub-problem with one column less.
  nqueens' (r : placed) rows (cols - 1)

guard :: Bool -> [()]
guard True = return ()
guard False = []

-------------------------------------------------------------------------------
-- Solving & printing the solution

-- | Prints the first solution.
solve :: Integer -> IO ()
solve n = case nqueens n of
  [] -> putStrLn "No solution found."
  solution : _ -> putStr $ displaySolution n n solution

-- | Prints all solutions.
solveAll :: Integer -> IO ()
solveAll n = case nqueens n of
  [] -> putStrLn "No solution found."
  solutions -> mapM_ (putStrLn . displaySolution n n) solutions

-------------------------------------------------------------------------------
-- Visualizing

-- | Prettyprints a solution as a grid with the specified size.
displaySolution :: Integer -> Integer -> Solution -> String
displaySolution rows cols solution =
  unlines
    . List.transpose
    . List.genericTake cols
    $ map displayColumn solution ++ repeat emptyColumn
  where
    displayColumn n =
      List.genericTake rows $
        List.genericReplicate n '⋅' ++ '♛' : List.genericReplicate (rows - n - 1) '⋅'
    emptyColumn =
      List.genericReplicate rows '⋅'