module sets

import StdBool
import StdEnum
from StdFunc import const, flip, id, o
import StdList
import StdOrdList
import StdOverloaded
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.List as L
from Data.List import instance Functor []
import qualified Data.Map as M
import Data.Maybe
from Text import <+, class Text(concat), instance Text String

// It's been a long time since I worked with dynamics so I thought it would be
// nice to do that again.
:: SetState :== 'M'.Map Ident Dynamic
:: Ident :== String

// Like in the previous assignment: we can also use something like
// :: Sem a = Sem (SetState -> MaybeErrorString (a, SetState))
// and implement the monad functionality like that of StateT.
:: Sem a :== StateT SetState (MaybeError String) a

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

store :: Ident a -> Sem a | TC, toString a
store id x = modify ('M'.put id (dynamic x)) $> x

read :: Ident String -> Sem a | TC, toString a
read id type = gets ('M'.get id) >>= \v -> case v of
	Just (v :: a^) -> pure v
	Just d  -> fail $ "Expected " <+ type <+ " for '" <+ id <+ "'; got " <+ typeCodeOfDynamic d
	Nothing -> fail $ "Unknown variable '" <+ id <+ "'"

eval :: ((Sem a) -> MaybeError String a)
eval = flip evalStateT 'M'.newMap

:: Element   :== Sem Int
:: Set       :== Sem [Int]
:: Logical   :== Sem Bool

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

// Convenience
INT    :: (Element -> Element); INT = id
SET    :: (Set -> Set);         SET = id
BOOL   :: (Logical -> Logical); BOOL = id
PRINT  :: (Print -> Print);     PRINT = id
intvar :: (String -> Element);  intvar = variable
setvar :: (String -> Set);      setvar = variable

// -- Integer expressions --

class integer a :: (Int -> a)
class set a :: ([Int] -> a)
class size a b :: (a -> b)

class (+:) infixl 6 a b :: (a b -> b)
class (:+) infixl 6 a b :: (a b -> a)
class (:-) infixl 6 a b :: (a b -> a)
class (*:) infixl 7 a b :: (a b -> b)

instance integer Element where integer = pure
instance set Set where set = pure
instance size Set Element where size = fmap length

instance + Element where + a b = liftA2 (+) a b
instance - Element where - a b = liftA2 (-) a b
instance * Element where * a b = liftA2 (*) a b

instance + Set where + a b = liftA2 'L'.union      a b
instance - Set where - a b = liftA2 'L'.difference a b
instance * Set where * a b = liftA2 'L'.intersect  a b

instance +: Element Set where +: = (+) o fmap pure
instance :+ Set Element where :+ = flip (+:)
instance :- Set Element where :- = flip $ flip (-) o fmap pure
instance *: Element Set where *: = liftA2 $ map o (*)

instance integer Print where integer = \i st -> ["integer ":toString i:st]
instance set Print where set = \xs st -> ["set [":'L'.intersperse "," (map toString xs)] ++ ["]":st]
instance size Print Print where size = \s st -> ["size (":s [")":st]]

instance + Print where + a b = \st -> ["(":a [" + ":b [")":st]]]
instance - Print where - a b = \st -> ["(":a [" - ":b [")":st]]]
instance * Print where * a b = \st -> ["(":a [" * ":b [")":st]]]

instance +: Print Print where +: = \a b st -> ["(":a [" +: ":b [")":st]]]
instance :+ Print Print where :+ = \a b st -> ["(":a [" :+ ":b [")":st]]]
instance :- Print Print where :- = \a b st -> ["(":a [" :- ":b [")":st]]]
instance *: Print Print where *: = \a b st -> ["(":a [" *: ":b [")":st]]]

// -- Set expressions --

class Variable a
where
	variable :: (String -> a)
	(=.) infix 2 :: (Ident a -> a)

instance Variable Element
where
	variable = flip read "Int"
	=. = (=<<) o store

instance Variable Set
where
	variable = flip read "[Int]"
	=. = (=<<) o store

instance Variable Print
where
	variable = \id st -> [id:st]
	=. = \id x st -> [id:" =. ":x st]

// -- Logical expressions --

class Logic a
where
	true :: a
	false :: a
	Not :: (a -> a)
	(||.) infixr 2 :: (a a -> a)
	(&&.) infixr 3 :: (a a -> a)

class Compare a b
where
	(==.) infix 4 :: (a a -> b)
	(<=.) infix 4 :: (a a -> b)

class (In) infix a b c :: (a b -> c)

instance Logic Logical
where
	true = pure True
	false = pure False
	Not = fmap not
	(||.) = liftA2 (||)
	(&&.) = liftA2 (&&)

instance Compare Element Logical
where
	==. = liftA2 (==)
	<=. = liftA2 (<=)

instance Compare Set Logical
where
	==. = liftA2 ((==) `on` sort)
	<=. = liftA2 (flip (all o flip isMember))

instance In Element Set Logical where In = liftA2 isMember

instance Logic Print
where
	true = \st -> ["true":st]
	false = \st -> ["false":st]
	Not = \b st -> ["not (":b [")":st]]
	(||.) = \a b st -> ["(":a [" ||. ": b [")":st]]]
	(&&.) = \a b st -> ["(":a [" &&. ": b [")":st]]]

instance Compare Print Print
where
	==. = \a b st -> a [" ==. ":b st]
	<=. = \a b st -> a [" <=. ":b st]

instance In Print Print Print where In = \e s st -> e [" in ":s st]

// -- Statements --
// The Expression and Logical from last week are not needed; we can just use
// Sem Int, Sem [Int] and Sem Bool in the same pipeline.

class If a b :: a b b -> b
class for a b c :: Ident a b -> c

instance If Logical (Sem a) where If b t e = b >>= \b -> if b t e
instance for Set     Element Set where for k xs body = xs >>= mapM (flip (>>|) body o store k)
instance for Element Element Set where for k n body = n >>= \n -> for k (SET (set [0..n-1])) body

instance If Print Print where If b t e = \st -> ["if (":b [") {":t ["} else {":e ["}":st]]]]
instance for Print Print Print where for x l b = \st -> ["for ":x:" in ":l [" {":b ["}":st]]]

(:.) infixl 1
(:.) :== (>>|)

print :: ([String] -> [String]) -> String
print f = concat (f [])

//Start =
//	( "\n", print $ integer 5
//	, "\n", print $ set [1,2,3,4,5]
//	, "\n", print $ size (PRINT (set [1..10]))
//	, "\n", print $ PRINT (set [1..3]) +: PRINT (integer 5 * integer 10)
//	, "\n", print $ "x" =. integer 5
//	, "\n", print $ factorial 5
//	, "\n"
//	)

Start =
	( "\n", eval $ SET (for "x" (SET (set [0,1,2,3,4])) (variable "x" * INT (integer 5)))
	, "\n", eval $ "y" =. INT (integer 5) :. SET (for "x" (intvar "y") (variable "x" * INT (integer 5)))
	//, "\n", eval $ factorial 5
	, "\n")

//factorial :: Int -> a b
//factorial n =
//	"k" =. integer n :.
//	"r" =. integer 1 :.
//	for "x" (variable "k")
//		(If (variable "x" ==. integer 0)
//			(variable "r") // for type-checking
//			("r" =. variable "r" * variable "x")) :.
//	variable "r"