{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fno-warn-tabs #-}

-- | 'MonadState' whose state is parameterized by the 'Monad' stack.
module Language.LOL.Typing.Lib.Control.Monad.Classes.StateFix where

import Control.Applicative (Applicative(..))
import Control.Monad
import Control.Monad.Classes
import Control.Monad.Trans.Class
import qualified Control.Monad.Trans.State.Lazy as SL
-- import qualified Control.Monad.Trans.State.Strict as SS -- TODO: when needed :)
import Data.Bool (Bool(..))
import Data.Function ((.), const)
import Data.Functor.Identity (Identity)
import GHC.Prim (Proxy#, proxy#)
import Prelude (seq)

import Language.LOL.Typing.Lib.Control.Monad.Classes.EffectsFix

-- * Type 'StateLazyFixT'

data StateLazyFixT
 (st::{-StateLazyFixT st m-}(* -> *) -> *)
 (m::{-a-}* -> *)
 (a:: *)
 =   StateLazyFixT
 { unStateLazyFixT :: SL.StateT (st (StateLazyFixT st m)) m a }
 deriving (Functor)
instance Monad m => Applicative (StateLazyFixT st m) where
	pure = return
	(<*>) = ap
instance Monad m => Monad (StateLazyFixT st m) where
	return  = StateLazyFixT . return
	m >>= f = StateLazyFixT (unStateLazyFixT m >>= unStateLazyFixT . f)
instance MonadTrans (StateLazyFixT st) where
	lift = StateLazyFixT . lift

-- ** Type 'StateLazyFix'
type StateLazyFix  st
 =   StateLazyFixT st Identity

-- * Type family 'StateFixCanDo'

type instance CanDo (StateLazyFixT s m) eff
 = StateFixCanDo s eff

type family StateFixCanDo s eff where
	StateFixCanDo s (EffStateFix s)  = 'True
	StateFixCanDo s (EffReaderFix s) = 'True
	StateFixCanDo s (EffLocalFix s)  = 'True
	StateFixCanDo s (EffWriterFix s) = 'True
	StateFixCanDo s eff              = 'False

-- * Class 'MonadStateFixN'

class Monad m => MonadStateFixN (n :: Peano) s m where
	stateFixN :: Proxy# n -> (s m -> (a, s m)) -> m a

-- | Warning: only work when 'StateLazyFixT'
-- is the outermost 'Monad' (i.e. when @n@ @~@ 'Zero'),
-- because the state is paramaterized by this 'Monad'.
instance Monad m => MonadStateFixN 'Zero s (StateLazyFixT s m) where
	stateFixN _ = StateLazyFixT . SL.state

-- ** Type 'MonadStateFixN'

-- | The @'MonadStateFix' s m@ constraint asserts that @m@ is a 'Monad' stack
-- that supports state operations on type @s@
type MonadStateFix (s::(* -> *) -> *) m
 =   MonadStateFixN (Find (EffStateFix s) m) s m

-- | Construct a state 'Monad' computation from a function
stateFix
 :: forall s m a. (MonadStateFix s m)
 => (s m -> (a, s m)) -> m a
stateFix = stateFixN (proxy# :: Proxy# (Find (EffStateFix s) m))

-- | @'put' s@ sets the state within the 'Monad' to @s@
putFix :: MonadStateFix s m => s m -> m ()
putFix s = stateFix (const ((), s))

-- | Fetch the current value of the state within the 'Monad'
getFix :: MonadStateFix s m => m (s m)
getFix = stateFix (\s -> (s, s))

-- | Gets specific component of the state,
-- using a projection function supplied.
getsFix :: MonadStateFix s m => (s m -> a) -> m a
getsFix f = do
	s <- getFix
	return (f s)

-- | Maps an old state to a new state inside a state 'Monad' layer
modifyFix :: MonadStateFix s m => (s m -> s m) -> m ()
modifyFix f = stateFix (\s -> ((), f s))

-- | A variant of 'modify' in which the computation
-- is strict in the new state.
modifyFix' :: MonadStateFix s m => (s m -> s m) -> m ()
modifyFix' f = stateFix (\s -> let s' = f s in s' `seq` ((), s'))

-- Return the 'Monad' parameter and the state.
runStateLazyFix
 :: st (StateLazyFixT st m)
 ->     StateLazyFixT st m  a
 ->                      m (a, st (StateLazyFixT st m))
runStateLazyFix s m = SL.runStateT (unStateLazyFixT m) s

-- Return the 'Monad' parameter.
evalStateLazyFix
 :: Monad m
 => st (StateLazyFixT st m)
 ->     StateLazyFixT st m a
 ->                      m a
evalStateLazyFix s m = SL.evalStateT (unStateLazyFixT m) s

-- Return the state.
execStateLazyFix
 :: Monad m
 => st (StateLazyFixT st m)
 ->     StateLazyFixT st m a
 ->                      m (st (StateLazyFixT st m))
execStateLazyFix s m = SL.execStateT (unStateLazyFixT m) s