{-# LANGUAGE InstanceSigs #-}

module Lecture09 where

import Data.Map ( Map )
import qualified Data.Map as Map
import Data.Monoid ( (<>) )
import Control.Applicative
import Data.Foldable ( foldrM, foldlM )
import Control.Monad.Trans.Reader
import Control.Monad.Trans.Class

type Varname = String

data AExpr = Con Integer
           | Var Varname
           | Add AExpr AExpr
           | Mul AExpr AExpr
           deriving ( Eq, Show, Read )

type Assignment = Map Varname Integer

val' :: String -> Assignment -> Integer
val' v = Map.findWithDefault 0 v


newtype Writer w a = Writer { runWriter :: (a, w) }

tell :: w -> Writer w ()
tell msg = Writer ((), msg)

instance Functor (Writer w) where
    fmap f x = let (r, w) = runWriter x in Writer (f r, w)

instance Monoid w => Applicative (Writer w) where
    pure x = Writer (x, mempty)
    (<*>) :: Writer w (a -> b) -> Writer w a -> Writer w b
    f <*> x = let (f', wf) = runWriter f
                  (x', wx) = runWriter x
              in Writer (f' x', wf <> wx)

instance Monoid w => Monad (Writer w) where
    (>>=) :: Writer w a -> (a -> Writer w b) -> Writer w b
    x >>= f = let (x', wx) = runWriter x
                  (r, wfx) = runWriter (f x')
              in Writer (r, wx <> wfx)


plusLog :: (Num a, Show a) => a -> a -> Writer String a
plusLog x y = tell (show x ++ " + " ++ show y ++ "; ") >> pure (x + y)

ex_writer0 = runWriter $ foldrM plusLog 0 [1..10]
ex_writer1 = runWriter $ foldlM plusLog 0 [1..10]

newtype WriterT w m a = WriterT { runWriterT :: m (a, w) }

writer :: Applicative m => (a, w) -> WriterT w m a
writer (x, w) = WriterT $ pure (x, w)

tellT :: Monad m => w -> WriterT w m ()
tellT msg = writer ((), msg)

-- or import Control.Arrow ( first )
first :: (a -> b) -> (a, c) -> (b, c)
first f (x, y) = (f x, y)

instance Functor m => Functor (WriterT w m) where
    fmap :: (a -> b) -> WriterT w m a -> WriterT w m b
    -- fmap (first f) :: m (a, w) -> m (b, w)
    fmap f x = WriterT . fmap (first f) $ runWriterT x

instance (Monoid w, Applicative m) => Applicative (WriterT w m) where
    pure x = writer (x, mempty)

    (<*>) :: WriterT w m (a -> b) -> WriterT w m a -> WriterT w m b
    -- runWriterT x :: m (a, w)
    -- runWriterT f :: m (a -> b, w)
    f <*> x = WriterT $ liftA2 app (runWriterT f) (runWriterT x)
      where
        -- app :: (a -> b, w) -> (a, w) -> (b, w)
        app (f, wf) (x, wx) = (f x, wf <> wx)

instance (Monoid w, Monad m) => Monad (WriterT w m) where
    (>>=) :: WriterT w m a -> (a -> WriterT w m b) -> WriterT w m b
    x >>= f = WriterT $ do
        (x', wx) <- runWriterT x
        (r, rfx) <- runWriterT (f x')
        pure (r, wx <> rfx)

instance Monoid w => MonadTrans (WriterT w) where
    lift act = WriterT $ fmap (\x -> (x, mempty)) act