-- | Semantical analysis (i.e. type check and variances test) for Tempus expressions.
module Tempus.TypeCheck (
    TypeEnv,
    TypeSynEnv,
    ContextError (..),

    checkProgram,
    getType,
    substVar,
    expandTypeSyn,
    expandMuType,
    expandNuType,

    Variance (..),
    invertVariance
) where


import Control.Applicative
import Control.Arrow
import Control.Monad.Error
import Control.Monad.State

import Data.Generics.Uniplate.Operations
import Data.Function
import Data.List
import Data.Maybe

import Tempus.Loc
import Tempus.Syntax


-- | Type check monad.
type T a = ErrorT ContextError (State TState) a

-- | A general context check error.
data ContextError = UndefinedVariable Var    -- ^ Undefined global variable
                  | DuplicateVariable Var    -- ^ Variable defined twice
                  | DuplicateType Var        -- ^ Type defined twice
                  | OccursCheck Type Type    -- ^ Occurs check failed when trying to unify the types
                  | SymbolClash Type Type    -- ^ The types cannot be unified due to different
                                             -- outer-most constructors
                  | IncorrectVariances Type  -- ^ Incorrect variances
                  | UndefinedType Var        -- ^ Undefined type synonym or free type variable used
                  | TypeArgsMismatch Var     -- ^ Wrong number of arguments for the type synonym
                  | NoMuType Type            -- ^ Type has to be a mu type but isn't
                  | NoNuType Type            -- ^ Type has to be a nu type but isn't
                  | NoRecType Type           -- ^ Type has to be a recursive type but isn't
                  deriving (Show)


instance Error ContextError where
    noMsg      = error "noMsg not implemented for Error ContextError"
    strMsg msg = error "strMsg not implemented for Error ContextError"

instance Error e => Error (Loc e) where
    noMsg      = error "noMsg not implemented for Error (Loc e)"
    strMsg msg = error "stdMsg not implemented for Error (Loc e)"


-- | Type check state.
data TState = TState {
                  -- | The ID to use for the next fresh type variable.
                  varID :: Integer,
                  -- | The ID to use for the next fresh type constructor variable.
                  conID :: Integer
              } deriving (Eq, Show)


-- | Initial type check state.
initTState :: TState
initTState = TState { varID = 0, conID = 0 }


{- |
    Performs a semantical analysis for a module using the given global variables with types and
    type synonyms. Either a context error with source location or lists with the global variables
    and their inferred types and type synonyms is returned.
-}
checkProgram :: TypeEnv -> TypeSynEnv -> Program
             -> Either (Loc ContextError) (TypeEnv, TypeSynEnv)
checkProgram tEnv sEnv prog =
    fst $ runState (runErrorT $ checkDecls tEnv [] sEnv [] prog) initTState

-- | Helper function for 'checkProgram'.
checkDecls :: TypeEnv -> TypeEnv -> TypeSynEnv -> TypeSynEnv -> [Decl]
           -> ErrorT (Loc ContextError) (State TState) (TypeEnv, TypeSynEnv)
checkDecls _    tEnv' _    sEnv'[] = return (reverse tEnv', reverse sEnv')
checkDecls tEnv tEnv' sEnv sEnv' (DeclType loc v vs t : ds) = do
    when (not . null $ filter ((== v) . fst) $ sEnv ++ sEnv') $
        throwError $ Loc loc $ DuplicateType v
    checkDecls tEnv tEnv' sEnv ((v, (vs, t)) : sEnv') ds
