module sets

import StdBool
import StdEnum
from StdFunc import flip, 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.Maybe
from Text import <+

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

:: Expression a
	=      New            (Bimap a [Int]) [Int]
	|      Elem           (Bimap a Int)   Int
	| E.b: Variable       (Bimap a b)     Ident
	|      Size           (Bimap a Int)   Set
	|      (+.)  infixl 6                 (Expression a) (Expression a) & + a
	|      AddElemAndSet  (Bimap a [Int]) Elem           Set
	|      AddSetAndElem  (Bimap a [Int]) Set            Elem
	|      (-.)  infixl 6                 (Expression a) (Expression a) & - a
	|      RemoveFromSet  (Bimap a [Int]) Set            Elem
	|      (*.)  infixl 7                 (Expression a) (Expression a) & * a
	|      MulElemAndSet  (Bimap a [Int]) Elem           Set
	|      (=.)  infixl 2                 Ident          (Expression a)

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

	|      If                             Logical (Expression a) (Expression a)
	|      For            (Bimap a [Int]) Ident Set (Expression Int)
	| E.b: (:.) infixl 1                  (Expression b) (Expression a) & TC b

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

(+:) infixl 6; (+:) :== AddSetAndElem bimapId
(:+) infixl 6; (:+) :== AddElemAndSet bimapId
(-:) infixl 6; (-:) :== RemoveFromSet bimapId
(:*) infixl 7; (:*) :== MulElemAndSet bimapId

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

:: Set :== Expression [Int]
:: Elem :== Expression Int
:: Logical :== Expression Bool
:: Ident :== String

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 <+ "'"

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

eval :: (Expression a) -> Sem a | TC a
eval (New bm xs)            = pure $ bm.map_from xs
eval (Elem bm x)            = pure $ bm.map_from x
eval (Variable bm 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 (Logical bm b)         = pure $ bm.map_from b
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 >>= mapM (\x -> store id x >>| eval e))
eval (s :. t)               = eval s >>| eval t

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