-----------------------------------------------------------------------------
-- |
-- Module      :  Data.FMList
-- Copyright   :  (c) Sjoerd Visscher 2009
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  sjoerd@w3future.com
-- Stability   :  experimental
-- Portability :  portable
--
-- FoldMap lists: lists represented by their 'foldMap' function.
--
-- Examples:
--
-- > -- A right-infinite list
-- > c = 1 `cons` c
--
-- > -- A left-infinite list
-- > d = d `snoc` 2
--
-- > -- A middle-infinite list ??
-- > e = c `append` d
--
-- > *> head e
-- > 1
-- > *> last e
-- > 2
-----------------------------------------------------------------------------
{-# LANGUAGE RankNTypes #-}


module Data.FMList ( 

    FMList(..)
  , transform
    
  -- * Construction  
  , empty
  , singleton
  , cons
  , snoc
  , pair
  , append
  
  , fromList
  , fromFoldable
  
  -- * Basic functions
  , null
  , length
  , genericLength
  
  , head
  , tail
  , last
  , init
  , reverse
  
  -- * Folding
  , toList
  , flatten
  , filter
  , take
  , drop
  , takeWhile
  , dropWhile
  
  , zip
  , zipWith
  
  -- * Unfolding
  , iterate
  , repeat
  , unfold
  , unfoldr
  
  ) where

import Prelude 
  ( (.), ($), ($!), flip, const, id, error
  , Maybe(..), maybe
  , Either(..), either
  , Bool(..), (&&)
  , Ord(..), Num(..), Int
  , Show(..), String, (++)
  )
import qualified Data.List as List
import Data.Monoid
import Data.Foldable
import Data.Traversable
import Control.Monad
import Control.Applicative


-- | 'FMList' is a 'foldMap' function wrapped up in a newtype.
--
newtype FMList a = FM { unFM :: forall m . Monoid m => (a -> m) -> m }

-- | The function 'transform' transforms a list by changing 
-- the map function that is passed to 'foldMap'.
--
-- It has the following property:
--
-- @transform a . transform b = transform (b . a)@
--
-- For example:
--
--  * @  m >>= g@
--
--  * @= flatten (fmap g m)@
--
--  * @= flatten . fmap g $ m@
--
--  * @= transform foldMap . transform (. g) $ m@
--
--  * @= transform ((. g) . foldMap) m@
--
--  * @= transform (\\f -> foldMap f . g) m@
-- 
transform :: (forall m. Monoid m => (a -> m) -> (b -> m)) -> FMList b -> FMList a
transform t (FM l) = FM (l . t)

-- shorthand constructors
nil          :: FMList a
nil          = FM mempty

one          :: a -> FMList a
one x        = FM ($ x)

(#)          :: FMList a -> FMList a -> FMList a
FM l # FM r  = FM (l `mappend` r)

-- exported constructors
singleton    :: a -> FMList a
singleton    = one

cons         :: a -> FMList a -> FMList a
cons x l     = one x # l

snoc         :: FMList a -> a -> FMList a
snoc l x     = l # one x

pair         :: a -> a -> FMList a
pair l r     = one l # one r

append       :: FMList a -> FMList a -> FMList a
append       = (#)



fromList     :: [a] -> FMList a
fromList     = fromFoldable

fromFoldable :: Foldable f => f a -> FMList a
fromFoldable l = FM $ flip foldMap l


mhead        :: FMList a -> Maybe a
mhead l      = getFirst (unFM l (First . Just))

null         :: FMList a -> Bool
null         = maybe True (const False) . mhead

length       :: FMList a -> Int
length       = genericLength

genericLength :: Num b => FMList a -> b
genericLength l = getSum $ unFM l (const $ Sum 1)


head         :: FMList a -> a
head l       = mhead l `fromJustOrError` "Data.FMList.head: empty list"

tail         :: FMList a -> FMList a
tail l       = if null l then error "Data.FMList.tail: empty list" else drop (1::Int) l

last         :: FMList a -> a
last l       = getLast (unFM l (Last . Just)) `fromJustOrError` "Data.FMList.last: empty list"

init         :: FMList a -> FMList a
init l       = if null l then error "Data.FMList.init: empty list" else reverse . drop (1::Int) . reverse $ l

reverse      :: FMList a -> FMList a
reverse l    = FM $ getDual . unFM l . (Dual .)


flatten      :: Foldable t => FMList (t a) -> FMList a
flatten      = transform foldMap

filter       :: (a -> Bool) -> FMList a -> FMList a
filter p     = transform (\f e -> if p e then f e else mempty)


-- transform the foldMap to foldr with state.
transformCS  :: (forall m. Monoid m => (b -> m) -> a -> (m -> s -> m) -> s -> m) -> s -> FMList a -> FMList b
transformCS t s0 l = FM $ \f -> foldr (\e r -> t f e (\a -> mappend a . r)) mempty l s0

take         :: (Ord n, Num n) => n -> FMList a -> FMList a
take n       = transformCS (\f e c i -> if i > 0 then c (f e) (i-1) else mempty) n

takeWhile    :: (a -> Bool) -> FMList a -> FMList a
takeWhile p  = transformCS (\f e c _ -> if p e then c (f e) True else mempty) True

drop         :: (Ord n, Num n) => n -> FMList a -> FMList a
drop n       = transformCS (\f e c i -> if i <= 0 then c (f e) 0 else c mempty (i-1)) n

dropWhile    :: (a -> Bool) -> FMList a -> FMList a
dropWhile p  = transformCS (\f e c ok -> if ok && p e then c mempty True else c (f e) False) True

zipWith      :: (a -> b -> c) -> FMList a -> FMList b -> FMList c
zipWith t    = transformCS (\f e2 c r1 -> foldr (\e1 _ -> c (f (t e1 e2)) (drop (1::Int) r1)) mempty r1)

zip          :: FMList a -> FMList b -> FMList (a,b)
zip          = zipWith (,)


iterate      :: (a -> a) -> a -> FMList a
iterate f x  = x `cons` iterate f (f x)

repeat       :: a -> FMList a
repeat x     = xs where xs = x `cons` xs

-- | 'unfoldr' builds an 'FMList' from a seed value from left to right.
-- The function takes the element and returns 'Nothing'
-- if it is done producing the list or returns 'Just' @(a,b)@, in which
-- case, @a@ is a appended to the result and @b@ is used as the next
-- seed value in a recursive call.
--
-- A simple use of 'unfoldr':
--
-- > *> unfoldr (\b -> if b == 0 then Nothing else Just (b, b-1)) 10
-- > fromList [10,9,8,7,6,5,4,3,2,1]
--
unfoldr      :: (b -> Maybe (a, b)) -> b -> FMList a
unfoldr g    = unfold (maybe empty (\(a, b) -> Right a `pair` Left b) . g)

-- | 'unfold' builds a list from a seed value.
-- The function takes the seed and returns an 'FMList' of values.
-- If the value is 'Right' @a@, then @a@ is appended to the result, and if the
-- value is 'Left' @b@, then @b@ is used as seed value in a recursive call.
--
-- A simple use of 'unfold' (simulating unfoldl):
--
-- > *> unfold (\b -> if b == 0 then empty else Left (b-1) `pair` Right b) 10
-- > fromList [1,2,3,4,5,6,7,8,9,10]
--
unfold       :: (b -> FMList (Either b a)) -> b -> FMList a
unfold g     = transform (\f -> either (foldMap f . unfold g) f) . g


instance Functor FMList where
  fmap g     = transform (. g)
  
instance Foldable FMList where
  foldMap    = flip unFM
  
instance Monad FMList where
  return     = one
  m >>= g    = transform (\f -> foldMap f . g) m

instance Applicative FMList where
  pure       = one
  gs <*> xs  = transform (\f g -> unFM xs (f . g)) gs
    
instance Monoid (FMList a) where
  mempty     = nil
  mappend    = (#)
  
instance MonadPlus FMList where
  mzero      = nil
  mplus      = (#)
  
instance Alternative FMList where
  empty      = nil
  (<|>)      = (#)
  
instance Show a => Show (FMList a) where
  show l     = "fromList " ++ (show $! toList l)
  

newtype WrapApp f m = WrapApp { unWrapApp :: f m }
instance (Applicative f, Monoid m) => Monoid (WrapApp f m) where
  mempty                          = WrapApp $ pure mempty
  mappend (WrapApp a) (WrapApp b) = WrapApp $ mappend <$> a <*> b

instance Traversable FMList where
  traverse f = unWrapApp . foldMap one_f where one_f x = WrapApp $ one <$> f x


fromJustOrError :: Maybe a -> String -> a
fromJustOrError ma e = maybe (error e) id ma