checkDecls tEnv tEnv' sEnv sEnv' (DeclVal loc v e : ds) = do
    let env      = tEnv ++ tEnv'
        builtIns = map (\v -> (Var v, undefined)) ["add", "mult"]
    when (not . null $ filter ((== v) . fst) (env ++ builtIns)) $
        throwError $ Loc loc $ DuplicateVariable v
    t <- addErrorLoc loc $ checkedTypeExpr env (sEnv ++ sEnv') e
    checkDecls tEnv ((v, generalize t) : tEnv') sEnv sEnv' ds


-- | Infer the most general type for an expression using the types from a list of global variables
-- and a list of type synonyms.
getType :: TypeEnv -> TypeSynEnv -> Loc Expr -> Either (Loc ContextError) Type
getType tEnv sEnv (Loc loc e) =
    fst $ runState (runErrorT $ addErrorLoc loc $ checkedTypeExpr tEnv sEnv e) initTState


-- TODO: Remove type schemes?
-- | Type schemes.
type TypeScheme = ([Integer], Type)

-- | A type environment as a list of variables with their types.
type TypeEnv = [(Var, TypeScheme)]

-- | A substitution as a list of pairs of type variable IDs and the types of the corresponding
-- type variables.
type Subst = [(Integer, Type)]

-- | A type synonym environment as a list of pairs of a type synonym name with its parameter
-- names and type.
type TypeSynEnv = [(Var, ([Var], Type))]


-- | Transforms an error monad by adding a location information to the error type.
addErrorLoc :: (Error e, Functor m) => SrcLoc -> ErrorT e m a -> ErrorT (Loc e) m a
addErrorLoc loc = mapErrorT (fmap . left $ Loc loc)


-- | Perform a type inference with variance test within the @T@ monad for an expression, using the
-- given global variables with types and type synonyms.
checkedTypeExpr :: TypeEnv -> TypeSynEnv -> Expr -> T Type
checkedTypeExpr tEnv sEnv e = do
    t <- snd <$> typeExpr tEnv [] sEnv e
    var <- correctVariances sEnv t
    when (not $ var) $
        throwError $ IncorrectVariances t
    return t

-- | Performs a type inference for an expression. Returns the inferred type with a substitution
-- for the used type variables.
typeExpr :: TypeEnv     -- ^ Global type environment
         -> TypeEnv     -- ^ Local type environment
         -> TypeSynEnv  -- ^ Type synonym environment
         -> Expr        -- ^ Expression
         -> T (Subst, Type)
typeExpr gamma lambda sEnv e =
    case e of
        -- TODO: Put somewhere else
        ExVar (Var v)
            | v `elem` ["add", "mult"]
                -> return ([], TyFun TyNat $ TyFun TyNat TyNat)
        ExVar v -> case lookup v (lambda ++ gamma) of
                       Just s  -> (\t -> ([], t)) <$> instantiate s
                       Nothing -> throwError $ UndefinedVariable v
        ExLam x e -> do
            b <- freshVar
            (s, t) <- typeExpr (deleteBy ((==) `on` fst) (x, undefined) gamma)
                               (replVar (x, b) lambda) sEnv e
            return (s, TyFun (substType s b) t)
        ExApp e1 e2 -> do
            b <- freshVar
            (s1, t1) <- typeExpr gamma lambda sEnv e1
            (s2, t2) <- typeExpr gamma (substEnv s1 lambda) sEnv e2
            s3 <- mgu sEnv (substType s2 t1) (TyFun t2 b) []
            return (compSubst s3 $ compSubst s2 s1, substType s3 b)
        ExPair e1 e2 -> do
            (s1, t1) <- typeExpr gamma lambda sEnv e1
            (s2, t2) <- typeExpr gamma (substEnv s1 lambda) sEnv e2
            return (compSubst s2 s1, TyPair (substType s2 t1) t2)
{-
        ExLiftApp e1 e2 -> do
            (s1, t1) <- typeExpr gamma lambda e1
            [c, d] <- replicateM 2 freshVar
            s2 <- mgu (TyBehav $ TyFun c d) t1 []
            (s3, t2) <- typeExpr gamma (substEnv (compSubst s2 s1) lambda) e2
            s4 <- mgu (substType (compSubst s2 s1) $ TyBehav c) t2 []
            return (compSubst s4 $ compSubst s3 $ compSubst s2 s1, substType s4 $ TyBehav d)
-}
        -- TODO: rewrite
        ExLiftAppB e1 e2 -> do
            (s2', t1) <- typeExpr gamma lambda sEnv e1
            [a, b, c, d] <- replicateM 4 freshVar
            s3' <- mgu sEnv (TyFun (TyBehav $ TyFun c d)
                                   (TyFun (TyBehav c) (TyBehav d))) (TyFun t1 a) []
            (s2, t2) <- typeExpr gamma (substEnv (compSubst s3' s2') lambda) sEnv e2
            s3 <- mgu sEnv (substType (compSubst s2 s3') a) (TyFun t2 b) []
            return (compSubst s3 $ compSubst s2 $ compSubst s3' s2', substType s3 b)
        -- TODO: rewrite
        ExLiftAppE e1 e2 -> do
            (s2', t1) <- typeExpr gamma lambda sEnv e1
            [a, b, c, d] <- replicateM 4 freshVar
            s3' <- mgu sEnv (TyFun (TyBehav $ TyFun c d)
                                   (TyFun (TyEvent c) (TyEvent d))) (TyFun t1 a) []
            (s2, t2) <- typeExpr gamma (substEnv (compSubst s3' s2') lambda) sEnv e2
            s3 <- mgu sEnv (substType (compSubst s2 s3') a) (TyFun t2 b) []
            return (compSubst s3 $ compSubst s2 $ compSubst s3' s2', substType s3 b)
        ExConst e -> do
            (_, t) <- typeExpr gamma [] sEnv e
            return ([], TyBehav t)
        ExBehav f -> do
            (_, t) <- typeExpr gamma [] sEnv f
            a <- freshVar
            s <- mgu sEnv (TyFun TyNat a) t []
            return (s, substType s (TyBehav a))
        ExEvent t e -> do
            (s1, t1) <- typeExpr gamma lambda sEnv t
            s2 <- mgu sEnv TyNat t1 []
            (_, t2) <- typeExpr gamma [] sEnv e
            return (compSubst s2 s1, TyEvent t2)
        ExFold t f -> do
            mu@(MuType a t1) <- either throwError return $ expandMuType sEnv t
            (_, t2) <- typeExpr gamma [] sEnv f
            b <- freshVar
            s <- mgu sEnv (TyFun (substVar a b t1) b) t2 []
            return (s, substType s $ TyFun (TyMu mu) b)
        ExUnfold t f -> do
            nu@(NuType a t1) <- either throwError return $ expandNuType sEnv t
            (_, t2) <- typeExpr gamma [] sEnv f
            b <- freshVar
            s <- mgu sEnv (TyFun b (substVar a b t1)) t2 []
            return (s, substType s $ TyFun b (TyNu nu))
        -- TODO: throw returned error, not NoRecType
        ExPack t -> case (expandMuType sEnv t, expandNuType sEnv t) of
            (Right s@(MuType x t), Left _) -> return ([], TyFun (substVar x (TyMu s) t) (TyMu s))
            (Left _, Right s@(NuType x t)) -> return ([], TyFun (substVar x (TyNu s) t) (TyNu s))
            (Left _, Left _)               -> throwError $ NoRecType t
        ExUnpack t -> case (expandMuType sEnv t, expandNuType sEnv t) of
            (Right s@(MuType x t), Left _) -> return ([], TyFun (TyMu s) (substVar x (TyMu s) t))
            (Left _, Right s@(NuType x t)) -> return ([], TyFun (TyNu s) (substVar x (TyNu s) t))
            (Left _, Left _)               -> throwError $ NoRecType t
        _ -> (\t -> ([], t)) <$> typeSimple e


-- | Looks up a variable name in the type synonym list and applies the type synonym to the list of
-- argument types. Either a context error or the expanded type is returned.
expandTypeSyn :: TypeSynEnv -> Var -> [Type] -> Either ContextError Type
expandTypeSyn env v ts = case lookup v env of
    Nothing -> Left $ UndefinedType v
    Just (vs, t)
        | length vs == length ts -> Right $ foldr (uncurry substVar) t $ zip vs ts
        | otherwise -> Left $ TypeArgsMismatch v

-- | Tries to expand a type to a mu type. If the resulting type isn't a mu type, an error is
-- returned, otherwise, the mu type is returned.
expandMuType :: TypeSynEnv -> Type -> Either ContextError MuType
expandMuType sEnv (TyMu mu)    = Right mu
expandMuType sEnv (TyApp v ts) = expandTypeSyn sEnv v ts >>= expandMuType sEnv
expandMuType sEnv t            = Left $ NoMuType t

-- | Tries to expand a type to a nu type. If the resulting type isn't a nu type, an error is
-- returned, otherwise, the nu type is returned.
expandNuType :: TypeSynEnv -> Type -> Either ContextError NuType
expandNuType sEnv (TyNu nu)    = Right nu
expandNuType sEnv (TyApp v ts) = expandTypeSyn sEnv v ts >>= expandNuType sEnv
expandNuType sEnv t            = Left $ NoNuType t

-- | Instantiates a type scheme to a type.
instantiate :: TypeScheme -> T Type
instantiate (vs, t) = do
    bs <- replicateM (length vs) freshVar
    return $ substType (zip vs bs) t

-- TODO: condense type vars
-- | Generalizes a type to a type scheme.
generalize :: Type -> TypeScheme
generalize t = (ftv t, t)

-- | Returns the free type variables of a type as a list of type variable IDs.
ftv :: Type -> [Integer]
ftv = nub . ftv'
    where
        ftv' (TyVar i) = [i]
        ftv' t         = concatMap ftv' $ children t


-- | Compose two substitutions to a new one.
compSubst :: Subst -> Subst -> Subst
compSubst sigma theta = let theta' = map (second $ substType sigma) theta
                            sigma' = let domTheta = map fst theta
                                     in [ (x,t) | (x,t) <- sigma, x `notElem` domTheta ]
                            theta'' = [ (x,t) | (x,t) <- theta', t /= TyVar x ]
                        in nub $ theta'' ++ sigma'

{-
-- defined in Baader
compSubst :: Subst' -> Subst' -> Subst'
compSubst sigma theta = let sigma' = map (second $ substType theta) sigma
                            theta' = let domSigma = map fst sigma
                                     in [ (x,t) | (x,t) <- theta, x `notElem` domSigma ]
                            sigma'' = [ (x,t) | (x,t) <- sigma', t /= TyVar x ]
                        in nub $ sigma'' ++ theta'
-}

-- | Apply a substitution to each type in the type environment.
substEnv :: Subst -> TypeEnv -> TypeEnv
substEnv s = map (second $ second $ substType s)

-- | Apply a substitution to a type.
substType :: Subst -> Type -> Type
substType s t = transform f t
    where
        f (TyVar i) = fromMaybe (TyVar i) $ lookup i s
        f t = t

-- | Add, or replace if already present, a variable with a type in a list of variables with types.
replVar :: (Var, Type) -> TypeEnv -> TypeEnv
replVar (v,t) = ((v, ([], t)):) . filter ((/= v) . fst)


-- version Baader, Unification Theory; see also Heeren, Top Quality Type Error Messages
-- | Try to unify two types using a type synonym list and an initial substitution.
mgu :: TypeSynEnv -> Type -> Type -> Subst -> T Subst
mgu sEnv t1 t2 s =
    case (condSubst t1, condSubst t2) of
        (TyVar x, TyVar y)
            | x == y -> return []
        (TyVar x, t)
            | x `occursIn` t -> throwError $ OccursCheck (TyVar x) t
            | otherwise      -> return $ compSubst [(x, t)] s
        (t1, t2@(TyVar _)) -> mgu sEnv t2 t1 s
        (TyMu (MuType x1 t1), TyMu (MuType x2 t2)) -> do
            c <- freshCon
            mgu sEnv (substVar x1 c t1) (substVar x2 c t2) []
        (TyNu (NuType x1 t1), TyNu (NuType x2 t2)) -> do
            c <- freshCon
            mgu sEnv (substVar x1 c t1) (substVar x2 c t2) []
        (TyApp v ts, t) -> either throwError (\t' -> mgu sEnv t' t s) $ expandTypeSyn sEnv v ts
        (t, TyApp v ts) -> either throwError (\t' -> mgu sEnv t t' s) $ expandTypeSyn sEnv v ts
        (t1, t2)
            | equalSymbols t1 t2 -> unifySubterms (children t1) (children t2) s
            | otherwise -> throwError $ SymbolClash t1 t2
    where
        condSubst t@(TyVar _) = substType s t
        condSubst t           = t

        equalSymbols t1 t2 = case (t1, t2) of
            (TyFun _ _, TyFun _ _)   -> True
            (TyPlus _ _, TyPlus _ _) -> True
            (TyPair _ _, TyPair _ _) -> True
            (TyBehav _, TyBehav _)   -> True
            (TyEvent _, TyEvent _)   -> True
            (TyNat, TyNat)           -> True
            (TyZero, TyZero)         -> True
            (TyUnit, TyUnit)         -> True
            (TyCon x, TyCon y)       -> x == y
            _                        -> False

        -- TODO: use foldM
        unifySubterms [] [] sigma = return sigma
        unifySubterms (s:ss) (t:ts) sigma = do
            sigma' <- mgu sEnv s t sigma
            unifySubterms ss ts $ compSubst sigma sigma'
        unifySubterms _ _ _ = error "unifySubterms: same symbols with different arities"

-- | Check if a type variable with the given ID occurs in the type.
occursIn :: Integer -> Type -> Bool
occursIn x t = x `elem` [ y | TyVar y <- universe t]

-- | @substVar v s t@ replaces all occurrences of @v@ in @t@ by @s@. @v@ is not replaced if it is
-- bound locally by a mu type.
substVar :: Var -> Type -> Type -> Type
substVar x s t@(TyMu (MuType y t'))
    | x == y    = t
    | otherwise = TyMu (MuType y $ substVar x s t')
substVar x s t@(TyNu (NuType y t'))
    | x == y    = t
    | otherwise = TyNu (NuType y $ substVar x s t')
substVar x s t@(TyApp y [])
    | x == y    = s
    | otherwise = t
substVar x s t = descend (substVar x s) t

-- | Returns the type of a Tempus primitive expression.
typeSimple :: Expr -> T Type
typeSimple e =
    case e of
        ExNatLit _ -> return $ TyNat
        ExNull  -> TyFun TyZero <$> freshVar
        ExUnit  -> return $ TyUnit
        ExLeft  -> fresh2 $ \t1 t2 -> TyFun t1 (TyPlus t1 t2)
        ExRight -> fresh2 $ \t1 t2 -> TyFun t2 (TyPlus t1 t2)
        ExCase  -> fresh3 $ \t1 t2 s -> TyFun (TyFun t1 s) $
                                            TyFun (TyFun t2 s) $ TyFun (TyPlus t1 t2) s
        ExFst   -> fresh2 $ \t1 t2 -> TyFun (TyPair t1 t2) t1
        ExSnd   -> fresh2 $ \t1 t2 -> TyFun (TyPair t1 t2) t2
        ExExpand -> fresh1 $ \t -> TyFun (TyBehav t) (TyBehav $ TyPair t (TyBehav t))
        ExNever -> return $ TyEvent TyZero
        ExRace  -> fresh2 $ \t1 t2 -> TyFun (TyEvent t1) $
                                          TyFun (TyEvent t2) $
                                              TyEvent $ TyPlus (TyPair t1 t2) $
                                                        TyPlus (TyPair t1 $ TyEvent t2)
                                                               (TyPair t2 $ TyEvent t1)
        -- TODO: Replace n, s by new var
        ExReflect -> return $ TyFun TyNat $ TyMu $
                                  MuType (Var "n") (TyPlus TyUnit (TyApp (Var "n") []))
        ExUJump   -> fresh1 $ \t -> TyFun (TyNu $ NuType (Var "s") $ TyEvent $
                                               TyPlus t (TyApp (Var "s") []))
                                          (TyEvent t)
        ExUSwitch  -> fresh1 $ \t -> TyFun (TyNu $ NuType (Var "s") $
                                               TyPair (TyBehav t) $ TyEvent $ TyPair t $
                                                   TyApp (Var "s") [])
                                           (TyBehav t)
        _ -> error "typeSimple: expression is not simple"

-- | Generate a new unused type variable as a type expression.
freshVar :: T Type
freshVar = do
    i <- gets varID
    modify (\st -> st { varID = varID st + 1 })
    return $ TyVar i

-- | Generate a new unused type variable as a type expression and apply the given function to that
-- type.
fresh1 :: (Type -> Type) -> T Type
fresh1 f = liftM f freshVar

-- | Like 'fresh1' but for two variables.
fresh2 :: (Type -> Type -> Type) -> T Type
fresh2 f = liftM2 f freshVar freshVar

-- | Like 'fresh1' but for three variables.
fresh3 :: (Type -> Type -> Type -> Type) -> T Type
fresh3 f = liftM3 f freshVar freshVar freshVar

-- | Generate a new unused type constructor variable as a type expression.
freshCon :: T Type
freshCon = do
    i <- gets conID
    modify (\st -> st { conID = conID st + 1 })
    return $ TyCon i


-- | A variance.
data Variance = CoVariant | ContraVariant deriving (Eq, Show)

-- | Inverts a variance.
invertVariance :: Variance -> Variance
invertVariance CoVariant     = ContraVariant
invertVariance ContraVariant = CoVariant

-- TODO: Distinguish between bound variables with wrong variances and unbound variables
-- | Perform a variances check for the type using a type synonym environment, and return @true@
-- iff the test was successful.
correctVariances :: TypeSynEnv -> Type -> T Bool
correctVariances = var []
    where
        var delta sEnv (TyApp x ts)
            | (x, CoVariant) `elem` delta     = return True
            | (x, ContraVariant) `elem` delta = return False
            | otherwise                       = either throwError (var delta sEnv) $
                                                    expandTypeSyn sEnv x ts
        var delta _    (TyVar x)           = return True
        var _     _    TyNat               = return True
        var _     _    TyZero              = return True
        var _     _    TyUnit              = return True
        var delta sEnv (TyBehav t)         = var delta sEnv t
        var delta sEnv (TyEvent t)         = var delta sEnv t
        var delta sEnv (TyPair t1 t2)      = liftM2 (&&) (var delta sEnv t1) (var delta sEnv t2)
        var delta sEnv (TyPlus t1 t2)      = liftM2 (&&) (var delta sEnv t1) (var delta sEnv t2)
        var delta sEnv (TyFun t1 t2)       = liftM2 (&&)
                                                 (var (map (second invertVariance) delta) sEnv t1)
                                                 (var delta sEnv t2)
        var delta sEnv (TyMu (MuType v t)) = var (replVar (v, CoVariant) delta) sEnv t
        var delta sEnv (TyNu (NuType v t)) = var (replVar (v, CoVariant) delta) sEnv t

        replVar :: (Var, Variance) -> [(Var, Variance)] -> [(Var, Variance)]
        replVar (v,t) = ((v, t):) . filter ((/= v) . fst)