{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS -XMultiParamTypeClasses -XTypeSynonymInstances #-}

module Data.Collections.BaseInstances (
-- * Concrete collection types
    Seq.Seq, 
    IntMap.IntMap, IntSet.IntSet,
    StdSet, StdMap) 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 Data.Sequence (ViewL(..), ViewR(..))
import qualified Data.Sequence as Seq
import qualified Data.Foldable as AltFoldable

import qualified Data.Array as Array
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Maybe as Maybe
import qualified Data.Set as Set
import qualified Data.ByteString as BS
--import qualified Data.ByteString.Char8 as BSC 
-- Char8 version cannot be made as long as all bytestrings use the same type.
import qualified Data.ByteString.Lazy as BSL
import Data.Word (Word8)
-- import Data.Int (Int64)
-- import Control.Monad.Identity

type StdSet = Set.Set
type StdMap = Map.Map

-----------------------------------------------------------------------------
-- Instances
-----------------------------------------------------------------------------


-- We follow with (sample) instances of the classes.

-----------------------------------------------------------------------------
-- Data.List

instance Unfoldable [a] a where
    empty = []
    singleton = return
    insert = (:)

instance Collection [a] a where
    filter = List.filter

instance Sequence [a] a where
    take = List.take
    drop = List.drop
    splitAt = List.splitAt
    reverse = List.reverse
    front (x:xs) = return (x,xs)
    front [] = fail "front: empty sequence"
    back s = return swap `ap` front (reverse s)
        where swap (x,s) = (reverse s,x)
    cons = (:)
    snoc xs x = xs List.++ [x]
    isPrefix = List.isPrefixOf

instance Indexed [a] Int a where
    index = flip (List.!!)
    adjust f k l = left >< (f x:right)
        where (left,x:right) = List.splitAt k l
    inDomain k l = k >= 0 && k < List.length l
    
--------------------------------------
-- Data.Sequence

instance Unfoldable (Seq.Seq a) a where
    empty = Seq.empty
    singleton = return
    insert = (<|)

instance Foldable (Seq.Seq a) a where
    foldr = AltFoldable.foldr
    foldl = AltFoldable.foldl
    foldr1 = AltFoldable.foldr1
    foldl1 = AltFoldable.foldl1
    foldMap = AltFoldable.foldMap
    null = Seq.null

instance Collection (Seq.Seq a) a where
    filter f = fromList . filter f . fromFoldable    

instance Sequence (Seq.Seq a) a where
    take = Seq.take
    drop = Seq.drop
    splitAt = Seq.splitAt
    reverse = Seq.reverse
    front s = case Seq.viewl s of
                EmptyL -> fail "front: empty sequence"
                a :< s -> return (a,s)
    back s = case Seq.viewr s of
                EmptyR -> fail "back: empty sequence"
                s :> a -> return (s,a)
    cons = (Seq.<|)
    snoc = (Seq.|>)

instance Indexed (Seq.Seq a) Int a where
    index = flip Seq.index
    adjust = Seq.adjust
    inDomain k l = k >= 0 && k < Seq.length l

------------------------
-- Data.ByteString

instance Foldable BS.ByteString Word8 where
    fold = foldr (+) 0
    foldr = BS.foldr
    foldl = BS.foldl
    foldr1 = BS.foldr1
    foldl1 = BS.foldl1
    null = BS.null
    size = BS.length

instance Unfoldable BS.ByteString Word8 where
    empty = BS.empty
    singleton = BS.singleton
    insert = BS.cons    

instance Collection BS.ByteString Word8 where
    filter = BS.filter

instance Sequence BS.ByteString Word8 where
    take = BS.take
    drop = BS.drop
    splitAt = BS.splitAt
    reverse = BS.reverse
    front s = if BS.null s then fail "front: empty ByteString" else return (BS.head s,BS.tail s)
    back s = if BS.null s 
             then fail "back: empty sequence" 
             else let (s',x) = BS.splitAt (BS.length s - 1) s in return (s', BS.head x)
    cons = BS.cons
    snoc = BS.snoc

instance Indexed BS.ByteString Int Word8  where
    index = flip BS.index
    adjust = error "Indexed.ajust: not supported by ByteString"
    inDomain k l = k >= 0 && k < BS.length l

------------------------
-- Data.ByteString.Lazy

instance Foldable BSL.ByteString Word8 where
    fold = foldr (+) 0
    foldr = BSL.foldr
    foldl = BSL.foldl
    foldr1 = BSL.foldr1
    foldl1 = BSL.foldl1
    null = BSL.null
    size = fromIntegral . BSL.length

instance Unfoldable BSL.ByteString Word8 where
    empty = BSL.empty
    singleton = BSL.singleton
    insert = BSL.cons
    
instance Collection BSL.ByteString Word8 where
    filter = BSL.filter

instance Sequence BSL.ByteString Word8 where
    take = BSL.take . fromIntegral
    drop = BSL.drop . fromIntegral
    splitAt = BSL.splitAt . fromIntegral
    reverse = BSL.reverse
    front s = if BSL.null s then fail "front: empty ByteString" else return (BSL.head s,BSL.tail s)
    back s = if BSL.null s 
             then fail "back: empty sequence" 
             else let (s',x) = BSL.splitAt (BSL.length s - 1) s in return (s', BSL.head x)
    cons = BSL.cons
    snoc = BSL.snoc

instance Indexed BSL.ByteString Int Word8  where
    index = flip BSL.index . fromIntegral
    adjust = error "Indexed.ajust: not supported by ByteString.Lazy yet"
    inDomain k l = k >= 0 && k < size l

--------------------------------------
-- Data.Array

instance Array.Ix i => Indexed (Array.Array i e) i e where
    index = flip (Array.!)
    adjust f k a = a Array.// [(k,f (a ! k))]
    inDomain k a = Array.inRange (Array.bounds a) k
    (//) a l = (Array.//) a (toList l)

instance Array.Ix i => Array (Array.Array i e) i e where
    array b l = Array.array b (toList l)
    bounds = Array.bounds

-----------------------------------------------------------------------------
-- Data.Map

-- TODO: write the instance based on foldMap
instance Foldable (Map.Map k a) (k,a) where
    foldr f i m = Map.foldWithKey (curry f) i m
    null = Map.null

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

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

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

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

instance Ord k => SortingCollection (Map.Map k a) (k,a) where
    minView = Map.minViewWithKey

-----------------------------------------------------------------------------
-- Data.IntMap
instance Foldable (IntMap.IntMap a) (Int,a) where
    null = IntMap.null
    size = IntMap.size
    foldr f i m = IntMap.foldWithKey (curry f) i m

instance Unfoldable (IntMap.IntMap a) (Int,a) where
    insert = uncurry IntMap.insert
    singleton (k,a) = IntMap.singleton k a
    empty = IntMap.empty

instance Collection (IntMap.IntMap a) (Int,a) where
    filter f = IntMap.filterWithKey (curry f)

instance Indexed (IntMap.IntMap a) Int a where
    index = flip (IntMap.!)
    adjust = IntMap.adjust
    inDomain = member

instance Map (IntMap.IntMap a) Int a where
    isSubmapBy = IntMap.isSubmapOfBy
    isSubset = IntMap.isSubmapOfBy (\_ _->True)
    member = IntMap.member
    union = IntMap.union
    difference = IntMap.difference
    delete = IntMap.delete
    intersection = IntMap.intersection
    lookup = IntMap.lookup
    alter = IntMap.alter
    insertWith = IntMap.insertWith
    unionWith = IntMap.unionWith
    intersectionWith = IntMap.intersectionWith
    differenceWith = IntMap.differenceWith
    mapWithKey = IntMap.mapWithKey

-----------------------------------------------------------------------------
-- Data.Set

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

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

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

instance Ord a => Map (Set.Set a) a () where
    isSubset = Set.isSubsetOf
    isSubmapBy f x y = isSubset x y && (f () () || null (intersection x y))
    member = Set.member
    union = Set.union
    difference = Set.difference
    intersection = Set.intersection
    delete = Set.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 (Set.Set a) a where
    minView c = if null c then fail "Data.Set.minView: empty set" else return (Set.findMin c, Set.deleteMin c)
    -- FIXME: add support for this in Data.Set

-----------------------------------------------------------------------------
-- Data.IntSet

instance Foldable IntSet.IntSet Int where
    foldr f i s = IntSet.fold f i s
    fold = foldl (+) 0
    null = IntSet.null
    size = IntSet.size

instance Unfoldable IntSet.IntSet Int where
    insert = IntSet.insert
    singleton = IntSet.singleton
    empty = IntSet.empty

instance Collection IntSet.IntSet Int where
    filter = IntSet.filter

instance Set IntSet.IntSet Int where
    haddock_candy = haddock_candy

instance Map IntSet.IntSet Int () where
    isSubmapBy f x y = isSubset x y && (f () () || null (intersection x y))
    isSubset = IntSet.isSubsetOf
    member = IntSet.member
    union = IntSet.union
    difference = IntSet.difference
    intersection = IntSet.intersection
    delete = IntSet.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