{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# language NoImplicitPrelude #-}
module MyQueens where
-- To run, type `stack ghci` then `:load MyQueens` then `main`

import Rebase.Prelude

type Rank = Int8
type MaxRank = Rank

minRank :: Rank
minRank = 1

data RankOptions = RankOptions [Rank] | RankChoice Rank
  deriving (Show)

type Board = [RankOptions]
type PartialBoard = [RankOptions]

select :: MaxRank -> Board -> [Board]
select _ []                    = error "select on invalid board"
select maxRank (x@(RankChoice _):xs) = map (x:) (select maxRank xs)
select maxRank ((RankOptions rs):xs) = dropHopeless $ map (restrict xs) rs

  where

    dropHopeless :: [Board] -> [Board]
    dropHopeless = filter (not . any hopeless)

    hopeless :: RankOptions -> Bool
    hopeless (RankOptions []) = True
    hopeless _                = False

    filterRankOptions :: (Rank -> Bool) -> RankOptions -> RankOptions
    filterRankOptions f (RankOptions rs) = RankOptions (filter f rs)
    filterRankOptions _ x                = x

    restrict :: PartialBoard -> Rank -> PartialBoard
    restrict rs r = ((RankChoice r) :) $
      (dropDiagonal pred minRank r) .
      (dropDiagonal succ maxRank r) .
      (map (filterRankOptions (/= r))) $
      rs

    dropDiagonal :: (Rank -> Rank) -> Rank -> Rank -> PartialBoard -> PartialBoard
    dropDiagonal _ _ _ []                 = []
    dropDiagonal _ bound r x | r == bound = x
    dropDiagonal next bound r (x:xs)      = (filterRankOptions (/= next r) x) : dropDiagonal next bound (next r) xs

queens :: MaxRank -> [Board]
queens maxRank = foldl1 (.) (take (fromIntegral maxRank) $ repeat select') $ [allOptionsBoard]

  where

    select' :: [Board] -> [Board]
    select' = concatMap (select maxRank)

    allOptionsColumn :: RankOptions
    allOptionsColumn = RankOptions [minRank .. maxRank]

    allOptionsBoard :: Board
    allOptionsBoard = take (fromIntegral maxRank) $ repeat allOptionsColumn

eightQueens :: IO ()
eightQueens = do
  forM_ qs print
  print $ length qs
  where
    qs = queens 8

main :: IO ()
main = do
  eightQueens
  -- finding one solution is quick up to 21.
  forM_ [4..21] $
    \n -> print (n, head (queens n))
  -- finding all solutions is only quick up to 12.
  forM_ [4..12] $
    \n -> print (n, length (queens n))