module Test.RedBlackTreeSpec where

import Control.Applicative
import Control.Monad
import Data.List
import Data.RedBlackTree
import System.Random
import Test.Hspec
import Test.Hspec.Core.QuickCheck (modifyMaxSuccess)
import Test.QuickCheck

redBlackTreesWithRightChild ::
     (Bounded a, Ord a, Random a) => Gen (RedBlackTree a)
redBlackTreesWithRightChild =
  (mkRedBlackTree <$> (listOf arbitraryBoundedRandom)) `suchThat`
  ((mkEmpty /=) . rightChild)

redBlackTreesWithLeftChild ::
     (Bounded a, Ord a, Random a) => Gen (RedBlackTree a)
redBlackTreesWithLeftChild =
  (mkRedBlackTree <$> (listOf arbitraryBoundedRandom)) `suchThat`
  ((mkEmpty /=) . leftChild)

redBlackTreesWithDistinctKeys ::
     (Bounded a, Ord a, Random a) => Gen (RedBlackTree a)
redBlackTreesWithDistinctKeys = do
  xs <-
    (listOf arbitraryBoundedRandom) `suchThat`
    (\ys -> (length $ nub ys) == length ys)
  return $ mkRedBlackTree xs

redBlackTreeInsertions ::
     (Bounded a, Ord a, Random a) => Gen ((RedBlackTree a), a)
redBlackTreeInsertions = do
  xs <- listOf arbitraryBoundedRandom
  x <- arbitraryBoundedRandom
  return $ ((Data.RedBlackTree.insert (mkRedBlackTree xs) x), x)

