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

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

data Triplet a = Triplet a a a

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

-- | Creates a @Cofree Triplet v@, based on a seed, and functions to split
--   and succeed it.
seedDelta
    :: v                -- ^ Delta seed
    -> (v -> v)         -- ^ Successor
    -> (v -> (v, v))    -- ^ Splitter
    -> Cofree Triplet v
seedDelta seed sukk splt = go seed
  where
    go tt = tt `seq` (tt :< Triplet (go ll) (go mm) (go rr))
      where
        mm = sukk tt
        (ll, rr) = splt tt

-- | A delta of integers.
integerDelta :: Integer -> Cofree Triplet Integer
integerDelta n = seedDelta n succ (\m -> let m' = 2 * m in (m', m' + 1))

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

instance Applicative m => Applicative (FreshT v m) where
    pure x = FreshT $ \s -> pure (x, s)
    FreshT ff <*> FreshT xx = FreshT $ \(_ :< Triplet l _ r)
        -> (\(f, sf) (x, _) -> (f x, sf)) <$> ff l <*> xx r

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 :< Triplet _ s' _) -> pure (s, s')