{-# LANGUAGE FlexibleContexts #-}
module Data.Collections (
	-- * Basic classes
	-- ** Monoids
	Monoid(..),
	Foldable(..),
	Functor(..),
	-- ** Collections with a size
	Sizeable(..),
	cnull,
	-- * Specialized collections
	-- ** Collections with a single type parameter
	Collection(..),
	cempty,
	cnotMember,
	-- ** Dictionaries
	Dictionary(..),
	dempty,
	dassocs,
	dkeys,
	dvalues,
	-- * Optional properties of collections
	-- ** List conversions
	Enumerable(..),
	Sequenceable(..),
	-- ** Set operations
	Unionable(..),
	cunions,
	Differenceable(..),
	Intersectable(..),
	-- ** Filter
	Filterable(..),
	-- * Existing Collections
	Array.Array,
	Map.Map,
	Seq.Seq,
	Set.Set
) where

import qualified Data.Array as Array
import Data.Foldable
import Data.Ix
import qualified Data.List as List
import qualified Data.Map as Map
import Data.Monoid
import Data.Sequence (ViewL(..), (<|))
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import Prelude hiding (foldl, elem)

-- TODO: generalize, when a class is parametrized and when not (c a vs. c), see Unionable, ...
-- TODO: think about Monoid, it needs mappend... -> replace it?
-- TODO: Filterable needs Monoid (or its replacement).
-- TODO: remove cempty/dempty and put it in a replacement for Monoid?

class Sizeable c where
	cisEmpty :: c -> Bool
	cisEmpty = (==0) . csize
	csize :: c -> Int -- ! (or cgenericSize, better: both)
	csize = cgenericSize
	cgenericSize :: (Num n) => c -> n -- ! (or csize, better: both)
	cgenericSize = fromInteger . toInteger . csize

cnull :: (Sizeable c) => c -> Bool
cnull = cisEmpty

class (Monoid c) => Collection c a | c -> a where
	cmember :: a -> c -> Bool -- !
	csingleton :: a -> c
	csingleton x = cinsert x mempty
	cinsert :: a -> c -> c -- !
	cdelete :: a -> c -> c -- !

cempty :: (Collection c a) => c
cempty = mempty

cnotMember :: (Collection c a) => a -> c -> Bool
cnotMember x c = not $ cmember x c

class (Monoid d, Enumerable d (k, a)) => Dictionary d k a | d -> k a where
	dmember :: k -> d -> Bool
	dmember key mp = maybe False (const True) $ dlookup key mp -- case dlookup key mp of
--		Just _  -> True
--		Nothing -> False
	dnotMember :: k -> d -> Bool
	dnotMember key mp = not $ dmember key mp

	dlookup :: (Monad m) => k -> d -> m a -- !
	dfindWithDefault :: a -> k -> d -> a
	dfindWithDefault def key mp = maybe def id $ dlookup key mp
	
	dsingleton :: k -> a -> d
	dsingleton key value = dinsert key value mempty
	dinsert :: k -> a -> d -> d
	dinsert = dinsertWith const
	dinsertWith :: (a -> a -> a) -> k -> a -> d -> d
	dinsertWith f = dinsertWithKey (const f)
	dinsertWithKey :: (k -> a -> a -> a) -> k -> a -> d -> d
	dinsertWithKey f key value mp = snd $ dinsertLookupWithKey f key value mp
	dinsertLookupWithKey :: (k -> a -> a -> a) -> k -> a -> d -> (Maybe a, d) -- !
	
	ddelete :: k -> d -> d
	ddelete = dalter (const Nothing)

	dadjust :: (a -> a) -> k -> d -> d
	dadjust f = dadjustWithKey (const f)
	dadjustWithKey :: (k -> a -> a) -> k -> d -> d
	dadjustWithKey f = dupdateWithKey (\k v -> Just $ f k v)
	
	dupdate :: (a -> Maybe a) -> k -> d -> d
	dupdate f = dupdateWithKey (const f)
	dupdateWithKey :: (k -> a -> Maybe a) -> k -> d -> d
	dupdateWithKey f key mp = snd $ dupdateLookupWithKey f key mp
	dupdateLookupWithKey :: (k -> a -> Maybe a) -> k -> d -> (Maybe a, d) -- !

	dalter :: (Maybe a -> Maybe a) -> k -> d -> d -- !

dempty :: (Dictionary d k a) => d
dempty = mempty

-- how to deal with ascending order of keys in Map?
dassocs :: (Dictionary d k a) => d -> [(k, a)]
dassocs = ctoList

dkeys :: (Dictionary d k a) => d -> [k]
dkeys = (map fst) . dassocs

dvalues :: (Dictionary d k a) => d -> [a]
dvalues = (map snd) . dassocs

class Enumerable c a where
	ctoList :: c -> [a] -- !
	ctoAscList :: (Ord a) => c -> [a]
	ctoAscList  = List.sort . ctoList
	ctoDescList :: (Ord a) => c -> [a]
	ctoDescList = reverse . ctoAscList

class Sequenceable c a | c -> a where -- find a better name, revise default implementations
	cfromList :: [a] -> c -- !
	cfromAscList :: (Ord a) => [a] -> c
	cfromAscList = cfromList
	cfromDistinctAscList :: (Ord a) => [a] -> c
	cfromDistinctAscList = cfromList

-- cunion == mappend, useful? see cunions
class (Monoid c) => Unionable c where
	cunion :: c -> c -> c
	cunion = mappend

-- satisfy: cunions [] == mempty (i.e. cempty)
cunions :: (Unionable c) => [c] -> c
cunions = foldl cunion mempty

class Differenceable c a b where
	cdifference :: c a -> c b -> c a -- !

class Intersectable c a b where
	cintersection :: c a -> c b -> c a -- !

-- need monoid here! cfilter (const False) gives an empty collection
class (Monoid (c a)) => Filterable c a where
	cfilter :: (a -> Bool) -> c a -> c a
	cfilter f = fst . (cpartition f)
	cpartition :: (a -> Bool) -> c a -> (c a, c a) -- !

-----------------------------------------------------------
--                 Instance declarations                 --
-----------------------------------------------------------
-- They belong into the collection modules rather than here

-- Array
instance (Ix i) => Sizeable (Array.Array i e) where
	csize = rangeSize . Array.bounds

instance (Ix i, Eq e) => Enumerable (Array.Array i e) e where
	ctoList = Array.elems

-- List
instance Sizeable [a] where
	csize = List.length
	cgenericSize = List.genericLength

instance (Eq a) => Collection [a] a where
	cmember = List.elem
	cinsert e [] = [e]
	cinsert e (x:xs) = if e == x
		then x:xs
		else x: cinsert e xs
	cdelete = List.delete

instance (Eq k) => Dictionary [(k, a)] k a where
	dlookup key xs = maybe (fail "key not found") return $ List.lookup key xs
	dinsertLookupWithKey _ key value [] = (Nothing, [(key, value)])
	dinsertLookupWithKey f key value ((key', value'):xs) = if key == key'
		then (Just value', (key, f key value value'):xs)
		else let (m, list) = dinsertLookupWithKey f key value xs
			in (m, (key', value'):list)
	dupdateLookupWithKey _ _ [] = (Nothing, [])
	dupdateLookupWithKey f key ((key', value):xs) = if key == key'
		then (Just value, maybe [] (dsingleton key) (f key value) ++ xs)
		else let (m, list) = dupdateLookupWithKey f key xs
			in (m, (key', value):list)
	dalter f key [] = maybe [] (dsingleton key) (f Nothing)
	dalter f key ((key', value):xs) = if key == key'
		then maybe [] (dsingleton key) (f $ Just value) ++ xs
		else (key', value) : dalter f key xs

instance Enumerable [a] a where
	ctoList = id

instance Sequenceable [a] a where
	cfromList = id
	cfromAscList = id
	cfromDistinctAscList = id

instance (Eq a) => Unionable [a] where
	cunion = List.union

instance (Eq a) => Differenceable [] a a where
	cdifference = (List.\\)

instance (Eq a) => Intersectable [] a a where
	cintersection = List.intersect

instance Filterable [] a where
	cpartition = List.partition

-- Map
instance Sizeable (Map.Map k a) where
	csize = Map.size

instance (Ord k) => Dictionary (Map.Map k a) k a where
	dlookup = Map.lookup
	dinsertLookupWithKey = Map.insertLookupWithKey
	dupdateLookupWithKey = Map.updateLookupWithKey
	dalter = Map.alter

instance Enumerable (Map.Map k a) a where
	ctoList = Map.elems

instance Enumerable (Map.Map k a) k where
	ctoList = Map.keys

instance Enumerable (Map.Map k a) (k, a) where
	ctoList = Map.assocs

instance (Ord k) => Sequenceable (Map.Map k a) (k, a) where
	cfromList = Map.fromList
	cfromAscList = Map.fromAscList
	cfromDistinctAscList = Map.fromDistinctAscList

instance (Ord k) => Unionable (Map.Map k a) where
	cunion = Map.union

instance (Ord k) => Differenceable (Map.Map k) a b where
	cdifference = Map.difference

instance (Ord k) => Intersectable (Map.Map k) a b where
	cintersection = Map.intersection

instance (Ord k) => Filterable (Map.Map k) a where
	cpartition = Map.partition

-- Set
instance Sizeable (Set.Set a) where
	csize = Set.size

instance (Ord a) => Collection (Set.Set a) a where
	cmember = Set.member
	cinsert = Set.insert
	cdelete = Set.delete

instance Enumerable (Set.Set a) a where
	ctoList = Set.toList
	ctoAscList = Set.toAscList

instance (Ord a) => Sequenceable (Set.Set a) a where
	cfromList = Set.fromList
	cfromAscList = Set.fromAscList
	cfromDistinctAscList = Set.fromDistinctAscList

instance (Ord a) => Unionable (Set.Set a) where
	cunion = Set.union

instance (Ord a) => Differenceable Set.Set a a where
	cdifference = Set.difference

instance (Ord a) => Intersectable Set.Set a a where
	cintersection = Set.intersection

instance (Ord a) => Filterable Set.Set a where
	cpartition = Set.partition

-- Seq

instance Sizeable (Seq.Seq a) where
	csize = Seq.length

instance Enumerable (Seq.Seq a) a where
	ctoList seq = case Seq.viewl seq of
		Seq.EmptyL  -> []
		(a :< seq') -> a : ctoList seq'

instance Sequenceable (Seq.Seq a) a where
	cfromList = Seq.fromList

instance Filterable (Seq.Seq) a where
	cpartition f seq = case Seq.viewl seq of
		Seq.EmptyL  -> (Seq.empty, Seq.empty)
		(x :< seq') -> let (left, right) = cpartition f seq'
			in if f x
				then (x <| left, right)
				else (left,      x <| right)