module skeleton8

// Camil Staps, s4498062

/**
 * Advanved Progrmming 2017, Assignment 8
 * Pieter Koopman, pieter@cs.ru.nl
 */

import StdBool
import StdEnum
from StdFunc import const, flip, id, o
import StdList
import StdString
import StdTuple

import Control.Applicative
import Control.Monad
from Control.Monad.State import :: StateT(StateT), gets, modify, runStateT,
	instance Functor (StateT s m), instance Applicative (StateT s m),
	instance Monad (StateT s m)
import Data.Bifunctor
import Data.Error
from Data.Func import $, on, `on`
import Data.Functor
import Data.List
import qualified Data.Map as M
import Data.Maybe
import qualified Data.Set as S
import Data.Tuple
from Text import <+, class Text(concat), instance Text String

from iTasks import
	class iTask, class toPrompt, class Publishable, instance Publishable (Task a),
	instance toPrompt String, instance Functor Task,
	class TApplicative, instance TApplicative Task,
	generic gEq, generic gDefault, generic JSONDecode, generic JSONEncode,
	generic gText, generic gEditor,
	:: JSONNode, :: TextFormat, :: Editor, :: Task,
	:: ViewOption(..), :: UpdateOption(..), :: EnterOption(..),
	:: Title(Title), instance toPrompt Title,
	class tune, <<@, :: UIAttributes, instance tune UIAttributes Editor,
	instance tune ArrangeHorizontal Task,
	:: ArrangeHorizontal(..), directionAttr, :: UIDirection(..),
	updateInformation, viewInformation, enterInformation, @, -||, startEngine
import qualified iTasks

(>>>=)     :== 'iTasks'.tbind

:: Expression
	= New      [Int]
	| Elem     Int
	| Variable Ident
	| Size     Set
	| (+.) infixl 6 Expression Expression
	| (-.) infixl 6 Expression Expression
	| (*.) infixl 7 Expression Expression
	| (=.) infixl 2 Ident Expression

:: Logical
	= TRUE | FALSE
	| (In) infix 4 Elem Set
	| (==.) infix 4 Expression Expression
	| (<=.) infix 4 Expression Expression
	| Not Logical
	| (||.) infixr 2 Logical Logical
	| (&&.) infixr 3 Logical Logical

:: Stmt
	= Logical Logical
	| If Logical Stmt Stmt
	| For Ident Set Stmt
	| Expression Expression

:: Set   :== Expression
:: Elem  :== Expression
:: Ident :== String

// === State

:: Val
	= VElem Int
	| VSet  ('S'.Set Int)

:: State :== 'M'.Map Ident Val

:: Sem a :== StateT State (MaybeError String) a
// Or define :: Sem a = Sem (s -> MaybeErrorString (a,s)) and copy the relevant
// instances from Control.Monad.State

// The types of store, read and fail are more general so that these functions
// can be used with the Check monad below as well.
store :: Ident v -> StateT ('M'.Map Ident v) m v | Monad m
store i v = modify ('M'.put i v) $> v

read :: Ident -> StateT ('M'.Map Ident v) (MaybeError String) v
read i = gets ('M'.get i) >>= \v -> case v of
	Nothing -> fail $ "unknown variable '" <+ i <+ "'"
	Just v  -> pure v

fail :: String -> StateT s (MaybeError String) a
fail e = StateT \_ -> Error e

class eval f t :: f -> Sem t

evalE :: (Expression -> Sem Val)
evalE = eval

