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

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

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

instance Functor m => Functor (FreshT v m) where
    fmap f (FreshT xx) = FreshT $ \sp s -> fmap (first f) (xx sp s)
    x <$ FreshT xx = FreshT $ \sp s -> fmap (setFst x) (xx sp s)
      where
        setFst :: b -> (a, c) -> (b, c)
        setFst x (_, y) = (x, y)

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

instance Monad m => Monad (FreshT v m) where
    return x = FreshT $ \_ s -> return (x, s)
    FreshT aa >>= f = FreshT $ \sp s -> do
        (a, s1) <- aa sp s
        let FreshT ff = f a
        ff sp 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 $ \sp s -> pure (sp s)