-- Copyright 2022 University of Freiburg
-- Janek Spaderna <janek.spaderna@pluto.uni-freiburg.de>
module Ex08 where

import Control.Category ((>>>))
import Control.Monad.Trans.State.Strict
import qualified Data.Map as Map
import Data.Maybe
import Test.QuickCheck

-------------------------------------------------------------------------------
-- Exercise 1 (template rendering)

data Template
  = -- | A literal string.
    T String
  | -- | A variable. The template resolves to the variables name.
    V Template
  | -- | Include a subtemplate with a set of new bindings. The name of the
    -- subtemplate is a template itself.
    I Template [Definition]
  | -- | Concatenation of multiple templates.
    C [Template]

data Definition = D Template Template

data Environment = Environment
  { envTemplates :: Map.Map String Template,
    envVariables :: Map.Map String String
  }

lookupVar :: String -> Environment -> Maybe String
lookupVar s = envVariables >>> Map.lookup s

lookupTemplate :: String -> Environment -> Maybe Template
lookupTemplate s = envTemplates >>> Map.lookup s

addDefs :: [(String, String)] -> Environment -> Environment
addDefs defs env = env {envVariables = Map.fromList defs <> envVariables env}

local :: (e -> e') -> (e' -> a) -> (e -> a)
local upd f a = f (upd a)

resolveDef :: Definition -> Environment -> (String, String)
resolveDef (D nameTmpl valueTmpl) = do
  name <- resolve nameTmpl
  value <- resolve valueTmpl
  return (name, value)

resolve :: Template -> Environment -> String
resolve (T s) = return s
resolve (V t) = do
  v <- resolve t
  s <- lookupVar v
  return $ "" `fromMaybe` s
resolve (I t defs) = do
  tname <- resolve t
  inner <- lookupTemplate tname
  resolvedDefs <- mapM resolveDef defs
  local (addDefs resolvedDefs) $ do
    maybe (return "") resolve inner
resolve (C ts) = do
  ss <- mapM resolve ts
  return $ concat ss

-------------------------------------------------------------------------------
-- Exercise 2 (joining monads)

join :: Monad m => m (m a) -> m a
join mma = mma >>= id

-- equivalent --

join' :: Monad m => m (m a) -> m a
join' mma = do
  ma <- mma
  ma

{-

m (m a) -> m a                        -- substitute ((->) r) for m

((->) r ((->) r a)) -> ((->) r a)     -- turn prefix form into infix operators

(r -> (r -> a)) -> (r -> a)           -- use right-associativity of (->)

(r -> r -> a) -> r -> a

-}

prop_FunctionJoinUsesArgTwice ::
  Fun (Integer, Integer) Integer -> Integer -> Bool
prop_FunctionJoinUsesArgTwice f x =
  join (applyFun2 f) x == applyFun2 f x x

{-

join (.) :: (a -> a) -> a -> a

-}

prop_JoininigComposeAppliesTwice ::
  Fun Integer Integer -> Integer -> Bool
prop_JoininigComposeAppliesTwice f x =
  join (.) (applyFun f) x == applyFun f (applyFun f x)

-------------------------------------------------------------------------------
-- Exercise 3 (random number generator)

type Random a = State Integer a

fresh :: Random Integer
fresh = do
  xprev <- get
  let xnext = (6364136223846793005 * xprev + 1442695040888963407) `mod` 2 ^ (64 :: Integer)
  put xnext
  return xnext

runPRNG :: Random a -> Integer -> a
runPRNG = evalState