module Set (Set
  , empty   -- Set a
  , isEmpty -- Set a-> Bool
  , insert  -- Ord a=> Set a-> a-> Set a
  , remove  -- Ord a=> Set a-> a-> Set a
  , elem    -- Ord a=> a-> Set a-> Bool
  , enum    -- Ord a=> Set a-> [a]
  ) where

import Prelude hiding(elem)

type Set a = AVLTree a

data AVLTree a = Null 
               | Node Int 
                         (AVLTree a) 
                         a 
                         (AVLTree a)

    deriving Eq
instance Show a=> Show (AVLTree a) where
  show t = shw 0 t where
    shw _ Null = []
    shw n (Node _ l a r) = (shw (n+1) l)
                              ++ spc n ++ show a  ++"\n" 
                              ++ (shw (n+1) r)
    spc n = concat (replicate n "   ")

empty :: AVLTree a
empty = Null    

isEmpty :: AVLTree a-> Bool
isEmpty Null = True
isEmpty _       = False

ht :: AVLTree a-> Int
ht Null           = 0
ht (Node h _ _ _) = h  

mkNode :: AVLTree a-> a-> AVLTree a-> AVLTree a
mkNode l n r = Node h l n r where
               h = 1+ max (ht l) (ht r)

rotl :: AVLTree a-> AVLTree a
rotl (Node _ xt y (Node _ yt x zt)) = 
  mkNode (mkNode xt y yt) x zt

rotr :: AVLTree a-> AVLTree a
rotr (Node _ (Node _ ut y vt) x rt) = 
  mkNode ut y (mkNode vt x rt)

bias :: AVLTree a-> Int
bias (Node _ lt _ rt) = ht lt - ht rt  

mkAVL :: AVLTree a-> a-> AVLTree a-> AVLTree a
mkAVL xt y zt 
  | hx+1< hz && 
    bias zt < 0 = rotl (mkNode xt y zt)
  | hx+1< hz    = rotl (mkNode xt y (rotr zt))
  | hz+1< hx && 
    bias xt > 0 = rotr (mkNode xt y zt)
  | hz+1< hx    = rotr (mkNode (rotl xt) y zt)
  | otherwise   = mkNode xt y zt
    where hx= ht xt; hz= ht zt

insert :: Ord a=> AVLTree a-> a-> AVLTree a
insert Null a = mkNode Null a Null
insert (Node n l a r) b
  | b < a  = mkAVL (insert l b) a r
  | b == a = Node n l a r
  | b > a  = mkAVL l a (insert r b)

remove :: Ord a=> AVLTree a-> a-> AVLTree a
remove Null x = Null
remove (Node h l y r) x
  | x < y  = mkAVL (remove l x) y r
  | x == y = join l r
  | x > y  = mkAVL l y (remove r x)

join ::  AVLTree a-> AVLTree a-> AVLTree a
join xt Null = xt
join xt yt   = mkAVL xt u nu where
  (u, nu) = splitTree yt
  splitTree :: AVLTree a-> (a, AVLTree a)
  splitTree (Node h Null a t) = (a, t)
  splitTree (Node h lt a rt) =
    (u, mkAVL nu a rt) where
       (u, nu) = splitTree lt

elem :: Ord a=> a-> AVLTree a-> Bool
elem x Null = False 
elem x (Node _ lt a rt) | x == a  = True
                        | x <  a  = x `elem` lt
                        | x >  a  = x `elem` rt

foldT :: (a-> b-> b-> b)-> b-> AVLTree a-> b    
foldT f e Null = e
foldT f e (Node _ l a r) =
   f a (foldT f e l) (foldT f e r)

enum :: Ord a=> AVLTree a-> [a]
enum = foldT (\x t1 t2-> t1++ [x]++ t2) []

testtree :: AVLTree Int
testtree =
 foldl insert empty  [4,5,7,2,1,3,6,8,9]