{-# LANGUAGE InstanceSigs #-}

-- Tree (leaf valued) examples

import qualified Data.Functor
import qualified Data.Foldable

-- define a (polymorphic) data-type of a (leaf-valued) tree
data Tree a  = Leaf a 
             | Branch (Tree a) (Tree a)
	     deriving (Show,Read) -- can use 'show' and 'read' for out-/in-put

-- a function that collects all values from the leaves of the tree
fringe                      :: Tree a -> [a]
fringe (Leaf x)             =  [x]
fringe (Branch left right)  =  fringe left ++ fringe right

-- a simple example: compute the depth of the tree
depth :: Tree a -> Int
depth (Leaf _) = 1
depth (Branch left right) = 1 + max (depth left) (depth right)

-- check whether this tree is balanced
isBalanced :: Tree a -> Bool
isBalanced (Leaf _) = True
isBalanced (Branch left right) = (abs (depth left - depth right)) <= 1 && isBalanced left && isBalanced right
-- Q: is this an efficient version? what's the complexity in terms of tree depth? can you make it more efficient?
-- A: no; for 2^d nodes, depth has a complexity of 2^d, and isBalanced \Sum_{i=0}^{d} 2^i = 2^{i+1}-1

-- a better, but longer solution is this
isBalanced1 :: Tree a -> Bool
isBalanced1 t = abs (pmax-pmin) <= 1
                where pmax = minimum ps
                      pmin = maximum ps
                      ps = allPathLengths t 
                      allPathLengths :: Tree a -> [Int]
                      allPathLengths (Leaf _) = [1]
                      allPathLengths (Branch left right) = map (1+) (allPathLengths left ++ allPathLengths right)

-- a better solution, avoiding intermediate lists, is this
isBalanced2 :: Tree a -> Bool
isBalanced2 t = fst3 (isBalanced'' t)
                where fst3 (x,_,_) = x
                      isBalanced'' :: Tree a -> (Bool, Int, Int)
		      isBalanced'' (Leaf _) = (True, 1, 1)
		      isBalanced'' (Branch left right) = (lb && rb && abs (newmax - newmin) <= 1 , 1+newmin, 1+newmax)
                                                         where (lb, lmin, lmax) = isBalanced'' left
		                                               (rb, rmin, rmax) = isBalanced'' right
                                                               newmin = min lmin rmin
                                                               newmax = max lmax rmax


{- test the performance of both versions like this:

*Main> let t8 = mkTree [(1::Int)..99000]
*Main> :set +s
(*Main> isBalanced t8
True
(1.57 secs, 429012280 bytes)
*Main> isBalanced1 t8
True
(0.54 secs, 268147848 bytes)
*Main> isBalanced2 t8
True
(0.43 secs, 136741408 bytes)
-}

-- isSorted cecks whether the leaf values are sorted
-- version with explicit recursion over a list
isSorted' :: Tree Int -> Bool
isSorted' = isSortedList . fringe
-- the above version avoids naming the argument, which only comes in from the right (better style)
-- the version below is equivalent
-- isSorted' t = isSortedList (fringe t)

-- aux function: a verbose version of checking whether a list is sorted
isSortedList :: [Int] -> Bool
isSortedList [] = True
isSortedList (x:xs) = isSortedList' x xs
                      where isSortedList' x [] = True 
                            isSortedList' x (y:ys) = x<=y && isSortedList' y ys

-- aux fct: a more compact, higher-order version of checking whether a list is sorted
isSorted :: Tree Int -> Bool
isSorted t = and (zipWith (<) xs (tail xs))
             where xs = fringe t

-- overload the == operator on trees
instance (Eq a) => Eq (Tree a) where
 Leaf a         ==  Leaf b          =   a == b
 (Branch l1 r1) == (Branch l2 r2)   =   (l1==l2) && (r1==r2)
 _              ==  _               =   False

-- insert a new value into a tree, whose leaf values are sorted
-- NB: efficiency is poor, due to repeated calculation of fringe
-- Exercise: how can you improve either algorithm or tree rep to avoid this
insert :: Tree Int -> Int -> Tree Int
insert (Leaf x) z | x<z       = Branch (Leaf x) (Leaf z)
                  | otherwise = Branch (Leaf z) (Leaf x)
insert (Branch left right) z | maximum (fringe left) < z = Branch left (insert right z)
                             | otherwise                 = Branch (insert left z) right  

-- pretty print a tree, using indentation to represent depth
pp :: (Show a) => Tree a -> IO ()
pp t = mapM_ putStrLn (pp' 0 t)
       where pp' n (Leaf x) = [(take n (repeat ' ')) ++ (show x)]
             pp' n (Branch left right) = pp' (n+1) left ++ [(take n (repeat ' ')) ++ ['.']] ++ pp' (n+1) right

-- build a balanced tree from a list
mkTree :: [a] -> Tree a
mkTree [] = error "mkTree: empty list"
mkTree [x] = Leaf x
mkTree xs = Branch left right
            where (ls,rs) = splitAt (length xs `div` 2) xs
                  left = mkTree ls
                  right = mkTree rs

-- sample trees
t1, t2, t10 :: Tree Int
t1 = Branch (Branch (Leaf 1) (Leaf 2)) (Branch (Leaf 3) (Leaf 4))
t2 = Branch (Leaf 1) (Branch (Leaf 2) (Branch (Leaf 3) (Leaf 4)))
t3 = Branch t1 t2
t10 = mkTree [1..10]