{-# LANGUAGE InstanceSigs #-}

module Lecture09 where

import Data.Map ( Map )
import qualified Data.Map as Map
import Control.Applicative

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

evala' :: AExpr -> Assignment -> Integer
evala' (Con x) _   = x
evala' (Var x) a   = val' x a
evala' (Add x y) a = evala' x a + evala' y a
evala' (Mul x y) a = evala' x a * evala' y a

newtype Reader r a = Reader { runReader :: r -> a }

ask :: Reader r r
ask = Reader id

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

instance Functor (Reader r) where
    fmap :: (a -> b) -> Reader r a -> Reader r b
    fmap f x = Reader $ f . runReader x

instance Applicative (Reader r) where
    (<*>) :: Reader r (a -> b) -> Reader r a -> Reader r b
    -- runReader f :: r -> (a -> b); runReader x :: r -> a
    f <*> x = Reader $ \r -> runReader f r (runReader x r)
    pure x  = Reader $ \_ -> x

instance Monad (Reader r) where
    (>>=) :: Reader r a -> (a -> Reader r b) -> Reader r b
    -- runReader x :: r -> a; runReader . f :: a -> r -> b
    x >>= f = Reader $ \r -> (runReader . f) (runReader x r) r

val :: String -> Reader Assignment Integer
val v = reader (Map.findWithDefault 0 v)

evala :: AExpr -> Reader Assignment Integer
evala (Con x)   = pure x
evala (Var x)   = val x
evala (Add x y) = liftA2 (+) (evala x) (evala y)
evala (Mul x y) = liftA2 (*) (evala x) (evala y)

evala'' :: AExpr -> Assignment -> Integer
evala'' (Con x)   = pure x
evala'' (Var x)   = val' x
evala'' (Add x y) = liftA2 (+) (evala'' x) (evala'' y)
evala'' (Mul x y) = liftA2 (*) (evala'' x) (evala'' y)

local :: (r -> r) -> Reader r a -> Reader r a
local f act = Reader $ \r -> runReader act (f r)

ex_local :: (Integer, Integer, Integer)
ex_local = runReader (liftA3 (,,) (evala (Var "x"))
                                  (local (Map.insert "x" 42) (evala (Var "x")))
                                  (evala (Var "x")))
                     Map.empty

newtype ReaderIO r a = ReaderIO { runReaderIO :: r -> IO a }

readerIO :: (r -> a) -> ReaderIO r a
readerIO f = ReaderIO (pure . f)

askIO :: ReaderIO r r
askIO = readerIO id

instance Functor (ReaderIO r) where
    fmap :: (a -> b) -> ReaderIO r a -> ReaderIO r b
    fmap f x = ReaderIO $ \r -> fmap f (runReaderIO x r)

instance Applicative (ReaderIO r) where
    pure x = ReaderIO (\_ -> pure x)
    (<*>) :: ReaderIO r (a -> b) -> ReaderIO r a -> ReaderIO r b
    -- runReaderIO f :: r -> IO (a -> b)
    -- runReaderIO x :: r -> IO a
    f <*> x = ReaderIO $ \r -> runReaderIO f r <*> runReaderIO x r

instance Monad (ReaderIO r) where
    (>>=) :: ReaderIO r a -> (a -> ReaderIO r b) -> ReaderIO r b
    -- runReaderIO x   :: r -> IO a
    x >>= f = ReaderIO $
              \r -> runReaderIO x r >>= \xx ->
                    -- xx :: a
                    runReaderIO (f xx) r

newtype ReaderT r m a = ReaderT { runReaderT :: r -> m a }

readerT :: Applicative m => (r -> a) -> ReaderT r m a
readerT f = ReaderT (pure . f)

askT :: Applicative m => ReaderT r m r
askT = readerT id

instance Functor m => Functor (ReaderT r m) where
    fmap :: (a -> b) -> ReaderT r m a -> ReaderT r m b
    fmap f x = ReaderT $ fmap f . runReaderT x

instance Applicative m => Applicative (ReaderT r m) where
    pure x = ReaderT (\_ -> pure x)

    (<*>) :: ReaderT r m (a -> b) -> ReaderT r m a -> ReaderT r m b
    -- runReaderT f :: r -> m (a -> b)
    -- runReaderT x :: r -> m a
    f <*> x = ReaderT $ \r -> runReaderT f r <*> runReaderT x r

instance Monad m => Monad (ReaderT r m) where
    (>>=) :: ReaderT r m a -> (a -> ReaderT r m b) -> ReaderT r m b
    -- runReaderT x   :: r -> m a
    x >>= f = ReaderT $
              \r -> runReaderT x r >>= \xx ->
                    -- xx :: a
                    runReaderT (f xx) r
    fail msg = ReaderT $ \_ -> fail msg

class MonadTrans t where
    lift :: Monad m => m a -> t m a

instance MonadTrans (ReaderT r) where
    lift :: Monad m => m a -> ReaderT r m a
    lift act = ReaderT $ \_ -> act