-- # Software Tools for Discrete Mathematics
module Stdm05Trees where

import Test.QuickCheck

-- # Chapter 5.  Trees

data Tree a = Leaf | Node a (Tree a) (Tree a)
  deriving (Show, Eq)

-- |
-- >>> t1
-- Node 6 Leaf Leaf
t1 :: Tree Int
t1 = Node 6 Leaf Leaf

-- |
-- >>> t2
-- Node 5 (Node 3 Leaf Leaf) (Node 8 (Node 6 Leaf Leaf) (Node 12 Leaf Leaf))
t2 :: Tree Int
t2 =
  Node
    5
    (Node 3 Leaf Leaf)
    ( Node
        8
        (Node 6 Leaf Leaf)
        (Node 12 Leaf Leaf)
    )

-- |
-- >>> nodeCount t2
-- 5
nodeCount :: Tree a -> Int
nodeCount Leaf = 0
nodeCount (Node x t1 t2) = 1 + nodeCount t1 + nodeCount t2

-- |
-- >>> height t2
-- 3
height :: Tree a -> Int
height Leaf = 0
height (Node x t1 t2) = 1 + max (height t1) (height t2)

-- |
-- >>> reflect t2
-- Node 5 (Node 8 (Node 12 Leaf Leaf) (Node 6 Leaf Leaf)) (Node 3 Leaf Leaf)
reflect :: Tree a -> Tree a
reflect Leaf = Leaf
reflect (Node a t1 t2) = Node a (reflect t2) (reflect t1)

-- |
-- >>> mapTree (5 *) t2
-- Node 25 (Node 15 Leaf Leaf) (Node 40 (Node 30 Leaf Leaf) (Node 60 Leaf Leaf))
mapTree :: (a -> b) -> Tree a -> Tree b
mapTree f Leaf = Leaf
mapTree f (Node a t1 t2) = Node (f a) (mapTree f t1) (mapTree f t2)

-- | This is like a dictionary of (fst, snd)
-- >>> t3
-- Node (5,10) (Node (3,6) (Node (1,1) Leaf Leaf) (Node (4,8) Leaf Leaf)) (Node (7,14) (Node (6,12) Leaf Leaf) (Node (8,16) Leaf Leaf))
t3 :: Tree (Int, Int)
t3 =
  Node
    (5, 10)
    ( Node
        (3, 6)
        (Node (1, 1) Leaf Leaf)
        (Node (4, 8) Leaf Leaf)
    )
    ( Node
        (7, 14)
        (Node (6, 12) Leaf Leaf)
        (Node (8, 16) Leaf Leaf)
    )

-- |
-- >>> balanced t2
-- False

-- |
-- >>> balanced t3
-- True
balanced :: Tree a -> Bool
balanced Leaf = True
balanced (Node x t1 t2) = balanced t1 && balanced t2 && (height t1 == height t2)

-- |
-- >>> inorder t2
-- [3,5,6,8,12]
inorder :: Tree a -> [a]
inorder Leaf = []
inorder (Node x t1 t2) = inorder t1 ++ [x] ++ inorder t2

-- |
-- >>> preorder t2
-- [5,3,8,6,12]
preorder :: Tree a -> [a]
preorder Leaf = []
preorder (Node x t1 t2) = [x] ++ preorder t1 ++ preorder t2

-- |
-- >>> postorder t2
-- [3,6,12,8,5]
postorder :: Tree a -> [a]
postorder Leaf = []
postorder (Node x t1 t2) = postorder t1 ++ postorder t2 ++ [x]

-- | This prints nicely; show it!
-- >>> putStrLn (inorderBackwards t2 "")
-- "\n\n    \n            Leaf\n        12\n            Leaf\n    8\n    \n            Leaf\n        6\n            Leaf\n5\n\n        Leaf\n    3\n        Leaf"
inorderBackwards :: (Show a) => Tree a -> String -> String
inorderBackwards Leaf indent = "    " ++ "Leaf"
inorderBackwards (Node x t1 t2) indent =
  ( "\n"
      ++ indent
      ++ inorderBackwards t2 (indent ++ "    ")
      ++ "\n"
      ++ indent
      ++ show x
      ++ "\n"
      ++ indent
      ++ inorderBackwards t1 (indent ++ "    ")
  )

-- Expression trees

data Exp
  = Const Int
  | Add Exp Exp
  | Mult Exp Exp
  deriving (Show, Eq)

-- |
-- >>> t4
-- Mult (Const 2) (Add (Const 5) (Const 3))
t4 :: Exp
t4 = Mult (Const 2) (Add (Const 5) (Const 3))

-- | (*) 2 ((+) 5 3)
-- >>> eval t4
-- 16
eval :: Exp -> Int
eval (Const n) = n
eval (Add e1 e2) = eval e1 + eval e2
eval (Mult e1 e2) = eval e1 * eval e2

-- Binary Search Trees are ordered left to right, and must remain ordered!
-- They are more efficient than search through lists, like this:

