module sets

// Laurens Kuijper, s4467299
// Camil Staps, s4498062

import StdArray
import StdBool
import StdEnum
from StdFunc import flip, id, o
import StdGeneric
import StdInt
import StdList
import StdString

import Control.Applicative
import Control.Monad
import Control.Monad.State
import Data.Error
from Data.Func import $, on, `on`
import Data.Functor
import qualified Data.Map as M
import Data.List
import Data.Maybe
from Text import <+, class Text(concat), instance Text String

instance + [a] | Eq a where + xs ys = xs - ys ++ ys
instance - [a] | Eq a where - xs ys = removeMembers xs ys
instance * [a] | Eq a where * xs ys = [x \\ x <- xs | hasElem x ys]

addElem :: a [a] -> [a] | Eq a
addElem x xs = if (isMember x xs) xs [x:xs]

removeElem :: (a [a] -> [a]) | Eq a
removeElem = removeMember

hasElem :: (a [a] -> Bool) | Eq a
hasElem = isMember

/**
 * This Expression type is agnostic with regards to the expression type, so it
 * can handle sets of booleans, sets of sets, reals, etc. Because of this, we
 * only need Elem and not New and TRUE/FALSE.
 */
:: Expression a
	=        Elem                           a
	|        Variable                       Ident
	| E.b:   Size           (Bimap a Int)   (Set b)                       & TC, Print b
	|        (+.)  infixl 6                 (Expression a) (Expression a) & + a
	| E.b:   AddElemAndSet  (Bimap a [b])   (Elem b)       (Set b)        & TC, Eq, Print b
	| E.b:   AddSetAndElem  (Bimap a [b])   (Set b)        (Elem b)       & TC, Eq, Print b
	|        (-.)  infixl 6                 (Expression a) (Expression a) & - a
	| E.b:   RemoveFromSet  (Bimap a [b])   (Set b)        (Elem b)       & TC, Eq, Print b
	|        (*.)  infixl 7                 (Expression a) (Expression a) & * a
	| E.b:   MulElemAndSet  (Bimap a [b])   (Elem b)       (Set b)        & TC, *, Print b
	|        (=.)  infixl 2                 Ident          (Expression a)

	|        Not            (Bimap a Bool)  (Elem Bool)
	|        Or             (Bimap a Bool)  (Elem Bool)    (Elem Bool)
	|        And            (Bimap a Bool)  (Elem Bool)    (Elem Bool)
	| E.b:   In             (Bimap a Bool)  (Elem b)       (Set b)        & TC, Eq, Print b
	| E.b:   Eq             (Bimap a Bool)  (Expression b) (Expression b) & TC, Eq, Print b
	| E.b:   Le             (Bimap a Bool)  (Expression b) (Expression b) & TC, Ord, Print b

	|        If                             (Elem Bool) (Expression a) (Expression a)
	| E.b c: For            (Bimap a [b])   Ident (Set c) (Expression b)  & TC, Print b & TC, Print c
	| E.b:   (:.) infixl 1                  (Expression b) (Expression a) & TC, Print b

// Convenience
elem :== Elem
var  :== Variable
size :== Size bimapId
for  :== For bimapId

(+:)  infixl 6; (+:)  :== AddSetAndElem bimapId
(:+)  infixl 6; (:+)  :== AddElemAndSet bimapId
(-:)  infixl 6; (-:)  :== RemoveFromSet bimapId
(:*)  infixl 7; (:*)  :== MulElemAndSet bimapId
(==.) infix  4; (==.) :== Eq  bimapId
(<=.) infix  4; (<=.) :== Le  bimapId
(||.) infixr 2; (||.) :== Or  bimapId
(&&.) infixr 3; (&&.) :== And bimapId

:: SetState :== 'M'.Map Ident Dynamic
:: Sem a :== StateT SetState (MaybeError String) a

:: Set a :== Expression [a]
:: Elem a :== Expression a
:: Ident :== String

fail :: (String -> StateT s (MaybeError String) a)
fail = StateT o const o Error

store :: Ident v -> StateT SetState m v | Monad m & TC v
store i v = modify ('M'.put i (dynamic v)) $> v

read :: Ident -> StateT SetState (MaybeError String) v | TC v
read i = gets ('M'.get i) >>= \v -> case v of
	Just (x :: v^) -> pure x
	Just d         -> fail $ "type error, " <+ typeCodeOfDynamic d <+ " for '" <+ i <+ "'"
	Nothing        -> fail $ "unknown variable '" <+ i <+ "'"

