-- Tree relabelling example from chapter 12 of Programming in Haskell,
-- Graham Hutton, Cambridge University Press, 2016.

-- The state monad

type State = Int

newtype ST a = S (State -> (a,State))

app :: ST a -> State -> (a,State)
app (S st) x = st x

instance Functor ST where
   -- fmap :: (a -> b) -> ST a -> ST b
   fmap g st = S (\s -> let (x,s') = app st s in (g x, s'))

instance Applicative ST where
   -- pure :: a -> ST a
   pure x = S (\s -> (x,s))

   -- (<*>) :: ST (a -> b) -> ST a -> ST b
   stf <*> stx = S (\s ->
      let (f,s')  = app stf s
          (x,s'') = app stx s' in (f x, s''))

instance Monad ST where
   -- (>>=) :: ST a -> (a -> ST b) -> ST b
   st >>= f = S (\s -> let (x,s') = app st s in app (f x) s')

-- Relabelling trees

inc :: ST String 
inc = S (\n -> ("", n+1))

double :: ST String
double = S (\n -> ("", 2*n))

report :: ST String
report = S (\n ->( show  n, n))

saveState :: ST State
saveState = S (\n -> (n,n))

resetState :: State -> ST State
resetState m = S (\n -> (0,m))


dothis = do
    inc
    double
    inc
    s <- saveState
    inc
    inc
    str1 <- report

    str2 <- if even s then do 
               resetState s
               str3 <- report
               return str3
            else do
               resetState (s+1)
               str3 <- report
               return str3

    return ("State was " ++ str1 ++ " then it became " ++ str2)