instance eval Expression Val
where
	eval :: Expression -> Sem Val
	eval (New xs) = pure $ VSet $ 'S'.fromList xs
	eval (Elem x) = pure $ VElem x
	eval (Variable i) = read i
	eval (Size s)     = eval s >>= \xs -> case xs of
		VSet xs -> pure $ VElem $ 'S'.size xs
		_       -> fail "Cannot apply Size to Elem"
	eval (a +. b) = (liftA2 valAdd `on` eval) a b
	where
		valAdd (VElem i) (VElem j) = VElem $ i + j
		valAdd (VElem i) (VSet xs) = VSet $ 'S'.insert i xs
		valAdd (VSet xs) (VElem i) = VSet $ 'S'.insert i xs
		valAdd (VSet xs) (VSet ys) = VSet $ 'S'.union xs ys
	eval (a -. b) = (liftA2 tuple `on` eval) a b >>= uncurry valSub
	where
		valSub (VElem i) (VElem j) = pure $ VElem $ i - j
		valSub (VElem i) (VSet xs) = fail "Cannot subtract Set from Elem"
		valSub (VSet xs) (VElem i) = pure $ VSet $ 'S'.delete i xs
		valSub (VSet xs) (VSet ys) = pure $ VSet $ 'S'.difference xs ys
	eval (a *. b) = (liftA2 tuple `on` eval) a b >>= uncurry valMul
	where
		valMul (VElem i) (VElem j) = pure $ VElem $ i * j
		valMul (VElem i) (VSet xs) = pure $ VSet  $ 'S'.mapSet ((*) i) xs
		valMul (VSet xs) (VElem i) = fail "Cannot multiply Elem with Set"
		valMul (VSet xs) (VSet ys) = pure $ VSet  $ 'S'.intersection xs ys
	eval (n =. e) = eval e >>= store n

// === semantics

evalB :: (Logical -> Sem Bool)
evalB = eval

instance eval Logical Bool
where
	eval :: Logical -> Sem Bool
	eval TRUE = pure True
	eval FALSE = pure False
	eval (e In s) = (liftA2 tuple `on` eval) e s >>= \t -> case t of
		(VElem i, VSet xs) -> pure $ 'S'.member i xs
		_                  -> fail "Can only apply In to Elem and Set"
	eval (a ==. b) = (liftA2 tuple `on` eval) a b >>= \t -> case t of
		(VElem i, VElem j) -> pure $ i == j
		(VSet xs, VSet ys) -> pure $ 'S'.isSubsetOf xs ys && 'S'.size xs == 'S'.size ys
		                      // Cannot use == due to name clash
		_                  -> fail "Cannot apply == to Elem and Set"
	eval (a <=. b) = (liftA2 tuple `on` eval) a b >>= \t -> case t of
		(VElem i, VElem j) -> pure $ i <= j
		(VSet xs, VSet ys) -> pure $ 'S'.isSubsetOf xs ys
		_                  -> fail "Cannot apply <= to Elem and Set"
	eval (Not b) = not <$> eval b
	eval (a ||. b) = (liftA2 (||) `on` eval) a b
	eval (a &&. b) = (liftA2 (&&) `on` eval) a b

:: Result
	= RElem Int
	| RSet  ('S'.Set Int)
	| RBool Bool

evalS :: (Stmt -> Sem Result)
evalS = eval

/**
 * For `For`, I use the following semantics:
 *  - `For x [e1,e2,...] body`
 *    Consecutively assigns e1, e2, ... to x and evaluates body to an element.
 *    The elements are placed in a set, which is the result.
 *  - `For x n body`
 *    Sugar for `For x [0..n-1] body`, or an error if n < 0.
 */
