{-# LANGUAGE FlexibleContexts #-}

module Main where

import Control.Lens
import Control.Applicative
import Linear.V2
import Linear.V3
import Data.Matrix
import Data.List
import qualified Data.Vector as V
--import Numeric.LinearAlgebra
import System.Random
import Foreign.Storable
import System.IO
import System.Environment

-- Open | Attacked | Queen
data Square = O | X | Qu deriving (Show, Eq)

initBoard = matrix 8 8 $ const O

main :: IO ()
main = do
  args <- getArgs
  let n = read $ head args :: Int
    in putStrLn $ show $ solve n

solve :: Int -> [Matrix Square]
solve n
  | n < 1 || n > 8 = []
  | otherwise = solveN n
  where
    place :: [Matrix Square] -> [Matrix Square]
    place bs = concat $ map (\b -> map (\p -> placeQueen p b) (openPositions b)) bs
    solveN 0 = []               -- ghc thinks I need this; I'm not so sure
    solveN 1 = place [initBoard]
    solveN n = place (solveN (n-1))


queenAllowed p b = all (== O) $ map snd $ getRowColDiag p b
placeQueen p b = setElem Qu p $ mapQueensPath (\_ _ -> X) p b

allPositions = [(i,j) | i <- [1..8], j <- [1..8]]
scanBoard s b = [(i,j) | (i,j) <- allPositions, (getElem i j b) == s]

openPositions b = scanBoard O b
mapQueensPath :: ((Int,Int) -> a -> a) -> (Int,Int) -> Matrix a -> Matrix a
mapQueensPath f p b = mapRowColDiag f p b

markX p b = setElem X p b

diag :: (Int,Int) -> Matrix a -> [(Int,Int)]
diag c m = apply c
  where
    ops = (\a b -> [(a,a),(a,b),(b,a),(b,b)]) (+1) (subtract 1)
    applyopR (x,y) (fa,fb) = traverse' (\(i,j) -> (fa i, fb j)) (x,y)
    traverse' dir p = takeWhile inBounds $ iterate dir p
    inBounds (i,j) = i >= 1 && i <= nrows m && j >= 1 && j <= ncols m
    apply x = concat $ map (applyopR x) ops

-- did this 2 ways; not sure why I like above better; mostly cuz
diag' r c = let rl = [1..r]
                rr = [r..8]
                cu = [1..c]
                cd = [c..8]
             in zip rl cu ++ zip rr cd ++
                zip (reverse rl) cd ++ zip (reverse rr) cu


-- some of this next shit would've been nice to have already been Data.Matrix 
mapRowColDiag :: ((Int,Int) -> a -> a) -> (Int,Int) -> Matrix a -> Matrix a
mapRowColDiag f (r,c) m =
  matrix (nrows m) (ncols m) $ \(i,j) ->
    let a = unsafeGet i j m
    in  if i == r || j == c || elem (i,j) ds
          then f (i,j) a
        else a
  where ds = diag (r,c) m

getRowColDiag (r,c) m = map (\(i,j) -> ((i,j), getElem i j m)) $ p
  where rows = [(r,j) | j <- [1..(ncols m)]]
        cols = [(i,c) | i <- [1..(nrows m)]]
        ds = diag (r,c) m
        p = nub $ rows ++ cols ++ ds

getRowL r b = map (\x -> b ! (r,x)) [1..8]
getColL c b = map (\x -> b ! (x,c)) [1..8]
getDiagL p b = map (\x -> b ! x) $ diag p b

mapDiag :: ((Int, Int) -> a -> a) -> (Int,Int) -> Matrix a -> Matrix a
mapDiag f p m =
  let diags = diag p m
   in matrix (nrows m) (ncols m) $ \(i, j) ->
        let a = unsafeGet i j m
         in if elem (i,j) diags
              then f (i, j) a
              else a