{-
 Redblack tree implementation taken from:
 http://www.cs.kent.ac.uk/people/staff/smk/redblack/Untyped.hs
-}

module RedBlackTree where

import VizHaskell.TreeRepresentation
import VizHaskell.StringRepresentation
import VizHaskell.Core



data Color = Red | Black deriving Show

data RedBlackTree a = Empty | Node Color (RedBlackTree a) a (RedBlackTree a)
      deriving Show


instance TreeRepresentable RedBlackTree where
  contents Empty          = Nothing
  contents (Node _ _ x _) = Just x

  children  Empty                 = []
  children (Node _ Empty x Empty) = []
  children (Node _ l x Empty)     = [l]
  children (Node _ Empty x r)     = [r]
  children (Node _ l x r)         = [l,r]

  className Empty = Just "emptyNode"
  className (Node Red _ _ _) = Just "redNode"
  className (Node Black _ _ _) = Just "blackNode"



{- Insertion and membership test as by Okasaki -}
insert :: Ord a => a -> RedBlackTree a -> RedBlackTree a
insert x s =
    Node Black a z b
    where
    Node _ a z b = ins s
    ins Empty = Node Red Empty x Empty
    ins s@(Node Black a y b)
        | x<y = balance (ins a) y b
        | x>y = balance a y (ins b)
        | otherwise = s
    ins s@(Node Red a y b)
        | x<y = Node Red (ins a) y b
        | x>y = Node Red a y (ins b)
        | otherwise = s

member :: Ord a => a -> RedBlackTree a -> Bool
member x Empty = False
member x (Node _ a y b)
    | x<y = member x a
    | x>y = member x b
    | otherwise = True

{- balance: first equation is new,
   to make it work with a weaker invariant -}
balance :: RedBlackTree a -> a -> RedBlackTree a -> RedBlackTree a
balance (Node Red a x b) y (Node Red c z d) = Node Red (Node Black a x b) y (Node Black c z d)
balance (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 (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 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 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 a x b = Node Black a x b

{- deletion a la SMK -}
delete :: Ord a => a -> RedBlackTree a -> RedBlackTree a
delete x t =
    case del t of {Node _ a y b -> Node Black a y b; _ -> Empty}
    where
    del Empty = Empty
    del (Node _ a y b)
        | x<y = delformLeft a y b
        | x>y = delformRight a y b
            | otherwise = app a b
    delformLeft a@(Node Black _ _ _) y b = balleft (del a) y b
    delformLeft a y b = Node Red (del a) y b
    delformRight a y b@(Node Black _ _ _) = balright a y (del b)
    delformRight a y b = Node Red a y (del b)

balleft :: RedBlackTree a -> a -> RedBlackTree a -> RedBlackTree a
balleft (Node Red a x b) y c = Node Red (Node Black a x b) y c
balleft bl x (Node Black a y b) = balance bl x (Node Red a y b)
balleft bl x (Node Red (Node Black a y b) z c) = Node Red (Node Black bl x a) y (balance b z (sub1 c))

balright :: RedBlackTree a -> a -> RedBlackTree a -> RedBlackTree a
balright a x (Node Red b y c) = Node Red a x (Node Black b y c)
balright (Node Black a x b) y bl = balance (Node Red a x b) y bl
balright (Node Red a x (Node Black b y c)) z bl = Node Red (balance (sub1 a) x b) y (Node Black c z bl)

sub1 :: RedBlackTree a -> RedBlackTree a
sub1 (Node Black a x b) = Node Red a x b
sub1 _ = error "invariance violation"

app :: RedBlackTree a -> RedBlackTree a -> RedBlackTree a
app Empty x = x
app x Empty = x
app (Node Red a x b) (Node Red c y d) =
    case app b c of
        Node Red b' z c' -> Node Red(Node Red a x b') z (Node Red c' y d)
        bc -> Node Red a x (Node Red bc y d)
app (Node Black a x b) (Node Black c y d) =
    case app b c of
        Node Red b' z c' -> Node Red(Node Black a x b') z (Node Black c' y d)
        bc -> balleft a x (Node Black bc y d)
app a (Node Red b x c) = Node Red (app a b) x c
app (Node Red a x b) c = Node Red a x (app b c)



test1 = foldr insert Empty [1..10]

{-
Try this:

putStrLn $ render RepresentationString test1
putStrLn $ render (RepresentationTree RepresentationString) test1
 -}
