{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
module AlBhed.Tree
    ( -- * Types
      RedBlackTree
    , TreeZipper
      -- * Constructors
    , empty
      -- * Manipulations
    , insert
    , remove
      -- * Search
    , elem
    , search
      -- * Zipper
    , fromZipper
    , left
    , right
    , toZipper
    , up
    ) where
import           Prelude hiding (elem)
data Color = Red | Black deriving (Show, Eq)
data RedBlackTree a
    = Node Color (RedBlackTree a) a (RedBlackTree a)
    | Leaf
    deriving (Show, Eq)
empty :: RedBlackTree a
empty = Leaf
insert :: Ord a => a -> RedBlackTree a -> RedBlackTree a
insert x = makeBlack . go
    where
        go node@(Node color left value right)
            | value > x = balance color (go left) value right
            | value == x = node
            | otherwise = balance color left value (go right)
        go Leaf = Node Red Leaf x Leaf
        makeBlack (Node _ l x r) = Node Black l x r
        makeBlack x              = error "Was expecting a node"
balance :: Color -> RedBlackTree a -> a -> RedBlackTree a -> RedBlackTree a
balance Black (Node Red (Node Red a x b) y c) z d = Node Red (Node Black a x b) y (Node Black c z d)
balance Black (Node Red a x (Node Red b y c)) z d = Node Red (Node Black a x b) y (Node Black c z d)
balance Black a x (Node Red (Node Red b y c) z d) = Node Red (Node Black a x b) y (Node Black c z d)
balance Black a x (Node Red b y (Node Red c z d)) = Node Red (Node Black a x b) y (Node Black c z d)
balance color l x r = Node color l x r
elem :: Ord a => a -> RedBlackTree a -> Bool
elem _ Leaf = False
elem x node@(Node _ left value right)
    | x < value = elem x left
    | x == value = True
    | otherwise = elem x right
data Direction = L | R deriving Show
type TreeZipper a = ([(Direction, RedBlackTree a)], RedBlackTree a)
toZipper :: RedBlackTree a -> TreeZipper a
toZipper = (,) []
left :: TreeZipper a -> TreeZipper a
left (ctx, node@(Node _ l _ _)) = (ctx ++ [(L, node)], l)
left (ctx, Leaf)                = (ctx, Leaf)
right :: TreeZipper a -> TreeZipper a
right (ctx, node@(Node _ _ _ r)) = (ctx ++ [(R, node)], r)
right (ctx, Leaf)                = (ctx, Leaf)
up :: TreeZipper a -> TreeZipper a
up ([], node) = ([], node)
up (xs, node) =
    let (dir, Node c l v r) = last xs
        in case dir of
            L -> (init xs, Node c node v r)
            R -> (init xs, Node c l v node)
fromZipper :: TreeZipper a -> RedBlackTree a
search :: forall a. Ord a => a -> RedBlackTree a -> Maybe (TreeZipper a)
search x root =
    let s = go . toZipper $ root
        in case snd s of
            Leaf -> Nothing
            _    -> Just s
    where
        go :: Ord a => TreeZipper a -> TreeZipper a
        go (path, Leaf) = (path, Leaf)
        go z@(path, Node _ l value r)
            | x < value = go $ left z
            | x == value = z
            | x > value = go $ right z
remove :: (Ord a) => a -> RedBlackTree a -> RedBlackTree a
remove x node@(Node _ Leaf value Leaf)
    | x == value = Leaf
    | otherwise = node
    where
        go D1 = undefined
        go _  = undefined
    = D1
    | D2
    | D3
    | D5
    | D6
    deriving (Show)
identifyCase zipper =
        in case colors of
            [Red, Black, Black, Black] -> D4
    let colors = map (colorOf . snd . ($ zipper)) [parent, sibling, distantNephew, closeNephew]
identifyCase :: TreeZipper a -> Case
    | D4
distantNephew :: TreeZipper a -> TreeZipper a
distantNephew z = case nodeDirection z of
    L -> right . sibling $ z
    R -> left . sibling $ z
closeNephew :: TreeZipper a -> TreeZipper a
closeNephew z = case nodeDirection z of
    L -> left . sibling $ z
    R -> right . sibling $ z
colorOf :: RedBlackTree a -> Color
colorOf Leaf = Black
colorOf (Node c _ _ _) = c
data Case
        switchColor (Node Black l v r) = Node Red l v r
        switchColor (Node Red l v r) = Node Black l v r
        switchColor n = n
nodeDirection :: TreeZipper a -> Direction
nodeDirection = fst . last . fst
parent :: TreeZipper a -> TreeZipper a
parent = up
sibling :: TreeZipper a -> TreeZipper a
sibling node = case nodeDirection node of
    L -> right . parent $ node
    R -> left . parent $ node
        go D4 =
            let actions = [modify (const Leaf), modify switchColor . sibling, modify switchColor . parent]
                zipper' = foldl (\acc f -> f acc) zipper actions
                in zipper'
        go :: Case -> TreeZipper a
remove' zipper = go $ identifyCase zipper
removeSimple :: forall a. (Ord a) => TreeZipper a -> RedBlackTree a
removeSimple = fromZipper . go
    where
        go :: TreeZipper a -> TreeZipper a
        go zip@([], Node c node@(Node {}) v r@(Node {})) = undefined
            -- let leftMost = min $ toZipper r
            --     leftMod = modify (const Leaf) leftMost
            --     rootSubTree = fromZipper leftMod
            --     switch = modify (const (Node lc node left rootSubTree)) zip
            --     in upMost leftMost
        go z = remove' z
        min :: TreeZipper a -> TreeZipper a
        min (path, node@(Node _ Leaf _ _)) = (path, node)
        min zipper = min $ left zipper
withChild :: RedBlackTree a -> Bool
withChild (Node _ (Node {}) _ _) = True
withChild (Node _ _ _ (Node {})) = True
withChild _ = False
isColor :: Color -> RedBlackTree a -> Bool
isColor color (Node c _ _ _) = c == color
isColor color Leaf = color == Black
withoutChildAndBlack :: RedBlackTree a -> Bool
withoutChildAndBlack = (isColor Black &&& (not . withChild)) >>> Arrow.arr (uncurry (&&))
remove' :: forall a. (Ord a) => TreeZipper a -> TreeZipper a
remove x node = maybe node removeSimple $ search x node
modify :: (RedBlackTree a -> RedBlackTree a) -> TreeZipper a -> TreeZipper a
modify = second
fromZipper = snd . upMost
upMost :: TreeZipper a -> TreeZipper a
upMost ([], node) = ([], node)
upMost z = upMost $ up z
import Data.Bifunctor
import qualified Control.Arrow as Arrow
import Control.Arrow ((&&&), (>>>))
import qualified Control.Category as Cat
    -- * Testing purposes
    , withChild
    , isColor
    , withoutChildAndBlack
{-# LANGUAGE LambdaCase #-}