eval :: (Expression a) -> Sem a | TC a
eval (Elem x)         = pure x
eval (Variable id)    = read id
eval (Size bm s)      = bm.map_from <$> length <$> eval s
eval (x +. y)         = (liftA2 (+) `on` eval) x y
eval (AddElemAndSet bm e s) = bm.map_from <$> liftA2 addElem (eval e) (eval s)
eval (AddSetAndElem bm s e) = bm.map_from <$> liftA2 addElem (eval e) (eval s)
eval (x -. y)         = (liftA2 (-) `on` eval) x y
eval (RemoveFromSet bm s e) = bm.map_from <$> liftA2 removeElem (eval e) (eval s)
eval (x *. y)         = (liftA2 (*) `on` eval) x y
eval (MulElemAndSet bm e s) = bm.map_from <$> liftA2 (map o (*)) (eval e) (eval s)
eval (id =. e)        = eval e >>= store id
eval (Not bm l)       = bm.map_from <$> not <$> eval l
eval (Or bm x y)      = bm.map_from <$> (liftA2 (||) `on` eval) x y
eval (And bm x y)     = bm.map_from <$> (liftA2 (&&) `on` eval) x y
eval (In bm e s)      = bm.map_from <$> liftA2 hasElem (eval e) (eval s)
eval (Eq bm x y)      = bm.map_from <$> liftA2 (==) (eval x) (eval y)
eval (Le bm x y)      = bm.map_from <$> liftA2 (<=) (eval x) (eval y)
eval (If b t e)       = eval b >>= \b -> eval $ if b t e
eval (For bm id s e)  = bm.map_from <$> (eval s >>= iterate e)
where
	iterate :: (Elem b) -> [a] -> Sem [b] | TC a & TC b
	iterate e = mapM (\x -> store id x >>| eval e)
eval (s :. t)         = eval s >>| eval t

:: Print :== [String] -> [String]

class Print a where pr :: a -> Print

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

between :: a b c -> Print | Print a & Print b & Print c
between a b c = pr b o pr a o pr c

surround :: a b c -> Print | Print a & Print b & Print c
surround a b c = pr a o pr c o pr b

parens :: (a -> Print) | Print a
parens = surround "(" ")"

braces :: (a -> Print) | Print a
braces = surround "{" "}"

interpr :: a [b] -> Print | Print a & Print b
interpr _ [] = id
interpr _ [x] = pr x
interpr g [x:xs] = between g x (interpr g xs)

instance Print String where pr s = (++) (pure s)
instance Print Char where pr c = pr {#c}
instance Print Int where pr i = pr (toString i)
instance Print Bool where pr b = pr (toString b)
instance Print [a] | Print a where pr xs = surround "[" "]" (interpr "," xs)
instance Print Print where pr p = p

instance Print (Expression a) | Print a
where
	pr (Elem x)              = pr x
	pr (Variable x)          = surround '"' '"' x
	pr (Size _ xs)           = surround "size (" ")" xs
	pr (x +. y)              = parens (between " +. " x y)
	pr (AddElemAndSet _ e s) = parens (between " :+ " e s)
	pr (AddSetAndElem _ s e) = parens (between " +: " s e)
	pr (x -. y)              = parens (between " -. " x y)
	pr (RemoveFromSet _ s e) = parens (between " -: " s e)
	pr (x *. y)              = parens (between " *. " x y)
	pr (MulElemAndSet _ e s) = parens (between " :* " e s)
	pr (id =. e)             = between " =. " id e
	pr (Not _ l)             = pr "~" o parens l
	pr (Or _ x y)            = parens (between " ||. " x y)
	pr (And _ x y)           = parens (between " &&. " x y)
	pr (In _ e s)            = parens (between " ∈ "   e s)
	pr (Eq _ x y)            = parens (between " == "  x y)
	pr (Le _ x y)            = parens (between " <= "  x y)
	pr (If b t e)            = pr "if " o pr b o pr " " o braces t o pr " " o braces e
	pr (For _ id s e)        = pr "for " o between " ∈ " id s o pr " " o braces e
	pr (s :. t)              = between " :. " s t

/**
 * We did not manage to get the iTasks simulator working, because the compiler
 * 'cannot build a generic representation of an existential type'.
 * We're sure it would be *possible* by implementing an Editor manually,
 * forcing all Bimaps to bimapId and using cast :: a -> b in case of type
 * errors, but did not try to do this due to time constraints.
 */

Start = (evalStateT (eval stmt) 'M'.newMap, print stmt)
where
	stmt :: Expression [Int]
	stmt =
		"x" =. elem [3..10] :.
		"y" =. elem 0 :.
		elem 10 :* for "i" (elem [1..5] +. var "x") (
			"y" =. var "y" +. elem 1 :.
			var "i" *. var "y"
		)