{-# OPTIONS -fglasgow-exts #-}

module LISS.Environment where

import LISS.Syntax

import qualified Data.Map as Map
import Control.Monad.State
import Control.Monad.Error

guardMsg :: (Error e, MonadError e m) => Bool -> String -> m ()
guardMsg b msg = if b then return () else throwError (strMsg msg)

type Seq = [Int]
type Set = (Int -> Bool)
type Array = [Int]

type IntFun   = (Args, Body, IntExp)
type BoolFun  = (Args, Body, BoolExp)
type SeqFun   = (Args, Body, SeqExp)
type SetFun   = (Args, Body, SetExp)
type ArrayFun = (Args, Body, ArrayExp)

data Env =
  Env { ints      :: Map.Map IntVar Int,
  	bools     :: Map.Map BoolVar Bool,
  	seqs      :: Map.Map SeqVar Seq,
	sets      :: Map.Map SetVar Set,
	arrays    :: Map.Map ArrayVar Array,
	intfuns   :: Map.Map ProgramName IntFun,
	boolfuns  :: Map.Map ProgramName BoolFun,
	seqfuns   :: Map.Map ProgramName SeqFun,
	setfuns   :: Map.Map ProgramName SetFun,
	arrayfuns :: Map.Map ProgramName ArrayFun
      }

emptyEnv = Env Map.empty Map.empty Map.empty Map.empty Map.empty Map.empty Map.empty Map.empty Map.empty Map.empty

instance Show Env where
	showsPrec _ env =
		showString "\nEnvironment:\n" .
			showString "\tIntegers: " . shows (ints env) .
			showString "\n\tBooleans: " . shows (bools env) .
			showString "\n\tSequences: " . shows (seqs env) .
			showString "\n\tSet variables: " . shows (Map.keys (sets env)) .
			showString "\n\tArrays: " . shows (arrays env) .
			showString "\n\tFunctions: " . shows ((Map.keys (intfuns env)) ++ (Map.keys (boolfuns env))
				++ (Map.keys (seqfuns env)) ++ (Map.keys (setfuns env)) ++ (Map.keys (arrayfuns env)))

genFresh :: [Var] -> Var
genFresh [] = "fresh"
genFresh [v] = if v == "fresh" then "fresher" else "fresh"
genFresh vs = concat vs

class SubEnv ref val | ref -> val where
	envName :: ref -> String
	access :: ref -> Env -> Map.Map Var val
	apply :: ref -> (Map.Map Var val -> Map.Map Var val) -> Env -> Env
	defaultVal :: ref -> val
	freshVar :: ref -> Env -> Var
	freshVar ref = genFresh . Map.keys . (access ref)
	contains :: (Error e, MonadError e m) => ref -> Var -> StateT Env m Bool
	contains ref var = gets (Map.member var . (access ref))
	getVal :: (Error e, MonadError e m) => ref -> Var -> StateT Env m val
	getVal ref var =
	     do declared <- contains ref var
		get >>= guardMsg declared .
		  (("Variable \'"++var++"\' undeclared in environment of "++envName ref++":")++) . show
		get >>= lift . (Map.lookup var . (access ref))
	setVal :: (Error e, MonadError e m) => ref -> Var -> val -> StateT Env m ()
	setVal ref var val =
	     do declared <- contains ref var
	     	get >>= guardMsg declared .
		  (("Variable \'"++var++"\' undeclared in environment of "++envName ref++":")++) . show
		modify (apply ref (Map.insert var val))
	setNew :: (Error e, MonadError e m) => ref -> Var -> val -> StateT Env m ()
	setNew ref var val =
	     do declared <- contains ref var
		get >>= guardMsg (not declared) .
		  (("Variable \'"++var++"\' already declared in environment of "++envName ref++":")++) . show
		modify (apply ref (Map.insert var val))
	addVar :: (Error e, MonadError e m) => ref -> Var -> StateT Env m ()
	addVar ref var = setNew ref var (defaultVal ref)

data IntEnv = IntEnv
instance SubEnv IntEnv Int where
	envName _ = "Integers"
	access _ = ints
	apply _ f env = Env (f (ints env)) (bools env) (seqs env) (sets env) (arrays env)
			    (intfuns env) (boolfuns env) (seqfuns env) (setfuns env) (arrayfuns env)
	defaultVal _ = 0
data BoolEnv = BoolEnv
instance SubEnv BoolEnv Bool where
	envName _ = "Booleans"
	access _ = bools
	apply _ f env = Env (ints env) (f (bools env)) (seqs env) (sets env) (arrays env)
			    (intfuns env) (boolfuns env) (seqfuns env) (setfuns env) (arrayfuns env)
	defaultVal _ = False
data SeqEnv = SeqEnv
instance SubEnv SeqEnv [Int] where
	envName _ = "Sequences"
	access _ = seqs
	apply _ f env = Env (ints env) (bools env) (f (seqs env)) (sets env) (arrays env)
			    (intfuns env) (boolfuns env) (seqfuns env) (setfuns env) (arrayfuns env)
	defaultVal _ = []
data SetEnv = SetEnv
instance SubEnv SetEnv (Int -> Bool) where
	envName _ = "Sets"
	access _ = sets
	apply _ f env = Env (ints env) (bools env) (seqs env) (f (sets env)) (arrays env)
			    (intfuns env) (boolfuns env) (seqfuns env) (setfuns env) (arrayfuns env)
	defaultVal _ = const False
data ArrayEnv = ArrayEnv
instance SubEnv ArrayEnv [Int] where
	envName _ = "Arrays"
	access _ = arrays
	apply _ f env = Env (ints env) (bools env) (seqs env) (sets env) (f (arrays env))
			    (intfuns env) (boolfuns env) (seqfuns env) (setfuns env) (arrayfuns env)
	defaultVal _ = []

getArrayPos :: ArrayVar -> Int -> StateT Env IO Int
getArrayPos var i =
     do arr <- getVal ArrayEnv var
     	guardMsg (i < (length arr) && i >= 0) "Array index out of bounds."
	return (arr !! i)

setArrayPos :: ArrayVar -> Int -> Int -> StateT Env IO ()
setArrayPos var i val=
     do arr <- getVal ArrayEnv var
     	guardMsg (i < (length arr) && i >= 0) "Array index out of bounds."
	setVal ArrayEnv var (replace arr i val)

data IntFunEnv = IntFunEnv
instance SubEnv IntFunEnv IntFun where
	envName _ = "Integer functions"
	access _ = intfuns
	apply _ f env = Env (ints env) (bools env) (seqs env) (sets env) (arrays env)
			    (f (intfuns env)) (boolfuns env) (seqfuns env) (setfuns env) (arrayfuns env)
	defaultVal _ = undefined
data BoolFunEnv = BoolFunEnv
instance SubEnv BoolFunEnv BoolFun where
	envName _ = "Boolean functions"
	access _ = boolfuns
	apply _ f env = Env (ints env) (bools env) (seqs env) (sets env) (arrays env)
			    (intfuns env) (f (boolfuns env)) (seqfuns env) (setfuns env) (arrayfuns env)
	defaultVal _ = undefined
data SeqFunEnv = SeqFunEnv
instance SubEnv SeqFunEnv SeqFun where
	envName _ = "Sequence functions"
	access _ = seqfuns
	apply _ f env = Env (ints env) (bools env) (seqs env) (sets env) (arrays env)
			    (intfuns env) (boolfuns env) (f (seqfuns env)) (setfuns env) (arrayfuns env)
	defaultVal _ = undefined
data SetFunEnv = SetFunEnv
instance SubEnv SetFunEnv SetFun where
	envName _ = "Set functions"
	access _ = setfuns
	apply _ f env = Env (ints env) (bools env) (seqs env) (sets env) (arrays env)
			    (intfuns env) (boolfuns env) (seqfuns env) (f (setfuns env)) (arrayfuns env)
	defaultVal _ = undefined
data ArrayFunEnv = ArrayFunEnv
instance SubEnv ArrayFunEnv ArrayFun where
	envName _ = "Array functions"
	access _ = arrayfuns
	apply _ f env = Env (ints env) (bools env) (seqs env) (sets env) (arrays env)
			    (intfuns env) (boolfuns env) (seqfuns env) (setfuns env) (f (arrayfuns env))
	defaultVal _ = undefined

-- auxiliars
replace l i v | (il,(_:rl)) <- splitAt i l = il++(v:rl)
