module sets

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

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

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

// -- Integer expressions --

integer :: (Int -> Element)
integer = pure

set :: ([Int] -> Set)
set = pure

size :: (Set -> Element)
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

(+:) infixl 6 :: (Element Set -> Set); (+:) = (+) o fmap pure
(:+) infixl 6 :: (Set Element -> Set); (:+) = flip (+:)

(:-) infixl 6 :: (Set Element -> Set); (:-) = flip $ flip (-) o fmap pure
(*:) infixl 6 :: (Element Set -> Set); (*:) = liftA2 $ map o (*)

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

// -- Set expressions --

class variable a :: (String -> a)

instance variable Element where variable = flip read "Int"
instance variable Set     where variable = flip read "[Int]"

// Convenience
intvar :: (String -> Element); intvar = variable
setvar :: (String -> Set);     setvar = variable

class (=.) infix 2 a :: (Ident a -> a)

instance =. Element where =. = (=<<) o store
instance =. Set     where =. = (=<<) o store

// -- Logical expressions --

true :: Logical
true = pure True

false :: Logical
false = pure False

(In) infix :: (Element Set -> Logical)
(In) = liftA2 isMember

class (==.) infix 4 a :: (a a -> Logical)
instance ==. Element where ==. = liftA2 (==)
instance ==. Set     where ==. = liftA2 ((==) `on` sort)

class (<=.) infix 4 a :: (a a -> Logical)
instance <=. Element where <=. = liftA2 (<=)
instance <=. Set     where <=. = liftA2 (flip (all o flip isMember))

Not :: (Logical -> Logical)
Not = fmap not

(||.) infixr 2 :: (Logical Logical -> Logical)
(||.) = liftA2 (||)

(&&.) infixr 3 :: (Logical Logical -> Logical)
(&&.) = liftA2 (&&)

// -- 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.

If :: Logical (Sem a) (Sem a) -> Sem a
If b t e = b >>= \b -> if b t e

class for a :: Ident a Element -> Set
instance for Set where for k xs body = xs >>= mapM (flip (>>|) body o store k)
instance for Element where for k n body = n >>= \n -> for k (set [0..n-1]) body

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

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

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