{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -XMultiParamTypeClasses -XTypeSynonymInstances -fno-warn-orphans #-}

module Data.Collections.AVL (

-- * Concrete collection types
    AvlSet, AvlMap
    ) where

import Prelude hiding (sum,concat,lookup,map,filter,foldr,foldr1,foldl,null,reverse,(++),minimum,maximum,all,elem,concatMap,drop,head,tail,init)

import Control.Monad
import Data.Monoid
import Data.Collections
import Data.Collections.Foldable

import qualified Data.Maybe as Maybe
import qualified Data.Set.AVL as AvlSet
import qualified Data.Map.AVL as AvlMap

type AvlSet = AvlSet.Set
type AvlMap = AvlMap.Map

-----------------------------------------------------------------------------
-- Data.AvlMap
instance Foldable (AvlMap.Map k a) (k,a) where
    foldr f i m = AvlMap.foldWithKey (curry f) i m
    null = AvlMap.null

instance Ord k => Unfoldable (AvlMap.Map k a) (k,a) where
    insert = uncurry AvlMap.insert
    singleton (k,a) = AvlMap.singleton k a
    empty = AvlMap.empty

instance Ord k => Collection (AvlMap.Map k a) (k,a) where
    filter f = AvlMap.filterWithKey (curry f)

instance (Ord k, Monoid a) => Indexed (AvlMap.Map k a) k a where
    index = flip (AvlMap.!)
    adjust = AvlMap.adjust
    inDomain = member

instance (Ord k, Monoid a) => Map (AvlMap.Map k a) k a where
    isSubmapBy = AvlMap.isSubmapOfBy
    isSubset = AvlMap.isSubmapOfBy (\_ _->True)
    member = AvlMap.member
    union = AvlMap.union
    difference = AvlMap.difference
    delete = AvlMap.delete
    intersection = AvlMap.intersection
    lookup = AvlMap.lookup
    alter = AvlMap.alter
    insertWith = AvlMap.insertWith
    unionWith = AvlMap.unionWith
    intersectionWith = AvlMap.intersectionWith
    differenceWith = AvlMap.differenceMaybe
    mapWithKey = AvlMap.mapWithKey

instance Ord k => SortingCollection (AvlMap.Map k a) (k,a) where
    minView c = if null c then fail "Data.AVL.Map.minView: empty map" else return (AvlMap.findMin c, AvlMap.deleteMin c)
    -- FIXME: add support for this in AvlMap.Map


---------------------------------------
-- AvlSet

instance Foldable (AvlSet.Set a) a where
    foldr f i s = AvlSet.fold f i s
    null = AvlSet.null

instance Ord a => Unfoldable (AvlSet.Set a) a where
    insert = AvlSet.insert
    singleton = AvlSet.singleton
    empty = AvlSet.empty

instance Ord a => Collection (AvlSet.Set a) a where
    filter = AvlSet.filter

instance Ord a => Set (AvlSet.Set a) a where
    haddock_candy = haddock_candy

instance Ord a => Map (AvlSet.Set a) a () where
    isSubmapBy f x y = isSubset x y && (f () () || null (intersection x y))
    isSubset = AvlSet.isSubsetOf
    member = AvlSet.member
    union = AvlSet.union
    difference = AvlSet.difference
    intersection = AvlSet.intersection
    delete = AvlSet.delete
    insertWith _f k () = insert k
    unionWith _f = union
    intersectionWith _f = intersection
    differenceWith f s1 s2 = if f () () == Nothing then difference s1 s2 else s1
    lookup k l = if member k l then return () else fail "element not found"
    alter f k m = case f (lookup k m) of
                      Just _ -> insert k m
                      Nothing -> delete k m
    mapWithKey _f = id

instance Ord a => SortingCollection (AvlSet.Set a) a where
    minView c = if null c then fail "Data.AVL.Set.minView: empty map" else return (AvlSet.findMin c, AvlSet.deleteMin c)
    -- FIXME: add support for this in AvlSet.Set