spec :: Spec
spec =
  modifyMaxSuccess (const 10000) $
  describe "Binary Search Trees" $ do
    context "when balanced externally" $ do
      context "when constructing trees" $ do
        it "preserves the BST ordering property" $
          property $ \xs ->
            inorderTraversal (mkRedBlackTree (xs :: [Int])) `shouldBe` (sort xs)
      context "when traversing" $ do
        it "can be traversed in order" $ do
          let bst = mkRedBlackTree [2, 1, 3]
          inorderTraversal bst `shouldBe` [1, 2, 3]
      context "when finding the smallest element" $ do
        it "returns Nothing given the empty tree" $ do
          Data.RedBlackTree.minimum (mkEmpty :: RedBlackTree Int) `shouldBe`
            Nothing
        it "can find the smallest element of non-empty trees" $ do
          let bst = mkRedBlackTree [2, 1, 3]
          let expectedMinimum = search bst 1
          Data.RedBlackTree.minimum bst `shouldBe` expectedMinimum
      context "when finding the successors of nodes" $ do
        it "can yield the successor of a node in the empty tree" $ do
          successor (mkEmpty :: RedBlackTree Int) 42 `shouldBe` Nothing
        it "can yield the successor of a node in a non-empty tree" $ do
          let bst = mkRedBlackTree [2, 1, 3]
          let expectedSuccessor = search bst 3
          successor bst 2 `shouldBe` expectedSuccessor
        it
          "can yield the successor of a node in a non-empty tree when the provided node has no right subtree" $ do
          let bst = mkRedBlackTree [2, 1, 3]
          let expectedSuccessor = search bst 2
          successor bst 1 `shouldBe` expectedSuccessor
        it "always yields the correct successor" $
          forAll (redBlackTreeInsertions :: Gen ((RedBlackTree Int), Int)) $ \(rbt, newKey) ->
            (keyOf =<< successor rbt newKey) `shouldBe`
            (find (> newKey) $ inorderTraversal rbt)
      context "when searching" $ do
        it "returns nothing if the tree is empty" $ do
          let emptyBST = mkEmpty :: RedBlackTree Int
          search emptyBST 1 `shouldBe` Nothing
        it "retrieves keys that are present in the singleton tree" $ do
          let bst = mkRoot 1
          search bst 1 `shouldBe` (Just bst)
        it "returns nothing if the key is not present in the tree" $ do
          let bst = mkRoot 1
          search bst 99 `shouldBe` Nothing
        it "can recurse down the left hand path of the tree" $ do
          let bst = mkRedBlackTree [2, 1] :: RedBlackTree Int
          let searchResult = search bst 1
          (keyOf =<< searchResult) `shouldBe` (Just 1)
        it "can recurse down the right hand path of the tree" $ do
          let bst = mkRedBlackTree [2, 3]
          let searchResult = search bst 3
          (keyOf =<< searchResult) `shouldBe` (Just 3)
      context "when inserting new elements" $ do
        it "turns the empty tree into a singleton tree" $ do
          Data.RedBlackTree.insert mkEmpty 1 `shouldBe` mkRoot 1
        it "makes keys less than the key of a singleton tree the left child" $ do
          let bst = mkRoot 1 :: RedBlackTree Int
          let bst' = Data.RedBlackTree.insert bst 0
          let newNode = search bst' 0
          (keyOf . leftChild $ bst') `shouldBe` (Just 0)
          rightChild bst' `shouldBe` mkEmpty
        it "makes keys greater than the key of a singleton tree the right child" $ do
          let bst = mkRoot 1
          let bst' = Data.RedBlackTree.insert bst 2
          let newNode = search bst' 2
          leftChild bst' `shouldBe` mkEmpty
          (keyOf . rightChild $ bst') `shouldBe` (Just 2)
        it "inserts keys less than the key of a node into the left subtree" $ do
          let bst = mkRedBlackTree [2, 1, 3]
          let bst' = Data.RedBlackTree.insert bst 0
          let left = leftChild bst'
          let leftLeft = leftChild . leftChild $ bst'
          let right = rightChild $ bst'
          (keyOf left) `shouldBe` (Just 1)
          (keyOf leftLeft) `shouldBe` (Just 0)
          (keyOf right) `shouldBe` (Just 3)
        it "inserts keys greater than the key of a node into the right subtree" $ do
          let bst = mkRedBlackTree [2, 1, 3]
          let bst' = Data.RedBlackTree.insert bst 4
          let left = leftChild bst'
          let right = rightChild $ bst'
          let rightRight = rightChild . rightChild $ bst'
          (keyOf left) `shouldBe` (Just 1)
          (keyOf right) `shouldBe` (Just 3)
          (keyOf rightRight) `shouldBe` (Just 4)
    context "when balancing internally" $ do
      context "when rotating left" $ do
        it "preserves the BST ordering property" $
          forAll (redBlackTreesWithRightChild :: Gen (RedBlackTree Int)) $ \rbt ->
            inorderTraversal (rotateLeft rbt) `shouldBe` (inorderTraversal rbt)
        it "is the inverse of rotateRight" $
          forAll (redBlackTreesWithRightChild :: Gen (RedBlackTree Int)) $ \rbt -> do
            (inorderTraversal (rotateRight . rotateLeft $ rbt)) `shouldBe`
              (inorderTraversal rbt)
            (preorderTraversal (rotateRight . rotateLeft $ rbt)) `shouldBe`
              (preorderTraversal rbt)
        it "makes the right child of a tree the new root" $ do
          let rbt = mkRedBlackTree [100, 25, 175, 0, 150, 50, 200]
          let rbt' = rotateLeft rbt
          (keyOf rbt') `shouldBe` (Just 175)
        it "makes the old root the new root's left subtree" $ do
          let rbt = mkRedBlackTree [100, 25, 175, 0, 150, 50, 200]
          let rbt' = rotateLeft rbt
          (keyOf $ leftChild rbt') `shouldBe` (Just 100)
        it "makes the right child's left subtree the old root's right subtree" $ do
          let rbt = mkRedBlackTree [100, 25, 175, 0, 150, 50, 200]
          let rbt' = rotateLeft rbt
          (keyOf . rightChild . leftChild $ rbt') `shouldBe` (Just 150)
        it "makes the new root's right subtree the right child's right subtree" $ do
          let rbt = mkRedBlackTree [100, 25, 175, 0, 150, 50, 200]
          let rbt' = rotateLeft rbt
          (keyOf . rightChild $ rbt') `shouldBe` (Just 200)
      context "when rotating right" $ do
        it "preserves the BST ordering property" $
          forAll (redBlackTreesWithLeftChild :: Gen (RedBlackTree Int)) $ \rbt ->
            inorderTraversal (rotateRight rbt) `shouldBe` (inorderTraversal rbt)
        it "is the inverse of rotateLeft" $
          forAll (redBlackTreesWithLeftChild :: Gen (RedBlackTree Int)) $ \rbt -> do
            (inorderTraversal (rotateLeft . rotateRight $ rbt)) `shouldBe`
              (inorderTraversal rbt)
            (preorderTraversal (rotateLeft . rotateRight $ rbt)) `shouldBe`
              (preorderTraversal rbt)
        it "makes the left child of a tree the new root" $ do
          let rbt = mkRedBlackTree [100, 25, 175, 0, 150, 50, 200]
          let rbt' = rotateRight rbt
          (keyOf rbt') `shouldBe` (Just 25)
        it "makes the old root the new root's right subtree" $ do
          let rbt = mkRedBlackTree [100, 25, 175, 0, 150, 50, 200]
          let rbt' = rotateRight rbt
          (keyOf $ rightChild rbt') `shouldBe` (Just 100)
        it "makes the left child's right subtree the old root's left subtree" $ do
          let rbt = mkRedBlackTree [100, 25, 175, 0, 150, 50, 200]
          let rbt' = rotateRight rbt
          (keyOf . leftChild . rightChild $ rbt') `shouldBe` (Just 50)
        it "makes the new root's left subtree the left child's left subtree" $ do
          let rbt = mkRedBlackTree [100, 25, 175, 0, 150, 50, 200]
          let rbt' = rotateRight rbt
          (keyOf . leftChild $ rbt') `shouldBe` (Just 0)
      context "when inserting new elements" $ do
        it "keeps the root node black" $
          property $ \xs ->
            colourOf (mkRedBlackTree (xs :: [Int])) `shouldBe` Black
        it "ensures all (sentinel) leaves are black" $ do
          let leavesAreBlack rbt =
                case mkEmpty == rbt of
                  True -> colourOf rbt == Black
                  _ ->
                    leavesAreBlack (leftChild rbt) &&
                    leavesAreBlack (rightChild rbt)
          property $ \xs -> leavesAreBlack (mkRedBlackTree (xs :: [Int]))
        it "ensures that both children of red nodes are black" $ do
          let bothChildrenAreBlack rbt =
                colourOf (leftChild rbt) == Black &&
                colourOf (rightChild rbt) == Black
          let bothChildrenOfRedNodesAreBlack rbt =
                case mkEmpty == rbt of
                  True -> True
                  _ ->
                    ((colourOf rbt == Black) ||
                     (colourOf rbt == Red && bothChildrenAreBlack rbt)) &&
                    bothChildrenOfRedNodesAreBlack (leftChild rbt) &&
                    bothChildrenOfRedNodesAreBlack (rightChild rbt)
          property $ \xs ->
            bothChildrenOfRedNodesAreBlack $ mkRedBlackTree (xs :: [Int])
        it
          "ensures that all paths to leaves have the same number of black nodes" $ do
          let blackHeightLeft rbt n =
                case mkEmpty == rbt of
                  True -> n
                  _ ->
                    case (colourOf $ leftChild rbt) of
                      Black -> blackHeightLeft (leftChild rbt) (n + 1)
                      Red -> blackHeightLeft (leftChild rbt) n
          let blackHeightRight rbt n =
                case mkEmpty == rbt of
                  True -> n
                  _ ->
                    case (colourOf $ rightChild rbt) of
                      Black -> blackHeightRight (rightChild rbt) (n + 1)
                      Red -> blackHeightRight (rightChild rbt) n
          let allPathsHaveSameNumberOfBlackNodes rbt =
                case mkEmpty == rbt of
                  True -> True
                  _ -> blackHeightLeft rbt 0 == blackHeightRight rbt 0
          forAll (redBlackTreesWithDistinctKeys :: Gen (RedBlackTree Int)) $
            allPathsHaveSameNumberOfBlackNodes
        it "colours the new node red" $ do
          let rbt = mkRedBlackTree [1]
          let rbt' = Data.RedBlackTree.insert rbt 0
          let newNode = search rbt' 0
          colourOf <$> newNode `shouldBe` (Just Red)
        it
          "handles the case in which the new node was inserted as a right child of a left child, and its uncle is red" $ do
          let rbt = mkRedBlackTree [50, 0, 100]
          let rbt' = Data.RedBlackTree.insert rbt 25
          let newNode = search rbt' 25
          let newNodeParent = search rbt' 0
          let newNodeGrandparent = search rbt' 50
          let newNodeUncle = search rbt' 100
          colourOf <$> newNode `shouldBe` (Just Red)
          colourOf <$> newNodeParent `shouldBe` (Just Black)
          colourOf <$> newNodeGrandparent `shouldBe` (Just Black)
          colourOf <$> newNodeUncle `shouldBe` (Just Black)
        it
          "handles the case in which the new node was inserted as a left child of a left child, and its uncle is red" $ do
          let rbt = mkRedBlackTree [100, 50, 200]
          let rbt' = Data.RedBlackTree.insert rbt 25
          let newNode = search rbt' 25
          let newNodeParent = search rbt' 50
          let newNodeGrandparent = search rbt' 100
          let newNodeUncle = search rbt' 200
          colourOf <$> newNode `shouldBe` (Just Red)
          colourOf <$> newNodeParent `shouldBe` (Just Black)
          colourOf <$> newNodeGrandparent `shouldBe` (Just Black)
          colourOf <$> newNodeUncle `shouldBe` (Just Black)
        it
          "handles the case in which the new node was inserted as a right child of a right child, and its uncle is red" $ do
          let rbt = mkRedBlackTree [50, 0, 100]
          let rbt' = Data.RedBlackTree.insert rbt 200
          let newNode = search rbt' 200
          let newNodeParent = search rbt' 100
          let newNodeGrandparent = search rbt' 50
          let newNodeUncle = search rbt' 0
          colourOf <$> newNode `shouldBe` (Just Red)
          colourOf <$> newNodeParent `shouldBe` (Just Black)
          colourOf <$> newNodeGrandparent `shouldBe` (Just Black)
          colourOf <$> newNodeUncle `shouldBe` (Just Black)
        it
          "handles the case in which the new node was inserted as a left child of a right child, and its uncle is red" $ do
          let rbt = mkRedBlackTree [100, 50, 200]
          let rbt' = Data.RedBlackTree.insert rbt 150
          let newNode = search rbt' 150
          let newNodeParent = search rbt' 200
          let newNodeGrandparent = search rbt' 100
          let newNodeUncle = search rbt' 50
          colourOf <$> newNode `shouldBe` (Just Red)
          colourOf <$> newNodeParent `shouldBe` (Just Black)
          colourOf <$> newNodeGrandparent `shouldBe` (Just Black)
          colourOf <$> newNodeUncle `shouldBe` (Just Black)
        it
          "handles the case in which the new node was inserted as a left child of a left child, and its uncle is black" $ do
          let rbt = mkRedBlackTree [200, 50]
          let rbt' = Data.RedBlackTree.insert rbt 0
          let newNode = search rbt' 0
          let newNodeParent = search rbt' 50
          let newNodeSibling = search rbt' 200
          colourOf <$> newNode `shouldBe` (Just Red)
          colourOf <$> newNodeParent `shouldBe` (Just Black)
          colourOf <$> newNodeSibling `shouldBe` (Just Red)
        it
          "handles the case in which the new node was inserted as a right child of a left child, and its uncle is black" $ do
          let rbt = mkRedBlackTree [50, 0]
          let rbt' = Data.RedBlackTree.insert rbt 25
          let newNode = search rbt' 25 -- Becomes the root after rebalancing
          let newNodeLeftChild = search rbt' 0
          let newNodeRightChild = search rbt' 50
          colourOf <$> newNode `shouldBe` (Just Black)
          colourOf <$> newNodeLeftChild `shouldBe` (Just Red)
          colourOf <$> newNodeRightChild `shouldBe` (Just Red)
        it
          "handles the case in which the new node was inserted as a left child of a right child, and its uncle is black" $ do
          let rbt = mkRedBlackTree [200, 500]
          let rbt' = Data.RedBlackTree.insert rbt 300
          let newNode = search rbt' 300 -- Becomes the root after rebalancing
          let newNodeLeftChild = search rbt' 200
          let newNodeRightChild = search rbt' 500
          colourOf <$> newNode `shouldBe` (Just Black)
          colourOf <$> newNodeLeftChild `shouldBe` (Just Red)
          colourOf <$> newNodeRightChild `shouldBe` (Just Red)
        it
          "handles the case in which the new node was inserted as a right child of a right child, and its uncle is black" $ do
          let rbt = mkRedBlackTree [0, 50]
          let rbt' = Data.RedBlackTree.insert rbt 100
          let newNode = search rbt' 100
          let newNodeParent = search rbt' 50
          let newNodeSibling = search rbt' 0
          colourOf <$> newNode `shouldBe` (Just Red)
          colourOf <$> newNodeParent `shouldBe` (Just Black)
          colourOf <$> newNodeSibling `shouldBe` (Just Red)