{-# LANGUAGE UndecidableInstances, TypeOperators, FlexibleContexts, MultiParamTypeClasses, FlexibleInstances, TypeFamilies #-}
module Data.Monoid.Generator
    ( module Data.Monoid.Reducer
    , Generator
    , Elem
    , mapReduce
    , mapTo
    , mapFrom
    , reduce
    , Keys(Keys, getKeys)
    , Values(Values, getValues)
    , Char8(Char8, getChar8)
    ) where

import Data.Word (Word8)
import Data.Text (Text)
import Data.Foldable (fold,foldMap)
import qualified Data.Text as Text
import qualified Data.ByteString as Strict (ByteString, foldl')
import qualified Data.ByteString.Char8 as Strict8 (foldl')
import qualified Data.ByteString.Lazy as Lazy (ByteString, toChunks)
import qualified Data.ByteString.Lazy.Char8 as Lazy8 (toChunks)
import qualified Data.Sequence as Seq
import Data.FingerTree (Measured, FingerTree)
import Data.Sequence (Seq)
import qualified Data.Set as Set
import Data.Set (Set)
import qualified Data.IntSet as IntSet
import Data.IntSet (IntSet)
import qualified Data.IntMap as IntMap
import Data.IntMap (IntMap)
import qualified Data.Map as Map
import Data.Map (Map)

import Control.Parallel.Strategies
import Data.Monoid.Reducer

-- minimal definition mapReduce or affixMapReduce
class Generator c where
    type Elem c :: * 
    mapReduce :: (e `Reducer` m) => (Elem c -> e) -> c -> m
    mapTo     :: (e `Reducer` m) => (Elem c -> e) -> m -> c -> m 
    mapFrom   :: (e `Reducer` m) => (Elem c -> e) -> c -> m -> m

    mapReduce f = mapTo f mempty
    mapTo f m = mappend m . mapReduce f
    mapFrom f = mappend . mapReduce f

instance Generator Strict.ByteString where
    type Elem Strict.ByteString = Word8
    mapTo f = Strict.foldl' (\a -> snoc a . f)

instance Generator Lazy.ByteString where
    type Elem Lazy.ByteString = Word8
    mapReduce f = fold . parMap rwhnf (mapReduce f) . Lazy.toChunks

newtype Char8 c = Char8 { getChar8 :: c } 

instance Generator (Char8 Strict.ByteString) where
    type Elem (Char8 Strict.ByteString) = Char
    mapTo f m = Strict8.foldl' (\a -> snoc a . f) m . getChar8

instance Generator (Char8 Lazy.ByteString) where
    type Elem (Char8 Lazy.ByteString) = Char
    mapReduce f = fold . parMap rwhnf (mapReduce f . Char8) . Lazy8.toChunks . getChar8

instance Generator Text where
    type Elem Text = Char
    mapTo f = Text.foldl' (\a -> snoc a . f)

instance Generator [c] where
    type Elem [c] = c
    mapReduce f = foldMap (unit . f)

instance Measured v e => Generator (FingerTree v e) where
    type Elem (FingerTree v e) = e
    mapReduce f = foldMap (unit . f)

instance Generator (Seq c) where
    type Elem (Seq c) = c
    mapReduce f = foldMap (unit . f)

instance Generator IntSet where
    type Elem IntSet = Int
    mapReduce f = mapReduce f . IntSet.toList

instance Generator (Set a) where
    type Elem (Set a) = a
    mapReduce f = mapReduce f . Set.toList

instance Generator (IntMap v) where
    type Elem (IntMap v) = (Int,v)
    mapReduce f = mapReduce f . IntMap.toList

instance Generator (Map k v) where
    type Elem (Map k v) = (k,v) 
    mapReduce f = mapReduce f . Map.toList

newtype Keys c = Keys { getKeys :: c } 

instance Generator (Keys (IntMap v)) where
    type Elem (Keys (IntMap v)) = Int
    mapReduce f = mapReduce f . IntMap.keys . getKeys

instance Generator (Keys (Map k v)) where
    type Elem (Keys (Map k v)) = k
    mapReduce f = mapReduce f . Map.keys . getKeys

newtype Values c = Values { getValues :: c } 

instance Generator (Values (IntMap v)) where
    type Elem (Values (IntMap v)) = v
    mapReduce f = mapReduce f . IntMap.elems . getValues

instance Generator (Values (Map k v)) where
    type Elem (Values (Map k v)) = v
    mapReduce f = mapReduce f . Map.elems . getValues

{-# SPECIALIZE reduce :: (Word8 `Reducer` m) => Strict.ByteString -> m #-}
{-# SPECIALIZE reduce :: (Word8 `Reducer` m) => Lazy.ByteString -> m #-}
{-# SPECIALIZE reduce :: (Char `Reducer` m) => Char8 Strict.ByteString -> m #-}
{-# SPECIALIZE reduce :: (Char `Reducer` m) => Char8 Lazy.ByteString -> m #-}
{-# SPECIALIZE reduce :: (c `Reducer` m) => [c] -> m #-}
{-# SPECIALIZE reduce :: (Generator (FingerTree v e), e `Reducer` m) => FingerTree v e -> m #-}
{-# SPECIALIZE reduce :: (Char `Reducer` m) => Text -> m #-}
{-# SPECIALIZE reduce :: (e `Reducer` m) => Seq e -> m #-}
{-# SPECIALIZE reduce :: (Int `Reducer` m) => IntSet -> m #-}
{-# SPECIALIZE reduce :: (a `Reducer` m) => Set a -> m #-}
{-# SPECIALIZE reduce :: ((Int,v) `Reducer` m) => IntMap v -> m #-}
{-# SPECIALIZE reduce :: ((k,v) `Reducer` m) => Map k v -> m #-}
{-# SPECIALIZE reduce :: (Int `Reducer` m) => Keys (IntMap v) -> m #-}
{-# SPECIALIZE reduce :: (k `Reducer` m) => Keys (Map k v) -> m #-}
{-# SPECIALIZE reduce :: (v `Reducer` m) => Values (IntMap v) -> m #-}
{-# SPECIALIZE reduce :: (v `Reducer` m) => Values (Map k v) -> m #-}
reduce :: (Generator c, Elem c `Reducer` m) => c -> m
reduce = mapReduce id