{- |
Module      : ConClusion.BinaryTree
Description : Custom binary tree type with some special functions
Copyright   : Phillip Seeber, 2023
License     : AGPL-3
Maintainer  : phillip.seeber@googlemail.com
Stability   : experimental
Portability : POSIX, Windows
-}
module ConClusion.BinaryTree (
  BinTree (..),
  root,
  takeBranchesWhile,
  takeLeafyBranchesWhile,
)
where

import Data.Aeson hiding (Array)
import Data.Massiv.Array as Massiv hiding (IndexException)
import RIO

-- | A binary tree.
data BinTree e = Leaf e | Node e (BinTree e) (BinTree e)
  deriving (Eq, Show, Generic)

instance (FromJSON e) => FromJSON (BinTree e)

instance (ToJSON e) => ToJSON (BinTree e)

instance Functor BinTree where
  fmap f (Leaf a) = Leaf (f a)
  fmap f (Node a l r) = Node (f a) (fmap f l) (fmap f r)

-- | Look at the root of a binary tree.
root :: BinTree e -> e
root (Leaf e) = e
root (Node e _ _) = e

{- | Steps down each branch of a tree until some criterion is satisfied or the
end of the branch is reached. Each end of the branch is added to a result.
-}
takeBranchesWhile :: (a -> Bool) -> BinTree a -> Massiv.Vector DL a
takeBranchesWhile chk tree = go tree (Massiv.empty @DL)
 where
  go (Leaf v) acc = if chk v then acc `snoc` v else acc
  go (Node v l r) acc =
    let vAcc = if chk v then acc `snoc` v else acc
        lAcc = go l vAcc
        rAcc = go r lAcc
     in if chk v then rAcc else vAcc

{- | Takes the first value in each branch, that does not fullfill the criterion
anymore and adds it to the result. Terminal leafes of the branches are always
taken.
-}
takeLeafyBranchesWhile :: (a -> Bool) -> BinTree a -> Massiv.Vector DL a
takeLeafyBranchesWhile chk tree = go tree (Massiv.empty @DL)
 where
  go (Leaf v) acc = acc `snoc` v
  go (Node v l r) acc =
    let vAcc = if chk v then acc else acc `snoc` v
        lAcc = go l vAcc
        rAcc = go r lAcc
     in if chk v then rAcc else vAcc