import "Prelude"

data Set a = Empty
           | Cons Integer (Set a) a (Set a)

@private
cons :: Set a -> a -> Set a -> Set a
cons l x r = Cons (size l + size r + 1) l x r

@private
rotateSingleL l x (Cons _ m y r) = cons (cons l x m) y r

@private
rotateSingleR (Cons _ l x m) y r = cons l x (cons m y r)

@private
rotateDoubleL l x (Cons _ (Cons _ m1 y m2) z r) = cons (cons l x m1) y (cons m2 z r)

@private
rotateDoubleR (Cons _ l x (Cons _ m1 y m2)) z r = cons (cons l x m1) y (cons m2 z r)

@private
balance :: Set a -> a -> Set a -> Set a
balance l x r =
    let ln = size l
        rn = size r
    in  if ln + rn < 2 then cons l x r
        else if rn > 4*ln then
            let (Cons _ rl _ rr) = r
            in  if size rl < 2 * size rr
                then rotateSingleL l x r
                else rotateDoubleL l x r
        else if ln > 4*rn then
            let (Cons _ ll _ lr) = l
            in  if size lr < 2 * size ll
                then rotateSingleR l x r
                else rotateDoubleR l x r
        else cons l x r

isEmpty :: Set a -> Boolean
isEmpty Empty = True
isEmpty _ = False

size :: Set a -> Integer
size Empty = 0
size (Cons s _ _ _) = s

member :: Ord a => a -> Set a -> Boolean
member _ Empty = False
member x (Cons _ l y r) =
    let cmp = compare x y in
    if cmp < 0
    then member x l
    else if cmp > 0
    then member x r
    else True

empty :: Set a
empty = Empty

singleton :: a -> Set a
singleton x = Cons 1 Empty x Empty

insert :: Ord a => Set a -> a -> Set a
insert Empty x = singleton x
insert c@(Cons i l y r) x = 
    let cmp = compare x y in
    if cmp < 0
    then balance (insert l x) y r
    else if cmp > 0
    then balance l y (insert r x)
    else c

delete :: Ord a => Set a -> a -> Set a
delete Empty _ = Empty
delete (Cons _ l y r) x = 
    let cmp = compare x y in
    if cmp < 0
    then balance (delete l x) y r
    else if cmp > 0
    then balance l y (delete r x)
    else aux l r
  where
    aux Empty r = r
    aux l Empty = l
    aux l r = let (x, rr) = delmin r
              in balance l x rr

delmin :: Ord a => Set a -> (a, Set a)
delmin (Cons _ Empty x r) = (x, r)
delmin (Cons _ l x r) = let (y, ll) = delmin l
                        in (y, balance ll x r)
delmin _ = fail "Empty set given to delmin."

union :: Ord a => Set a -> Set a -> Set a
union Empty b = b
union a Empty = a
union a (Cons _ l x r) =
    let l' = splitL a x
        r' = splitR a x
    in concat3 (union l' l) x (union r' r)

splitL :: Ord a => Set a -> a -> Set a
splitL Empty _ = Empty
splitL (Cons _ l x r) y =   
    let cmp = compare x y in
    if cmp < 0
    then concat3 l x (splitL r y)
    else if cmp > 0
    then splitL l y
    else l

splitR :: Ord a => Set a -> a -> Set a
splitR Empty _ = Empty
splitR (Cons _ l x r) y =   
    let cmp = compare x y in
    if cmp < 0
    then splitR r y
    else if cmp > 0
    then concat3 (splitR l y) x r
    else r

concat3 :: Ord a => Set a -> a -> Set a -> Set a
concat3 Empty x r = insert r x
concat3 l x Empty = insert l x
concat3 l@(Cons n1 l1 v1 r1) x r@(Cons n2 l2 v2 r2) =
    if 4*n1 < n2
    then balance (concat3 l x l2) v2 r2
    else if 4*n2 < n1
    then balance l1 v1 (concat3 r1 x r)
    else cons l x r

fold :: (a -> b -> a) -> a -> Set b -> a
fold f init Empty = init
fold f init (Cons _ l x r) = fold f (f (fold f init l) x) r

instance (Show a) => Show (Set a) where
    sb <+ t = p True True t
      where
        p lm rm Empty = sb <<  
            if lm 
            then (if rm then "empty" else "set [")
            else (if rm then "]" else ", ") 
        p lm rm (Cons _ l x r) = do
            p lm False l
            sb <+ x
            p False rm r

set :: Ord a => [a] -> Set a
set = foldl insert empty