module Control.Replicate(
  replicateN
, replicateO
) where  

import Control.Applicative(liftA2)
import Control.Category((.))
import Control.Lens(AsEmpty(_Empty), Cons, (<|), ( # ))
import Control.Monad(Monad(return))
import Data.Functor((<$>))
import Data.Ord(Ord((<=)))
import Prelude(Num((-)))

-- $setup
-- >>> import Prelude

-- | Replicate, potentially changing the remaining amount of replication.
--
-- >>> let replicateN' :: (Int -> [(Int, String)]) -> Int -> [[String]]; replicateN' = replicateN in replicateN' (\n -> [(n, "abc")]) 3
-- [["abc","abc","abc"]]
--
-- >>> let replicateN' :: (Int -> [(Int, String)]) -> Int -> [[String]]; replicateN' = replicateN in replicateN' (\n -> [(n, "abc"), (n-1, "def")]) 3
-- [["abc","abc","abc"],["abc","abc","def"],["abc","def"],["def","abc"],["def","def"]]
replicateN ::
  (Monad m, Num n, Ord n, AsEmpty t, Cons t t a a) =>
  (n -> m (n, a))
  -> n
  -> m t
replicateN e n =
  if n <= 0
    then
      return (_Empty # ())
    else
      do (o, a) <- e n
         b <- replicateN e (o - 1)
         return (a <| b)

-- | Replicate with access to the current replication number.
--
-- >>> let replicateO' :: (Int -> [String]) -> Int -> [[String]]; replicateO' = replicateO in replicateO' (\n -> [show n]) 3
-- [["3","2","1"]]
--
-- >>> let replicateO' :: (Int -> [String]) -> Int -> [[String]]; replicateO' = replicateO in replicateO' (\n -> [show n, show (n-1)]) 3
-- [["3","2","1"],["3","2","0"],["3","1","1"],["3","1","0"],["2","2","1"],["2","2","0"],["2","1","1"],["2","1","0"]]
replicateO ::
  (Monad m, Num n, Ord n, AsEmpty t, Cons t t a a) =>
  (n -> m a)
  -> n
  -> m t
replicateO =
  replicateN . liftA2 (<$>) (,)