{-# LANGUAGE Haskell2010 #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Control.Monad.Trans.Fresh.Stream where

import Control.Applicative
import Control.Arrow (first)
import Control.Comonad.Cofree
import Control.Monad.Trans.Class
import Data.Fresh
import Data.Functor.Identity

newtype FreshT v m a = FreshT
    { runFreshT :: Cofree Identity v -> m (a, Cofree Identity v) }

-- | Creates a @Cofree Identity v@, based on a seed, and a function to split it.
seedStream
    :: v                -- ^ Stream seed
    -> (v -> v)         -- ^ Successor
    -> Cofree Identity v
seedStream seed sukk = go seed
  where
    go tt = tt `seq` (mm :< Identity (go mm))
      where
        mm = sukk tt

-- | A stream of integers.
integerStream :: Integer -> Cofree Identity Integer
integerStream n = seedStream n succ

-- | Creates a @Cofree Identity v@ based on an infinite list.
--
--   This function will cause an error when the list turns out to be finite.
listStream :: [v] -> Cofree Identity v
listStream (x:xs) = x :< Identity (listStream xs)
listStream [] = error "Can't create a stream from a finite list."

instance Functor m => Functor (FreshT v m) where
    fmap f (FreshT xx) = FreshT $ \s -> fmap (first f) (xx s)

instance (Functor m, Monad m) => Applicative (FreshT v m) where
    pure = return
    FreshT ff <*> FreshT xx = FreshT $ \s -> do
        (fv, s1) <- ff s
        (xv, s2) <- xx s1
        return (fv xv, s2)

instance Monad m => Monad (FreshT v m) where
    return x = FreshT $ \s -> return (x, s)
    FreshT aa >>= f = FreshT $ \s -> do
        (a, s1) <- aa s
        let FreshT ff = f a
        ff s1

-- TODO Alternative, MonadPlus
-- TODO MonadFix

instance MonadTrans (FreshT v) where
    lift m = FreshT $ \s -> do
                r <- m
                return (r, s)

-- TODO MonadIO

instance Applicative m => Fresh v (FreshT v m) where
    fresh = FreshT $ \(s :< Identity s') -> pure (s, s')