import Control.Monad (ap, liftM)
import Data.Traversable (fmapDefault, foldMapDefault)

-- Pretend this file is much larger and complicated...

data Tree a
  = Leaf a
  | Node (Tree a) (Tree a)
  deriving (Show)

someTree :: Tree Int
someTree =
  Node (Leaf 2) (Node (Node (Leaf 1) (Leaf 0)) (Leaf 5))

--                 *
--                / \
--               2   *
--                  / \
--                 *  5
--                / \
--               1  0

-- instance Foldable Tree where
--   foldMap :: (Monoid m) => (a -> m) -> Tree a -> m
--   foldMap f (Leaf a) = f a
--   foldMap f (Node lc rc) = mappend (foldMap f lc) (foldMap f rc)

transformTree :: Int -> Tree Int -> Tree Bool
transformTree x (Leaf y) = Leaf (y >= x)
transformTree x (Node lc rc) = Node (transformTree x lc) (transformTree x rc)

newtype Reader r a = Reader (r -> a)

instance Functor (Reader r) where
  fmap f (Reader r) = Reader (f . r)

instance Applicative (Reader r) where
  pure = Reader . const
  (Reader r1) <*> (Reader r2) = Reader $ \r -> let a = r2 r in r1 r a

instance Monad (Reader r) where
  (Reader ra) >>= arb = Reader $ \r ->
    let a = ra r
     in case arb a of
          (Reader rb) -> rb r

ask :: Reader r r
ask = Reader id

runReader :: Reader r a -> r -> a
runReader (Reader f) r = f r

-- transformTreeM :: Tree Int -> Reader Int (Tree Bool)
-- transformTreeM (Leaf v) = (Leaf . (v >=)) <$> ask
-- transformTreeM (Node lhs rhs) = Node <$> transformTreeM lhs <*> transformTreeM rhs
transformTreeM :: Tree Int -> Reader Int (Tree Bool)
transformTreeM (Leaf v) = do
  t <- ask
  pure $ Leaf $ v >= t
transformTreeM (Node lhs rhs) = do
  lhs' <- transformTreeM lhs
  rhs' <- transformTreeM rhs
  pure $ Node lhs' rhs'

-- newtype Writer w a = Writer (w, a)
--
-- instance Functor (Writer w) where
--   fmap f (Writer (w, a)) = Writer (w, f a)
--
-- instance (Monoid w) => Applicative (Writer w) where
--   pure :: (Monoid w) => a -> Writer w a
--   pure a = Writer (mempty, a)
--   (<*>) :: (Monoid w) => Writer w (a -> b) -> Writer w a -> Writer w b
--   (<*>) (Writer (w1, f)) (Writer (w2, a)) = Writer (mappend w1 w2, f a)
--
-- instance (Monoid w) => Monad (Writer w) where
--   (>>=) :: Writer w a -> (a -> Writer w b) -> Writer w b
--   (>>=) (Writer (w1, a)) f =
--     let (Writer (w2, b)) = f a
--      in (Writer (mappend w1 w2, b))
--
-- tell :: w -> Writer w ()
-- tell x = Writer (x, ())
--
-- tellOne :: w -> Writer [w] ()
-- tellOne x = Writer ([x], ())
--
-- runWriter :: Writer w a -> (w, a)
-- runWriter (Writer x) = x

-- liftM -- Defines fmap by only using Monad's implementation
fmapMonad :: (Monad m) => (a -> b) -> m a -> m b
fmapMonad f a = a >>= \a' -> pure (f a')

-- ap -- Defined (<*>) by only using Monad's implementation
applyMonad :: (Monad m) => m (a -> b) -> m a -> m b
applyMonad f x = f >>= \f' -> x >>= \a' -> pure (f' a')

newtype State s a = State (s -> (s, a))

instance Functor (State s) where
  -- fmap f (State g) = State $ \x -> let (s, a) = g x in (s, f a)
  fmap = liftM

instance Applicative (State s) where
  pure :: a -> State s a
  pure a = State $ \s -> (s, a)

  -- (<*>) :: State s (a -> b) -> State s a -> State s b
  -- (<*>) (State f) (State g) = State $ \y ->
  --   let
  --     (s1, f1) = f y
  --     (s2, a) = g s1
  --     in (s2, f1 a)
  (<*>) = ap

instance Monad (State s) where
  (State sa) >>= asb = State $ \s ->
    let (s', a) = sa s
        State sb = asb a
     in sb s'

put :: s -> State s ()
put s = State $ const (s, ())

get :: State s s
get = State (\s -> (s, s))

frobTreeM :: Tree Int -> State Int (Tree Int)
frobTreeM (Leaf x) = do
  s <- get
  put (s + 1)
  pure (Leaf (x + s))
frobTreeM (Node lhs rhs) = do
  lhs' <- frobTreeM lhs
  rhs' <- frobTreeM rhs
  pure (Node lhs' rhs')

runState :: State s a -> s -> (s, a)
runState (State x) = x

instance Functor Tree where
  fmap = fmapDefault

instance Foldable Tree where
  foldMap = foldMapDefault

instance Traversable Tree where
  traverse :: (Applicative f) => (a -> f b) -> Tree a -> f (Tree b)
  traverse f (Leaf a) = let fb = f a in Leaf <$> fb
  -- traverse f (Node lc rc) = Node <$> (traverse f lc) <*> (traverse f rc)
  -- traverse f (Node lc rc) = liftA2 Node (traverse f lc) (traverse f rc)
  traverse f (Node lc rc) = pure Node <*> (traverse f lc) <*> (traverse f rc)

-- let lb = traverse f lc
--     rb = traverse f rc
-- in Node <$> lb <*> rb

-- Traversable

-- Spaaaaace :pea_pod: . :coconut: =  :peanuts: