{-# LANGUAGE LambdaCase #-}
module TypeCheck (inferProgram) where

import           AST
import           Control.Monad              (foldM, replicateM)
import           Control.Monad.Except       (runExceptT, throwError)
import           Control.Monad.Identity     (Identity (runIdentity))
import           Control.Monad.State        (StateT, get, put, runStateT)
import           Control.Monad.Trans.Except (ExceptT)
import qualified Data.Bifunctor
import qualified Data.Map                   as M
import qualified Data.Set                   as S


data Modifier = W | R MonoType deriving (Eq, Show)
type Env = M.Map String (Modifier, PolyType)
type Refinement = M.Map String MonoType

type TypeCheckerT e m a = StateT Int (ExceptT e m) a
runTypeCheckerT :: StateT a1 (ExceptT e m) a2 -> a1 -> m (Either e (a2, a1))
runTypeCheckerT m = runExceptT . runStateT m

type TypeChecker e a = TypeCheckerT e Identity a
runTypeChecker :: StateT a (ExceptT e Identity) a2 -> a -> Either e (a2, a)
runTypeChecker m = runIdentity . runTypeCheckerT m

inferExpr :: Env -> Term -> Either String ((Refinement, MonoType), Int)
inferExpr env expr = runTypeChecker (typeCheck expr env W) 0

inferProgram :: Program -> (Env, Either String MonoType)
inferProgram (Program (gadts, t)) =
    case inferExpr env t of
        Right ((_, t'), _) -> (env, Right t')
        Left x             -> (env, Left x)
      where
        env = foldr aux M.empty gadts
        aux (GADT (_, _, cs)) m = foldr (\(x, tc) m' -> M.insert x (W, PTForall (S.toList $ ftv tc) tc) m') m cs

typeCheck :: Term -> Env -> Modifier -> TypeChecker String (Refinement, MonoType)

typeCheck (TLit (LInt _)) _ _ = return (M.empty, intType)
typeCheck (TLit (LBool _)) _ _ = return (M.empty, boolType)

typeCheck (TBinop op u v) env m = do
    opType <- case op of
            Add   -> return $ MTFunc intType (MTFunc intType intType)
            Sub   -> return $ MTFunc intType (MTFunc intType intType)
            Mul   -> return $ MTFunc intType (MTFunc intType intType)
            Equal -> do
                tEq <- freshTv
                return $ MTFunc tEq (MTFunc tEq boolType)
    tFresh <- freshTv
    (ru, tu) <- typeCheck u env m
    let newEnv = applyEnv ru env
    (rv, tv) <- typeCheck v newEnv m
    r <- unify (apply (rv `compose` ru) (MTFunc tu (MTFunc tv tFresh))) opType
    return (r `compose` rv `compose` ru, apply r tFresh)

typeCheck (TIf u v w) env m = do
    (ru, tu) <- typeCheck u env m
    (rv, tv) <- typeCheck v (applyEnv ru env) m
    (rw, tw) <- typeCheck w (applyEnv (rv `compose` ru) env) m
    r <- unify boolType tu
    r2 <- unify (apply r tv) tw
    return (r2 `compose` r `compose` rw `compose` rv `compose` ru, apply r2 tw)

-- ATM
typeCheck (TVar x) env _ = do
    case M.lookup x env of
      Just (_, sig) -> do
        t <- instantiate sig
        return (M.empty, t)
      Nothing                -> throwError $ "Variable " ++ x ++ " not in scope"

typeCheck (TConstr c) env _ = do
    case M.lookup c env of
      Just (_, sig) -> do
        t <- instantiate sig
        return (M.empty, t)
      Nothing                -> throwError $ "Constr " ++ c ++ " not in scope"

-- APP
typeCheck (TApp u v) env _ = do
    t1 <- freshTv
    (ru, tu) <- typeCheck u env W
    let newEnv = applyEnv ru env
    (rv, tv) <- typeCheck v newEnv W
    r <- unify (apply rv tu) (MTFunc tv t1)
    return (r `compose` rv `compose` ru, apply r t1)

-- ABS
typeCheck (TLam x u) env modif = do
    (newModx, newModu) <- case modif of
        W                -> return (W, W)
        R (MTFunc t1 t2) -> return (R t1, R t2)
        R _              -> return (R $ MTVar "''useless", R $ MTVar "''useless") -- todo fresh and unify
        -- R t              -> throwError $ "Error in rule ABS " ++ show t ++ "\n" ++ show (TLam x u)
    tf <- freshTv
    let newEnv = M.insert x (newModx, PTForall [] tf) env
    (ru, tu) <- typeCheck u newEnv newModu
    return (ru, MTFunc (apply ru tf) tu)

-- LET-W
typeCheck (TLet x u v) env modif = do
    tx <- freshTv
    let newEnv = M.insert x (W, PTForall [] tx) env
    (ru, tu) <- typeCheck u newEnv W
    let newEnv2 = applyEnv ru newEnv
    let sigU = generalize newEnv2 tu
    let newEnv3 = M.insert x (W, sigU) newEnv2
    (rv, tv) <- typeCheck v newEnv3 modif
    return (ru `compose` rv, tv)

-- LET-R
typeCheck (TLetAnn x a u v) env modif = do
    (PTForall as t1, envWithMappings) <- pAnnToType a env
    let newEnv = M.insert x (R t1, PTForall as t1) envWithMappings
    (ru, tu) <- typeCheck u newEnv (R t1)
    let newEnv2 = applyEnv ru newEnv
    r <- unify tu (apply ru t1)
    let newEnv3  = M.insert x (R (apply r tu), PTForall as (apply r tu)) newEnv2
    (rv, tv) <- typeCheck v newEnv3 modif
    return (ru `compose` r `compose` rv, tv)

-- CASE
typeCheck (TCase _ []) _ _ = throwError "Empty case statement"
typeCheck (TCase u cs) env modif = do
    (mp, rp, tp) <- scrCheck u env
    let newEnv = applyEnv rp env
    tt <- freshTv
    (rUnif, tUnif) <- aux cs (rp, MTFunc tp tt) newEnv mp
    case tUnif of
        MTFunc _ tRes-> do return (rUnif, tRes)
        _             -> throwError $ "Error in rule CASE " ++ show (TCase u cs)

      where
        aux :: [(Pattern, Term)] -> (Refinement, MonoType) -> Env -> Modifier -> TypeChecker String (Refinement, MonoType)
        aux [] (r, tc) _ _ = return (r, tc)
        aux (c@(PConstr _ _, _):tl) (r, tc) currEnv mp = do
            (rC, tC) <- pConCheck c tc currEnv mp modif
            let newEnv = applyEnv rC currEnv
            aux tl (r `compose` rC, tC) newEnv mp

        aux (_:_) _ _ _ = throwError "Nested patterns not implemented"


scrCheck :: Term -> Env -> TypeChecker String (Modifier, Refinement, MonoType)

-- SCR-VAR
scrCheck (TVar x) env = do
    case M.lookup x env of
      Just (modif, PTForall _ t) -> return (modif, M.empty, t)
      Nothing     -> throwError $ "Variable " ++ x ++ " not in scope"

-- SCR-OTHER
scrCheck u env = do
    (r, t) <- typeCheck u env W
    return (W, r, t)


pConCheck :: (Pattern, Term) -> MonoType -> Env -> Modifier -> Modifier -> TypeChecker String (Refinement, MonoType)

-- PCON-W
pConCheck (PConstr constr xs, u) tExp env W modif = do
    (rPat, tPat) <- typeCheck (TConstr constr) env (R $ MTVar "''useless")
    (t1s, constrT, t2s) <- auxUnfold tPat []

    xsNames <- mapM (\case
                        PVar xN -> return xN
                        _ -> throwError "Nested patterns not implemented"
                    ) xs
    rUnif1 <- case tExp of
        MTFunc tLeftExp _ -> unify (MTConstr constrT t2s) tLeftExp
        _ -> throwError $ "Error in rule PCON-W unexpected type: " ++ show tExp

    let newEnv = foldr (\(x, t) m -> M.insert x (W, PTForall [] t) m) env (zip xsNames (map (apply rUnif1) t1s))
    (ru, tu) <- typeCheck u newEnv modif
    rUnif2 <- case tExp of
        MTFunc _ tResExp -> unify tu (apply ru tResExp)
        _ -> throwError $ "Error in rule PCON-W unexpected type: " ++ show tExp
    return (rUnif2 `compose` ru `compose` rUnif1 `compose` rPat, MTFunc (apply rUnif1 $ MTConstr constrT t2s) (apply rUnif2 tu))


-- PCON-R
pConCheck (PConstr constr xs, u) _ env (R tRigid) modif = do
    (_, tPat) <- typeCheck (TConstr constr) env (R $ MTVar "''useless")
    (t1s, constrT, t2s) <- auxUnfold tPat []

    xsNames <- mapM (\case
                        PVar xN -> return xN
                        _ -> throwError "Nested patterns not implemented"
                    ) xs

    rUnif1 <- freshUnify tRigid (MTConstr constrT t2s)

    let newEnv = applyEnv rUnif1 $ foldr (\(x, t) m -> M.insert x (R t, PTForall [] t) m) env (zip xsNames t1s)
    (ru, tu) <- typeCheck u newEnv modif
    tRes <- case modif of
        W   -> return tu
        R t -> do
            _ <- freshUnify tu (apply ru t)
            return t

    return (M.empty, MTFunc tRigid tRes)

pConCheck _ _ _ _ _               = throwError "Nested patterns not implemented"

auxUnfold :: MonoType -> [MonoType] -> TypeChecker String ([MonoType], String, [MonoType])
auxUnfold (MTConstr constrT t2s) t1s = return (reverse t1s, constrT, t2s)
auxUnfold (MTFunc t1 t2) t1s = auxUnfold t2 (t1:t1s)
auxUnfold t _ = throwError $ "Error when unfolding" ++ show t


-- 'a' ... 'z' 'aa' ... 'az' 'ba' ...
-- Taken from Learn You a Haskell
letters :: [String]
letters = [1..] >>= flip replicateM ['a'..'z']

freshString :: Int -> String
freshString n = '\'' : letters !! n

freshTv :: TypeChecker String MonoType
freshTv = do
    freshCounter <- get
    let tv = freshString freshCounter
    put $ freshCounter + 1
    return $ MTVar tv

-- | Return the free variables in a type
ftv :: MonoType -> S.Set String
ftv (MTVar x)         = S.singleton x
ftv (MTFunc t1 t2)    = S.union (ftv t1) (ftv t2)
ftv (MTConstr _ args) = S.unions (map ftv args)

ftvP :: PolyType -> S.Set String
ftvP (PTForall xs t) = S.difference (ftv t) (S.fromList xs)

ftvEnv :: Env -> S.Set String
ftvEnv env = S.unions $ map (ftvP . snd) $ M.elems env

-- | Apply a substitution
apply :: Refinement -> MonoType -> MonoType
apply s (MTConstr c args) = MTConstr c (map (apply s) args)
apply s t@(MTVar a)       = M.findWithDefault t a s
apply s (MTFunc t1 t2)    = MTFunc (apply s t1) (apply s t2)

applyP :: Refinement -> PolyType -> PolyType
applyP s (PTForall xs t) = PTForall xs (apply (foldr M.delete s xs) t)

applyEnv :: Refinement -> Env -> Env
applyEnv r = M.map (Data.Bifunctor.second (applyP r))

checkftv :: String -> MonoType -> Bool
checkftv a t = a `S.member` ftv t

compose :: Refinement -> Refinement -> Refinement
s1 `compose` s2 = M.map (apply s1) s2 `M.union` s1

-- | Unifier inspired by Write You a Haskell
unify :: MonoType -> MonoType -> TypeChecker String Refinement
unify (MTFunc l  r) (MTFunc l'  r') = do
    s1 <- unify l l'
    s2 <- unify (apply s1 r) (apply s1 r')
    return $ M.map (apply s2) s1 `M.union` s2
unify (MTVar a) t = bind a t
unify t (MTVar a) = bind a t
unify (MTConstr a argsa) (MTConstr b argsb) | a == b = foldM (\s (x, y) -> unify (apply s x) (apply s y)) M.empty (zip argsa argsb)
unify t1 t2 = throwError $ "Cannot unify " ++ show t1 ++ " with " ++ show t2

bind ::  String -> MonoType -> TypeChecker String Refinement
bind a t | t == MTVar a     = return M.empty
         | checkftv a t     = throwError $ "Cannot bind " ++ show (MTVar a) ++ " with " ++ show t ++ " (infinite type?)"
         | otherwise        = return $ M.singleton a t

instantiate :: PolyType -> TypeChecker String MonoType
instantiate (PTForall as t) = do
    as' <- mapM (const freshTv) as
    let s = M.fromList $ zip as as'
    return $ apply s t

generalize :: Env -> MonoType -> PolyType
generalize env t  = PTForall as t
    where as = S.toList $ ftv t `S.difference` ftvEnv env

mAnnToType :: MonoType -> Env -> TypeChecker String MonoType
mAnnToType (MTVar x) env = case M.lookup x env of
    Just (R _, PTForall _ t) -> return t
    _                        -> throwError $ "Variable " ++ x ++ " not in scope"
mAnnToType (MTFunc t1 t2) env = do
    t1' <- mAnnToType t1 env
    t2' <- mAnnToType t2 env
    return $ MTFunc t1' t2'
mAnnToType (MTConstr c args) env = do
    args' <- mapM (`mAnnToType` env) args
    return $ MTConstr c args'

pAnnToType :: PolyType -> Env -> TypeChecker String (PolyType, Env)
pAnnToType (PTForall oldAs t) env = do
    (newEnv, newAs) <- aux (env, []) oldAs
    let ftvInEnv = ftvEnv env
    if not (S.disjoint ftvInEnv (S.fromList newAs)) then throwError $ "Error when converting Ann " ++ show (PTForall oldAs t)
    else do
        t' <- mAnnToType t newEnv
        return (PTForall newAs t', newEnv)
      where aux (currEnv, currAs) [] = return (currEnv, currAs)
            aux (currEnv, currAs) (hd:tl) = do
                newA <- freshTv
                let v = (R newA, PTForall [] newA)
                case newA of
                    MTVar x -> aux (M.insert hd v currEnv, x : currAs) tl
                    _       -> throwError $ "Error when converting Ann " ++ show (PTForall oldAs t)


freshUnify :: MonoType -> MonoType -> TypeChecker String Refinement
freshUnify t1 t2 = do
    fmguStar <- auxFmgu [(t1, t2)] (ftv t2)
    let restrictedSpace = S.union (ftv t1) (ftv t2) -- Not sure why this restriction is necessary?
    return $ M.foldrWithKey (\k t acc -> if S.member k restrictedSpace then M.insert k t acc else acc ) M.empty fmguStar
        where
            auxFmgu :: [(MonoType, MonoType)] -> S.Set String -> TypeChecker String Refinement
            auxFmgu [] _ = return M.empty
            auxFmgu ((MTVar a, MTVar b):tl) bSet
                | a == b = auxFmgu tl bSet
                | S.member a bSet && S.notMember b bSet  = do
                    eps <- auxFmgu (auxApplyL tl a (MTVar b)) bSet
                    return $ M.insert a (MTVar b) eps
                | S.member b bSet && S.notMember a bSet  = do
                    eps <- auxFmgu (auxApplyL tl b (MTVar a)) bSet
                    return $ M.insert b (MTVar a) eps
                | otherwise = do
                    tc <- freshTv
                    eps <- auxFmgu (auxApplyL (auxApplyL tl a tc) b tc) bSet
                    return $ M.insert a tc (M.insert b tc eps)
            auxFmgu ((MTVar a, tau):tl) bSet | S.notMember a (ftv tau) = do
                eps <- auxFmgu (auxApplyL tl a tau) bSet
                return $ M.insert a tau eps
            auxFmgu ((tau, MTVar a):tl) bSet | S.notMember a (ftv tau) = do
                eps <- auxFmgu (auxApplyL tl a tau) bSet
                return $ M.insert a tau eps
            auxFmgu ((MTFunc l  r, MTFunc l'  r'):tl) bSet = auxFmgu ((l, l'):(r, r'):tl) bSet
            auxFmgu ((MTConstr a argsa, MTConstr b argsb):tl) bSet | a == b = auxFmgu (zip argsa argsb ++ tl) bSet
            auxFmgu ((t1a, t2a):_) _ = throwError $ "Cannot freshUnify " ++ show t1a ++ " with " ++ show t2a

            auxApplyL :: [(MonoType, MonoType)] -> String -> MonoType -> [(MonoType, MonoType)]
            auxApplyL [] _ _ = []
            auxApplyL ((t1a, t2a):tl) n ts = (auxApply t1a n ts, auxApply t2a n ts) : auxApplyL tl n ts

            auxApply :: MonoType -> String -> MonoType -> MonoType
            auxApply (MTVar n1) n2 ts
                | n1 == n2 = ts
                | otherwise = MTVar n1
            auxApply (MTFunc l  r) n ts = MTFunc (auxApply l n ts) (auxApply r n ts)
            auxApply (MTConstr a argsa) n ts = MTConstr a (map (\x -> auxApply x n ts) argsa)