instance eval Stmt Result
where
	eval :: Stmt -> Sem Result
	eval (If b t e) = eval b >>= \b -> eval $ if b t e
	eval (For v s e) = eval s >>= \s -> case s of
		VElem n -> if (n < 0)
			(fail "Cannot iterate until a negative value")
			(eval $ For v (New [0..n-1]) e)
		VSet xs -> sequence [store v (VElem x) >>| eval e \\ x <- xs`] >>= \xs ->
			if (all (\x -> x=:(RElem _)) xs)
				(pure $ RSet $ 'S'.fromList [x \\ RElem x <- xs])
				(fail "Not all results of For loop were Elems")
			with xs` = 'S'.toList xs
	eval (Expression e) = vtor <$> eval e
	where
		vtor (VElem x) = RElem x
		vtor (VSet xs) = RSet xs
	eval (Logical l) = RBool <$> eval l

// === printing

class print a :: a [String] -> [String]

printToString :: (a -> String) | print a
printToString = concat o flip print []

instance print Int where print i st = [toString i:st]
instance print String where print s st = [s:st]

instance print Expression
where
	print (New xs) st = ["[":intersperse "," (map toString xs)] ++ ["]":st]
	print (Elem x) st = print x st
	print (Variable i) st = print i st
	print (Size s) st = ["(Size ":"(":print s ["))":st]]
	print (a +. b) st = ["(":print a ["+":print b [")":st]]]
	print (a -. b) st = ["(":print a ["-":print b [")":st]]]
	print (a *. b) st = ["(":print a ["*":print b [")":st]]]
	print (v =. e) st = print v [" = ":print e st]

instance print Logical
where
	print TRUE st = ["TRUE":st]
	print FALSE st = ["FALSE":st]
	print (x In xs) st = ["(":print x [" In ":print xs [")":st]]]
	print (a ==. b) st = ["(":print a ["==":print b  [")":st]]]
	print (a <=. b) st = ["(":print a ["<=":print b  [")":st]]]
	print (Not b)   st = ["~(":print b [")":st]]
	print (a ||. b) st = ["(":print a ["||":print b  [")":st]]]
	print (a &&. b) st = ["(":print a ["&&":print b  [")":st]]]

instance print Stmt
where
	print (If b t e) st = ["If ":print b [" {":print t ["} {":print e ["}":st]]]]
	print (For v s e) st = ["For ":print v [" {":print s ["} {":print e ["}":st]]]]
	print (Expression e) st = ["Expression (":print e [")":st]]
	print (Logical l) st = ["Logical (":print l [")":st]]

// === simulation

derive class iTask Expression, Logical, Stmt, Val, Result, 'S'.Set

simulate :: Stmt -> Task Stmt
simulate stmt =
	updateInformation (Title "Program") [UpdateUsing id (flip const) stmtEditor] stmt
	-|| 'iTasks'.allTasks
	[ viewInformation (Title "String representation") [ViewAs printToString] stmt
	, viewInformation (Title "Type")                  [ViewAs showtype]      stmt
	, viewInformation (Title "Execution")             [ViewAs execute]       stmt
	] <<@ ArrangeHorizontal
where
	execute = fmap (bifmap showResult (map (appSnd showVal) o 'M'.toList))
		o flip runStateT 'M'.newMap o evalS
	showtype = fmap fst o flip runStateT 'M'.newMap o type

	showVal :: Val -> String
	showVal (VElem i) = toString i
	showVal (VSet s)  = concat (intersperse ", " [toString i \\ i <- 'S'.toList s])

	showResult :: Result -> String
	showResult (RElem i) = toString i
	showResult (RSet s)  = concat (intersperse ", " [toString i \\ i <- 'S'.toList s])
	showResult (RBool b) = toString b

// NB: I want the input fields to be put vertically but don't see how that is
// possible; neither this nor this with Vertical nor
// <<@ Arrange{Horizontal,Vertical} on the task works.
stmtEditor :: Editor Stmt
stmtEditor = gEditor{|*|} <<@ directionAttr Horizontal

// === type checking

:: Type
	= TElem
	| TSet
	| TBool

:: Check a :== StateT ('M'.Map Ident Type) (MaybeError String) a

derive class iTask Type

class type a :: a -> Check Type

instance type Expression
where
	type (New _) = pure TSet
	type (Elem _) = pure TElem
	type (Variable i) = read i
	type (Size e) = type e >>= \e -> case e of
		TSet -> pure TElem
		_    -> fail "Can only apply Size to Set"
	type (a +. b) = (liftA2 tuple `on` type) a b >>= \ts -> case ts of
		(TElem, TElem) -> pure TElem
		(TBool, _)     -> fail "Cannot apply + to a boolean"
		(_,     TBool) -> fail "Cannot apply + to a boolean"
		_              -> pure TSet
	type (a -. b) = (liftA2 tuple `on` type) a b >>= \ts -> case ts of
		(TElem, TElem) -> pure TElem
		(TBool, _)     -> fail "Cannot apply - to a boolean"
		(_,     TBool) -> fail "Cannot apply - to a boolean"
		(TElem, TSet)  -> fail "Cannot apply - to Elem and Set"
		_              -> pure TSet
	type (a *. b) = (liftA2 tuple `on` type) a b >>= \ts -> case ts of
		(TElem, TElem) -> pure TElem
		(TBool, _)     -> fail "Cannot apply * to a boolean"
		(_,     TBool) -> fail "Cannot apply * to a boolean"
		(TSet,  TElem) -> fail "Cannot apply * to Set and Elem"
		_              -> pure TSet
	type (v =. e) = type e >>= \t -> case t of
		TBool -> fail "Cannot assign a boolean to a variable"
		_     -> modify ('M'.put v t) $> t

instance type Logical
where
	type TRUE = pure TBool
	type FALSE = pure TBool
	type (x In xs) = (liftA2 tuple `on` type) x xs >>= \ts -> case ts of
		(TElem, TSet) -> pure TBool
		_             -> fail "Can only apply In to Elem and Set"
	type (a ==. b) = (liftA2 tuple `on` type) a b >>= \ts -> case ts of
		(TElem, TElem) -> pure TBool
		(TSet,  TSet)  -> pure TBool
		_              -> fail "Cannot apply == to Elem and Set"
	type (a <=. b) = (liftA2 tuple `on` type) a b >>= \ts -> case ts of
		(TElem, TElem) -> pure TBool
		(TSet,  TSet)  -> pure TBool
		_              -> fail "Cannot apply <= to Elem and Set"
	type (Not b) = type b
	type (a ||. b) = type a >>| type b
	type (a &&. b) = type a >>| type b

/**
 * The cases for If and For are tricky. To type-check, we need the then and
 * else block to have the same type. To run, this is not necessary, as long as
 * the type of the used block is correct.
 * We require the For block to have type Elem here, but at runtime it's OK if
 * it has another type as long as the block is not used (the iterator is 0 or
 * empty).
 */
instance type Stmt
where
	type (If b t e) = type b >>| type t >>= \t -> type e >>= \e ->
		if (t === e) (pure t) (fail "Types of then and else blocks must match")
	type (For v s e) = type s >>| store v TElem >>| type e >>= \e -> case e of
		TElem -> pure TSet
		_     -> fail "Results of For loop should be Elem"
	type (Expression e) = type e
	type (Logical l) = type l

/*
// === Testing the type checker
// Because of the reason explained at the instance type Stmt, we cannot check
// whether type-check-ok iff run-ok. So we check that every program that can be
// type-checked correctly can be run correctly instead.
from Gast import
	quietnm, aStream, generic ggen, generic genShow,
	class TestArg, class Testable, instance Testable (a -> Bool), instance Testable Bool,
	:: GenState, :: RandomStream

derive ggen Stmt, Expression, Logical
derive genShow Stmt, Expression, Logical
derive bimap []

test :: Stmt -> (MaybeErrorString Result, MaybeErrorString Type)
test stm = (fst <$> runStateT (evalS stm) 'M'.newMap, fst <$> runStateT (type stm) 'M'.newMap)

check :: Stmt -> Bool
check stm = isOk result || isError type
where (result,type) = test stm

Start _ = quietnm 100000 20 aStream check
*/

Start w = startEngine (enterInformation (Title "Program") [EnterUsing id stmtEditor] >>>= sim) w
where
	sim :: Stmt -> Task Stmt
	sim st = simulate st >>>= sim