module Hasklet1 where

import Data.List (elemIndex)
import Data.Maybe (fromJust)


-- | A generic binary tree with values at internal nodes.
data Tree a = Node a (Tree a) (Tree a)
            | Leaf
  deriving (Eq,Show)


-- | Build a balanced binary tree from a list of values.
tree :: [a] -> Tree a
tree []     = Leaf
tree (x:xs) = Node x (tree l) (tree r)
  where (l,r) = splitAt (length xs `div` 2) xs


-- | Encode a secret message.
encode :: String -> String
encode s = [s !! ix i | i <- [0..n]]
  where
    n = length s - 1
    ls = concat (levels (tree [0..n]))
    ix i = fromJust (elemIndex i ls)


-- Some example trees containing integers.
t1, t2, t3, t4 :: Tree Int
t1 = Node 1 Leaf (Node 2 Leaf Leaf)
t2 = Node 3 (Node 4 Leaf Leaf) Leaf
t3 = Node 5 t1 t2
t4 = tree (filter odd [1..100])

-- An example tree containing a secret message!
t5 :: Tree Char
t5 = tree " bstyoouu rd oerrvialentikne"


-- | Define a recursive function that sums the numbers in a tree.
--
--   >>> sumTree Leaf
--   0
--
--   >>> sumTree t3
--   15
--
--   >>> sumTree t4
--   2500
--
sumTree :: Num a => Tree a -> a
sumTree Leaf         = 0
sumTree (Node i l r) = i + sumTree l + sumTree r


-- | Define a recursive function that checks whether a given element is
--   contained in a tree.
--
--   >>> contains 57 t4
--   True
--
--   >>> contains 58 t4
--   False
--
--   >>> contains 'k' t5
--   True
--
--   >>> contains 'z' t5
--   False
--
contains :: Eq a => a -> Tree a -> Bool
contains _ Leaf         = False
contains a (Node b l r) = a == b || contains a l || contains a r


-- | Define a function for converting a binary tree of type 'Tree a' into
--   a value of type 'b' by folding an accumulator function over the tree.
--   You should start by writing a type definition for the function.
--
--   Note there is more than one correct type for this function! Part of your
--   task is to figure out the type. For inspiration, think about the types of
--   the functions `foldl` and `foldr` for lists.
-- 
--   There are conceptually (at least) two ways to "fold" this tree data
--   structure.
--
--   First, for any algebraic data type, we can systematically define a
--   corresponding "elimination" function that applies a "descructor"
--   argument to each constructor. Such elimination functions are sometimes
--   called "catamorphisms" from a corresponding concept in category theory
--   (https://en.wikipedia.org/wiki/Catamorphism).
--   
elimTree :: (a -> b -> b -> b) -> b -> Tree a -> b
elimTree f b (Node a l r) = f a (elimTree f b l) (elimTree f b r)
elimTree _ b Leaf         = b


-- data Tree a
--    = Node a (Tree a) (Tree a)
--    | Trunk a (Tree a)
--    | Leaf
-- 
-- elimTree :: (a -> b -> b -> b) -> (a -> b -> b) -> b -> Tree a -> b
-- elimTree f g b (Node a l r) = f a (elimTree f g b l) (elimTree f g b r)
-- elimTree f g b (Trunk a t)  = g a (elimTree f g b t)
-- elimTree _ _ b Leaf         = b


-- foldr :: (a -> b -> b) -> b -> [a] -> b
--
-- data List a
--    = Cons a (List a)
--    | Nil


-- | Second, for many data types, we can define a more convenient fold using
--   the same interface as the standard 'foldr' function on lists. This
--   interface is captured by the 'Foldable' type class that we'll see later.
--   Implementing this interface is typically trickier than the catamorphism,
--   but the interface is simpler and making your type an instance of
--   'Foldable' has many benefits!
--
foldTree :: (a -> b -> b) -> b -> Tree a -> b
foldTree _ b Leaf         = b
foldTree f b (Node a l r) = foldTree f (f a (foldTree f b r)) l

instance Foldable Tree where
  foldr = foldTree

genericSum :: Foldable t => t Int -> Int
genericSum = foldr (+) 0


-- | Use 'foldTree' to define a new version of 'sumTree'.
--
--   >>> sumTreeFold Leaf
--   0
--
--   >>> sumTreeFold t3
--   15
--
--   >>> sumTreeFold t4
--   2500
--
sumTreeFold :: Num a => Tree a -> a
sumTreeFold = foldTree (+) 0

-- | And using our other fold 'elimTree'.
sumTreeElim :: Num a => Tree a -> a
sumTreeElim = elimTree (\i l r -> i + l + r) 0


-- | Use 'foldTree' to define a new version of 'contains'.
--
--   >>> containsFold 57 t4
--   True
--
--   >>> containsFold 58 t4
--   False
--
--   >>> containsFold 'v' t5
--   True
--
--   >>> containsFold 'q' t5
--   False
--
containsFold :: Eq a => a -> Tree a -> Bool
containsFold x = foldTree (\y b -> x == y || b) False

-- | And using our other fold 'elimTree'.
containsElim :: Eq a => a -> Tree a -> Bool
containsElim x = elimTree (\y l r -> x == y || l || r) False


-- | Implement a function that returns a list of values contained at each
--   level of the tree. That is, it should return a nested list where the
--   first list contains the value at the root, the second list contains the
--   values at its children, the third list contains the values at the next
--   level down the tree, and so on.
--
--   Apply this function to 't5' to reveal the secret message!
--
--   >>> levels Leaf
--   []
--   
--   >>> levels t1
--   [[1],[2]]
--   
--   >>> levels t2
--   [[3],[4]]
--   
--   >>> levels t3
--   [[5],[1,3],[2,4]]
--
--   >>> levels (tree [1..10])
--   [[1],[2,6],[3,4,7,9],[5,8,10]]
--
levels :: Tree a -> [[a]]
levels = elimTree (\a l r -> [a] : merge l r) []
  where
    merge ls     []     = ls
    merge []     rs     = rs
    merge (l:ls) (r:rs) = (l ++ r) : merge ls rs