main.hs

module Main where

import qualified Data.Set as S
import qualified Data.Map.Strict as M
import Data.Semigroup
import Control.Monad

data Piece
    = Rook
    | Bishop
    | Knight
    | Queen
    | King
    deriving (Show, Eq, Ord)

data Pos
    = Pos
    { _row :: !Int
    , _col :: !Int
    } deriving (Show, Eq, Ord)

type Dimension = Pos

attacks :: Dimension -> Piece -> Pos -> S.Set Pos
attacks d@(Pos mx my) p src@(Pos x y)
  = case p of
      Rook -> S.fromList [ Pos x' y | x' <- [1..mx] ] <> S.fromList [ Pos x y' | y' <- [1..my] ]
      Bishop -> S.fromList $ do
        x' <- [1..mx]
        let dx = x - x'
        y' <- [y+dx, y-dx]
        guard (y' > 0 && y' <= my)
        return (Pos x' y')
      Knight -> S.fromList [ Pos (x + x') (y + y') | x' <- [-2, -1, 1, 2] , y' <- [-2, -1, 1, 2] , abs x' /= abs y' ]
      Queen -> attacks d Bishop src <> attacks d Rook src
      King -> S.fromList [ Pos x' y' | x' <- [x-1..x+1], y' <- [y-1, y+1], x' > 0 && y' > 0 && x' <= mx && y' <= my ]

pieceChar :: Piece -> Char
pieceChar p =
    case p of
     Rook   -> 'R'
     Bishop -> 'B'
     Knight -> 'N'
     Queen  -> 'Q'
     King   -> 'K'

data Board
    = Board
    { _dimension :: Pos
    , _pieces :: M.Map Pos Piece
    } deriving (Eq, Ord)

showBoard :: Board -> String
showBoard (Board (Pos mx my) mp) = unlines $ do
    y <- [1..my]
    let pchar p = case M.lookup p mp of
                    Nothing -> '-'
                    Just pc -> pieceChar pc
    return [ pchar (Pos x y) | x <- [1..mx]]

place :: Dimension -> M.Map Pos Piece -> Maybe Piece -> [Pos] -> [Piece] -> S.Set Board
place dim pieces _ _ [] = S.singleton (Board dim pieces)
place dim@(Pos mx my) pieces lastpiece remaining (p:ps) = go toTraverse
    where
        toTraverse = if lastpiece == Just p
                     then remaining
                     else [ Pos x y | x <- [1..mx], y<- [1..my] ]
        go [] = S.empty
        go (pos : remaining') =
          let pieceSet = M.keysSet pieces
              attackSet = attacks dim p pos
              sol1 = place dim (M.insert pos p pieces) (Just p) remaining' ps
              next = go remaining'
          in if null (S.intersection pieceSet attackSet)
              then sol1 <> next
              else next

makeSolution :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> S.Set Board
makeSolution m n kings queens bishops rooks knights = place (Pos m n) M.empty Nothing [] toplace
  where
    toplace = replicate queens Queen
           ++ replicate rooks Rook
           ++ replicate bishops Bishop
           ++ replicate knights Knight
           ++ replicate kings King

main :: IO ()
main = mapM_ (putStrLn . showBoard) (makeSolution 8 8 0 8 0 0 0)