content: {-# LANGUAGE OverloadedStrings, FlexibleContexts #-}
  module Main ( main ) where
  
  import Control.Monad ( forM_ )
  import Data.List ( sort )
  import qualified Data.Set as S
  import Test.Framework ( defaultMain, testGroup, Test )
  import Test.Framework.Providers.HUnit
  import Test.HUnit hiding ( Test )
  
  import Database.Datalog
  
  main :: IO ()
  main = defaultMain tests
  
  tests :: [Test]
  tests = [ testGroup "t1" [ testCase "4queens" t4
                           , testCase "5queens" t5
                           , testCase "6queens" t6
                           ] ]
  
  type Position = (Int, Int)
  
  -- dbN :: (Failure DatalogError m) => Int -> m (Database Position)
  dbN :: Int -> IO (Database Position)
  dbN n = makeDatabase $ do
    let posTuples = [ (x, y) | x <- [1..n], y <- [1..n] ]
    position <- addRelation "position" 1
    forM_ posTuples $ \(x, y) -> assertFact position [ (x, y) ]
  
  -- Note, the restrictions on x and y equality also imply that the
  -- position rule can't select the same position more than once in a
  -- solution
  posCanAttack :: Position -> Position -> Bool
  posCanAttack (x1, y1) (x2, y2) =
    x1 == x2 || y1 == y2 || (abs (x1 - x2) == abs (y1 - y2))
  
  unique :: [[Position]] -> [[Position]]
  unique = S.toList . S.fromList
  
  -- Return False if any position can attack any others
  noneCanAttack :: [Position] -> Bool
  noneCanAttack [] = True
  noneCanAttack [_] = True
  noneCanAttack (p:ps) = not (any (posCanAttack p) ps) && noneCanAttack ps
  
  t4 :: Assertion
  t4 = do
    db4 <- dbN 4
    res <- queryDatabase db4 q
    let res' = unique $ map sort res
    print res'
    assertBool "t4" $ all noneCanAttack res' && length res' == 2
    where
      q = do
        position <- relationPredicateFromName "position"
        canAttack <- inferencePredicate "canAttack"
        let x = LogicVar "X"
            y = LogicVar "Y"
        (canAttack, [x, y]) |- [ lit position [x]
                               , lit position [y]
                               , cond2 posCanAttack (x, y)
                               ]
        let p1 = LogicVar "P1"
            p2 = LogicVar "P2"
            p3 = LogicVar "P3"
            p4 = LogicVar "P4"
        queens <- inferencePredicate "queens"
        (queens, [p1, p2, p3, p4]) |- [ lit position [p1]
                                      , lit position [p2]
                                      , negLit canAttack [p1, p2]
                                      , lit position [p3]
                                      , negLit canAttack [p1, p3]
                                      , negLit canAttack [p2, p3]
                                      , lit position [p4]
                                      , negLit canAttack [p1, p4]
                                      , negLit canAttack [p2, p4]
                                      , negLit canAttack [p3, p4]
                                      ]
        issueQuery queens [p1, p2, p3, p4]
  
  t5 :: Assertion
  t5 = do
    db5 <- dbN 5
    res <- queryDatabase db5 q
    let res' = unique $ map sort res
    print res'
    assertBool "t5" $ all noneCanAttack res' && length res' == 10
    where
      q = do
        position <- relationPredicateFromName "position"
        canAttack <- inferencePredicate "canAttack"
        let x = LogicVar "X"
            y = LogicVar "Y"
        (canAttack, [x, y]) |- [ lit position [x]
                               , lit position [y]
                               , cond2 posCanAttack (x, y)
                               ]
        let p1 = LogicVar "P1"
            p2 = LogicVar "P2"
            p3 = LogicVar "P3"
            p4 = LogicVar "P4"
            p5 = LogicVar "P5"
        queens <- inferencePredicate "queens"
        (queens, [p1, p2, p3, p4, p5]) |- [ lit position [p1]
                                          , lit position [p2]
                                          , negLit canAttack [p1, p2]
                                          , lit position [p3]
                                          , negLit canAttack [p1, p3]
                                          , negLit canAttack [p2, p3]
                                          , lit position [p4]
                                          , negLit canAttack [p1, p4]
                                          , negLit canAttack [p2, p4]
                                          , negLit canAttack [p3, p4]
                                          , lit position [p5]
                                          , negLit canAttack [p1, p5]
                                          , negLit canAttack [p2, p5]
                                          , negLit canAttack [p3, p5]
                                          , negLit canAttack [p4, p5]
                                          ]
        issueQuery queens [p1, p2, p3, p4, p5]
  
  t6 :: Assertion
  t6 = do
    db6 <- dbN 6
    res <- queryDatabase db6 q
    let res' = unique $ map sort res
    print res'
    assertBool "t6" $ all noneCanAttack res' && length res' == 4
    where
      q = do
        position <- relationPredicateFromName "position"
        canAttack <- inferencePredicate "canAttack"
        let x = LogicVar "X"
            y = LogicVar "Y"
        (canAttack, [x, y]) |- [ lit position [x]
                               , lit position [y]
                               , cond2 posCanAttack (x, y)
                               ]
        let p1 = LogicVar "P1"
            p2 = LogicVar "P2"
            p3 = LogicVar "P3"
            p4 = LogicVar "P4"
            p5 = LogicVar "P5"
            p6 = LogicVar "P6"
        queens <- inferencePredicate "queens"
        (queens, [p1, p2, p3, p4, p5, p6]) |- [ lit position [p1]
                                              , lit position [p2]
                                              , negLit canAttack [p1, p2]
                                              , lit position [p3]
                                              , negLit canAttack [p1, p3]
                                              , negLit canAttack [p2, p3]
                                              , lit position [p4]
                                              , negLit canAttack [p1, p4]
                                              , negLit canAttack [p2, p4]
                                              , negLit canAttack [p3, p4]
                                              , lit position [p5]
                                              , negLit canAttack [p1, p5]
                                              , negLit canAttack [p2, p5]
                                              , negLit canAttack [p3, p5]
                                              , negLit canAttack [p4, p5]
                                              , lit position [p6]
                                              , negLit canAttack [p1, p6]
                                              , negLit canAttack [p2, p6]
                                              , negLit canAttack [p3, p6]
                                              , negLit canAttack [p4, p6]
                                              , negLit canAttack [p5, p6]
                                              ]
        issueQuery queens [p1, p2, p3, p4, p5, p6]
  