module sets

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

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

:: 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]

pr :: a -> Print | toString a
pr x = \st -> [toString x:st]

prsperse :: a [b] -> Print | toString a & toString b
prsperse g xs = (++) ('L'.intersperse (toString g) (map toString xs))

between :: a Print Print -> Print | toString a
between g x y = x o pr g o y

surround :: a b Print -> Print | toString a & toString b
surround x y p = pr x o p o pr y

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

pBetween :: a Print Print -> Print | toString a
pBetween g x y = parens (between g x y)

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

// -- 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 = pr
instance set     Print    where set     = surround "[" "]" o prsperse ","
instance size Print Print where size    = surround "size (" ")"

instance + Print where + a b = pBetween " + " a b
instance - Print where - a b = pBetween " - " a b
instance * Print where * a b = pBetween " * " a b

instance +: Print Print where +: = pBetween " +: "
instance :+ Print Print where :+ = pBetween " :+ "
instance :- Print Print where :- = pBetween " :- "
instance *: Print Print where *: = pBetween " *: "

// -- 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 = pr
	=. = between " =. " o pr

// -- 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 = pr "true"
	false = pr "false"
	Not = surround "not (" ")"
	(||.) = pBetween " ||. "
	(&&.) = pBetween " ||. "

instance Compare Print Print
where
	==. = between " ==. "
	<=. = between " <=. "

instance In Print Print Print where In = between " in "

// -- 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
class (:.) infixl 1 a b :: (a b -> b)

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 :. (Sem a) (Sem b) where :. = (>>|)

instance If Print Print
where
	If b t e = pr "If (" o b o pr ") {" o t o pr "} else {" o e o pr "}"

instance For Print Print Print
where
	For x l b = pr "For " o pr x o pr " in " o l o pr " {" o b o pr "}"

instance :. Print Print where :. = between " :. "

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

import GenPrint
derive gPrint MaybeError

// Run with -b
Start = map (\x -> printToString x +++ "\n")
	[ print $ integer 5
	, print $ set [1,2,3,4,5]
	, print $ size (PRINT (set [1..10]))
	, print $ PRINT (set [1..3]) +: PRINT (integer 5 * integer 10)
	, print $ "x" =. integer 5
	, print $ factorial PRINT PRINT PRINT 5
	, printToString $ eval $ SET (For "x" (SET (set [0,1,2,3,4])) (variable "x" * INT (integer 5)))
	, printToString $ eval $ "y" =. INT (integer 5) :. SET (For "x" (INT (variable "y")) (variable "x" * INT (integer 5)))
	, printToString $ eval $ factorial INT SET BOOL 5
	]

/**
 * @param cast for integers
 * @param cast for sets
 * @param cast for booleans
 * @param integer to compute factorial for
 */
factorial :: (a -> a) (b -> b) (c -> c) Int -> a
	| Variable, integer, * a
	& :. a a & :. a b & :. b a
	& Compare a c
	& If c a
	& For a a b
factorial castint castset castlog n =
	"k" =. castint (integer n) :.
	"r" =. castint (integer 1) :.
	castset (For "x" (castint (variable "k"))
		(If (castlog (castint (variable "x") ==. castint (integer 0)))
			(variable "r") // for type-checking
			("r" =. castint (variable "r") * castint (variable "x")))) :.
	castint (variable "r")