-- | This searches through a list of (fst, snd) for fst, and returns snd.
-- >>> linSearch 11 (zip [1..] ['a'..'z'])
-- Just 'k'
linSearch :: (Eq a) => a -> [(a, b)] -> Maybe b
linSearch k [] = Nothing
linSearch k ((x, y) : xs) =
  if k == x
    then Just y
    else linSearch k xs

-- Tree t3 tree is ordered. See:
-- putStrLn (inorderBackwards t3 "")

-- | This searches through a tree of (fst, snd) for fst, and returns snd.
-- >>> bstSearch 7 t3
-- Just 14
bstSearch :: (Ord a) => a -> Tree (a, b) -> Maybe b
bstSearch key Leaf = Nothing
bstSearch key (Node (x, y) t1 t2) =
  if key == x
    then Just y
    else
      if key < x
        then bstSearch key t1
        else bstSearch key t2

-- | This function preserves the ordering of the BST!
-- >>> insertTree (2,5) t3
-- Node (5,10) (Node (3,6) (Node (1,1) Leaf (Node (2,5) Leaf Leaf)) (Node (4,8) Leaf Leaf)) (Node (7,14) (Node (6,12) Leaf Leaf) (Node (8,16) Leaf Leaf))
insertTree :: (Ord a) => (a, b) -> Tree (a, b) -> Tree (a, b)
insertTree (key, d) Leaf = Node (key, d) Leaf Leaf
insertTree (key, d) (Node (x, y) t1 t2) =
  if key == x
    then Node (key, d) t1 t2
    else
      if key < x
        then Node (x, y) (insertTree (key, d) t1) t2
        else Node (x, y) t1 (insertTree (key, d) t2)

-- Tree t3 tree is still ordered. Show this in ghci in class:
-- putStrLn (inorderBackwards (insertTree (2,5) t3) "")

-- Induction on trees

-- Theorem 29. Let t :: BinTree a.
-- Then reflect (reflect t) = t.
-- Proof. The proposition P (t) that we need to prove says reflect (reflect t) = t.
-- The theorem is proved by induction over t.

-- | The base case
-- >>> reflectBase Leaf
-- True
reflectBase :: (Eq a) => Tree a -> Bool
reflectBase t = (reflect . reflect) t == t

-- For the inductive case, let t1, t2 :: BinTree a be trees, and assume P (t1) and P (t2).
-- These are the inductive hypotheses.
-- The aim is to prove that the proposition holds for a larger tree P (Node x t1 t2).

-- | Inductive case + 1
-- >>> quickCheck reflectInductive
-- +++ OK, passed 100 tests.
reflectInductive :: (Eq a) => a -> Bool
reflectInductive val =
  let t = Node val Leaf Leaf
   in (reflect . reflect) t == t

-- Theorem 30. Let t :: BinTree a be an arbitrary binary tree of finite size.
-- Then inorder (reflect t) = reverse (inorder t).

-- | Base case
-- >>> reverseBase Leaf
-- True
reverseBase :: (Eq a) => Tree a -> Bool
reverseBase t = inorder (reflect t) == reverse (inorder t)

-- | Inductive case + 1
-- >>> quickCheck reverseInductive
-- +++ OK, passed 100 tests.
reverseInductive :: (Eq a) => a -> Bool
reverseInductive val =
  let t = Node val Leaf Leaf
   in inorder (reflect t) == reverse (inorder t)

-- Theorem 31. Let h = height t. If balanced t, then size t = 2h − 1.
-- Proof. The proposition we want to prove is: balanced t → size t = 2h − 1.
-- The proof is an induction over the tree structure.
-- For the base case, we need to prove that the theorem holds for a leaf.

-- | Base case
-- >>> sizeHeightBase Leaf
-- True
sizeHeightBase :: (Eq a) => Tree a -> Bool
sizeHeightBase t =
  if balanced t
    then nodeCount t == (2 ^ (height t)) - 1
    else error "Starting assumption not met"

-- | Inductive case + 1
-- >>> quickCheck sizeHeightInductive
-- +++ OK, passed 100 tests.
sizeHeightInductive :: (Eq a) => a -> Bool
sizeHeightInductive val =
  let t = Node val Leaf Leaf
   in if balanced t
        then nodeCount t == (2 ^ (height t)) - 1
        else error "Starting assumption not met"

-- Theorem 32. Let t :: BinTree a be any finite binary tree. Then length (inorder t) = size t.

-- | Base case
-- >>> flatLengthBase Leaf
-- True
flatLengthBase :: (Eq a) => Tree a -> Bool
flatLengthBase t = length (inorder t) == nodeCount t

-- | Inductive case + 1
-- >>> quickCheck flatLengthInductive
-- +++ OK, passed 100 tests.
flatLengthInductive :: (Eq a) => a -> Bool
flatLengthInductive val =
  let t = Node val Leaf Leaf
   in length (inorder t) == nodeCount t