--State s is going to be our monadic type.
--In order to fit the design pattern, functions
--return a State type, which is a state transformer
newtype State s a = State (s -> (a, s))

--The next function is a helper function which
--evaluates the state transformer to a final state.
--That is think of the first argument as the state transformer
--The second argument as an initial state
--And the return value is the final state after evaluating
--the initial state with the state transformer
runState :: (State s a) -> s -> (a, s)
runState (State transform) initialState = transform initialState

instance Monad (State s) where
  --return gives a transformer which doesn't affect the state
  return x = 
    State (\inState ->
      let outState = inState in
      (x, outState))

  --The job of >>= is to take a state transformer and a
  --function which produces another transformer and
  --produce a transformer which is the composition of the
  --two
  transform >>= f = 
    --The return value is a transformer
    State (\s -> 
      --which first runs the input transform
      let (x, intermediateState) = runState transform s in
      --Then gets the transform from f
      let transform' = f x in
      --Then runs the intermediateState through f's transform
      runState transform' intermediateState)

--Deriving the Applicative definition from Monad
instance Applicative (State s) where
  pure = return
  (<*>) mf ma = do 
    f <- mf
    a <- ma
    return (f a)

--Deriving fmap from Monad
instance Functor (State s) where
  fmap f state = state >>= (\x -> return (f x))

--get and put and helper functions. get returns the
--current state which can be operated on. It also is of
--type State so it can be composed with later functions
get :: State s s
--the first s is the one which is used by the outiside
--the second s in the pair just maintains the calling state
get = State (\s -> (s, s))

--put takes a state and produces a transform which sets
--the state to the one provided to put.
put :: s -> State s ()
put s' = State (\s -> ((), s'))

--This is a helper function which is used to produce a
--final answer given an initial state and transformer
evalState :: State s a -> s -> a
evalState transform initial = fst (runState transform initial)

--This helper function does the same thing but returns 
--the final state
execState :: State s a -> s -> s
execState transform initial = snd (runState transform initial)


--Here is an example using the state monad
type GameState = (Bool, Int)

--This example mimics a very simple parser
--The game state considers of being on or off and
--a score. If the game is off the score stays the 
--same. If the game is on and the character is an
--a the score is incremented by 1. If the game is
--on and the character is a b the score is decremented.
--A c toggles whether the game is on or off
playGame :: String -> State GameState Int
--If the list is empty all characters have been consumed
--and we return the final score
playGame [] = do
  (_, score) <- get
  return score

playGame (x:xs) = do
  --get the current game state
  (on, score) <- get
  --pattern match on the current character and whether
  --the game is on or off
  case x of
    'a' | on -> put (on, score + 1)
    'b' | on -> put (on, score - 1)
    'c' -> put (not on, score)
    _ -> put (on, score)
  playGame xs

--An example initialState for the game
startState = (False, 0)

--Generate the state transformer for the string below
--then run startState through that transformer.
main = print $ evalState (playGame "abcaaacbbcabbab") startState