{-# Language DeriveFunctor#-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Data.BinaryTree
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--
-- Several types of Binary trees.
--
--------------------------------------------------------------------------------
module Data.BinaryTree where

import           Algorithms.DivideAndConquer
import           Control.DeepSeq
import           Data.Bifunctor.Apply
import           Data.List.NonEmpty (NonEmpty)
import           Data.Maybe (mapMaybe)
import           Data.Measured.Class
import           Data.Measured.Size
import           Data.Semigroup.Foldable
import qualified Data.Tree as Tree
import           Data.Tree.Util (TreeNode(..))
import qualified Data.Vector as V
import           GHC.Generics (Generic)
import           Test.QuickCheck
--------------------------------------------------------------------------------

-- | Binary tree that stores its values (of type a) in the leaves. Internal
-- nodes store something of type v.
data BinLeafTree v a = Leaf !a
                     | Node (BinLeafTree v a) !v (BinLeafTree v a)
                     deriving (Show,Read,Eq,Ord,Functor,Generic)

instance (NFData v, NFData a) => NFData (BinLeafTree v a)

-- | smart constructor
node     :: Measured v a => BinLeafTree v a -> BinLeafTree v a -> BinLeafTree v a
node l r = Node l (measure l <> measure r) r


instance Bifunctor BinLeafTree where
  bimap f g = \case
    Leaf x     -> Leaf $ g x
    Node l k r -> Node (bimap f g l) (f k) (bimap f g r)

instance Measured v a => Measured v (BinLeafTree v a) where
  measure (Leaf x)     = measure x
  measure (Node _ v _) = v


instance Foldable (BinLeafTree v) where
  foldMap f (Leaf a)     = f a
  foldMap f (Node l _ r) = foldMap f l `mappend` foldMap f r

instance Foldable1 (BinLeafTree v)

instance Traversable (BinLeafTree v) where
  traverse f (Leaf a)     = Leaf <$> f a
  traverse f (Node l v r) = Node <$> traverse f l <*> pure v <*> traverse f r

instance Measured v a => Semigroup (BinLeafTree v a) where
  l <> r = node l r

instance (Arbitrary a, Arbitrary v) => Arbitrary (BinLeafTree v a) where
  arbitrary = sized f
    where f n | n <= 0    = Leaf <$> arbitrary
              | otherwise = do
                              l <- choose (0,n-1)
                              Node <$> f l <*> arbitrary <*> f (n-l-1)

-- | Create a balanced tree, i.e. a tree of height \(O(\log n)\) with the
-- elements in the leaves.
--
-- \(O(n)\) time.
asBalancedBinLeafTree :: NonEmpty a -> BinLeafTree Size (Elem a)
asBalancedBinLeafTree = divideAndConquer1 (Leaf . Elem)
-- -- the implementation below produces slightly less high trees, but runs in
-- -- \(O(n \log n)\) time, as on every level it traverses the list passed down.
-- asBalancedBinLeafTree ys = asBLT (length ys') ys' where ys' = toList ys

--     asBLT _ [x] = Leaf (Elem x)
--     asBLT n xs  = let h       = n `div` 2
--                       (ls,rs) = splitAt h xs
--                   in node (asBLT h ls) (asBLT (n-h) rs)

-- | Given a function to combine internal nodes into b's and leafs into b's,
-- traverse the tree bottom up, and combine everything into one b.
foldUp                  :: (b -> v -> b -> b) -> (a -> b) -> BinLeafTree v a -> b
foldUp _ g (Leaf x)     = g x
foldUp f g (Node l x r) = f (foldUp f g l) x (foldUp f g r)


-- | Traverses the tree bottom up, recomputing the assocated values.
foldUpData     :: (w -> v -> w -> w) -> (a -> w) -> BinLeafTree v a -> BinLeafTree w a
foldUpData f g = foldUp f' Leaf
  where
    f' l v r = Node l (f (access' l) v (access' r)) r

    access' (Leaf x)     = g x
    access' (Node _ v _) = v

-- | Takes two trees, that have the same structure, and uses the provided
-- functions to "zip" them together
zipExactWith                                  :: (u -> v -> w)
                                              -> (a -> b -> c)
                                              -> BinLeafTree u a
                                              -> BinLeafTree v b
                                              -> BinLeafTree w c
zipExactWith _ g (Leaf x)     (Leaf y)        = Leaf (x `g` y)
zipExactWith f g (Node l m r) (Node l' m' r') = Node (zipExactWith f g l l')
                                                     (m `f` m')
                                                     (zipExactWith f g r r')
zipExactWith _ _ _            _               =
    error "zipExactWith: tree structures not the same "



--------------------------------------------------------------------------------
-- * Converting into a Data.Tree

-- | \( O(n) \) Convert binary tree to a rose tree, aka 'Tree.Tree'.
toRoseTree              :: BinLeafTree v a -> Tree.Tree (TreeNode v a)
toRoseTree (Leaf x)     = Tree.Node (LeafNode x) []
toRoseTree (Node l v r) = Tree.Node (InternalNode v) (map toRoseTree [l,r])

-- | 2-dimensional ASCII drawing of a tree.
drawTree :: (Show v, Show a) => BinLeafTree v a -> String
drawTree = Tree.drawTree . fmap show . toRoseTree

--------------------------------------------------------------------------------
-- * Internal Node Tree

-- | Binary tree in which we store the values of type a in internal nodes.
data BinaryTree a = Nil
                  | Internal (BinaryTree a) !a (BinaryTree a)
                  deriving (Show,Read,Eq,Ord,Functor,Foldable,Traversable,Generic)
instance NFData a => NFData (BinaryTree a)

instance Arbitrary a => Arbitrary (BinaryTree a) where
  arbitrary = sized f
    where f n | n <= 0    = pure Nil
              | otherwise = do
                              l <- choose (0,n-1)
                              Internal <$> f l <*> arbitrary <*> f (n-l-1)

-- | Get the element stored at the root, if it exists
access                  :: BinaryTree a -> Maybe a
access Nil              = Nothing
access (Internal _ x _) = Just x

-- | Create a balanced binary tree.
--
-- running time: \(O(n)\)
asBalancedBinTree :: [a] -> BinaryTree a
asBalancedBinTree = mkTree . V.fromList
  where
    mkTree v = let n = V.length v
                   h = n `div` 2
                   x = v V.! h
               in if n == 0 then Nil
                            else Internal (mkTree $ V.slice 0 h v) x
                                          (mkTree $ V.slice (h+1) (n - h -1) v)

-- | Fold function for folding over a binary tree.
foldBinaryUp                      :: b -> (a -> b -> b -> b)
                                  -> BinaryTree a -> BinaryTree (a,b)
foldBinaryUp _ _ Nil              = Nil
foldBinaryUp e f (Internal l x r) = let l' = foldBinaryUp e f l
                                        r' = foldBinaryUp e f r
                                        g  = maybe e snd . access
                                        b  = f x (g l') (g r')
                                    in Internal l' (x,b) r'

-- | Convert a @BinaryTree@ into a RoseTree
toRoseTree'                  :: BinaryTree a -> Maybe (Tree.Tree a)
toRoseTree' Nil              = Nothing
toRoseTree' (Internal l v r) = Just $ Tree.Node v $ mapMaybe toRoseTree' [l,r]

-- | Draw a binary tree.
drawTree' :: Show a => BinaryTree a -> String
drawTree' = maybe "Nil" (Tree.drawTree . fmap show) . toRoseTree'