{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveFunctor #-}
{-# OPTIONS_GHC -fno-warn-overlapping-patterns #-}  -- TODO: remove
{-# OPTIONS_GHC -fno-warn-unused-binds #-}  -- TODO: remove
-- {-# OPTIONS_GHC -O0 #-}
module LambdaCube.Compiler.Infer
    ( inference
    , neutType'
    , makeCaseFunPars'
    ) where

import Data.Function
import Data.Monoid
import Data.Maybe
import Data.List
import qualified Data.Set as Set

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.Writer
import Control.Arrow hiding ((<+>))

import LambdaCube.Compiler.Utils
import LambdaCube.Compiler.DeBruijn
import LambdaCube.Compiler.Pretty hiding (braces, parens)
import LambdaCube.Compiler.DesugaredSource hiding (getList)
import LambdaCube.Compiler.Core
import LambdaCube.Compiler.InferMonad

------------------------------------------------------------------------------------

varType :: String -> Int -> Env -> (Binder, Exp)
varType err n_ env = f n_ env where
    f n (EAssign i (ET x _) es) = second (subst i x) $ f (if n < i then n else n+1) es
    f n (EBind2 b t es)  = if n == 0 then (b, up 1 t) else second (up 1) $ f (n-1) es
    f n (ELet2 _ (ET x t) es) = if n == 0 then (BLam Visible{-??-}, up 1 t) else second (up 1) $ f (n-1) es
    f n e = either (error $ "varType: " ++ err ++ "\n" ++ show n_ ++ "\n" ++ ppShow env) (f n) $ parent e

mkELet n x xt env = mkFun fn (Var <$> reverse vs) x
  where
    fn = FunName (FName n) (length vs) (ExpDef $ foldr addLam x vs) (foldr addPi xt vs)

    addLam v x = Lam $ rMove v 0 x
    addPi v x = Pi Visible (snd $ varType "mkELet" v env) $ rMove v 0 x

    vs = nub . concat $ grow [] mempty $ free x <> free xt

    -- TODO: avoid infinite loop if variable types refer each-other mutually
    grow accu acc s
        | Set.null s = accu
        | otherwise = grow (Set.toList s: accu) acc' s'
      where
        acc' = s <> acc
        s' = mconcat (free . snd . flip (varType "mkELet2") env <$> Set.toList s)

instance MkDoc a => PShow (CEnv a) where
    pShow = mkDoc (False, False)

instance PShow Env where
    pShow e = envDoc e $ underline $ text "<<HERE>>"

showEnvExp :: Env -> ExpType -> String
showEnvExp e c = show $ envDoc e $ underline $ pShow c

showEnvSExp :: (PShow a, HasFreeVars a) => Env -> SExp' a -> String
showEnvSExp e c = show $ envDoc e $ underline $ pShow c

showEnvSExpType :: (PShow a, HasFreeVars a) => Env -> SExp' a -> Exp -> String
showEnvSExpType e c t = show $ envDoc e $ underline $ (shAnn (pShow c) (pShow t))

envDoc :: Env -> Doc -> Doc
envDoc x m = case x of
    EGlobal{}           -> m
    EBind1 _ h ts b     -> envDoc ts $ shLam (usedVar' 0 b "") h m (pShow b)
    EBind2 h a ts       -> envDoc ts $ shLam (Just "") h (pShow a) m
    EApp1 _ h ts b      -> envDoc ts $ shApp h m (pShow b)
    EApp2 _ h (ET (Lam (Var 0)) (Pi Visible TType _)) ts -> envDoc ts $ shApp h (text "tyType") m
    EApp2 _ h a ts      -> envDoc ts $ shApp h (pShow a) m
    ELet1 _ ts b        -> envDoc ts $ shLet_ m (pShow b)
    ELet2 _ x ts        -> envDoc ts $ shLet_ (pShow x) m
    EAssign i x ts      -> envDoc ts $ shLet i (pShow x) m
    CheckType t ts      -> envDoc ts $ shAnn m $ pShow t
    CheckIType t ts     -> envDoc ts $ shAnn m (text "??") -- mkDoc ts' t
--    CheckSame t ts      -> envDoc ts $ shCstr <$> m <*> mkDoc ts' t
    CheckAppType si h t te b -> envDoc (EApp1 si h (CheckType_ (sourceInfo b) t te) b) m
    ERHS ts             -> envDoc ts $ shApp Visible "rhs" m
    ELHS n ts           -> envDoc ts $ shApp Visible ("lhs" `DApp` pShow n) m
    x   -> error $ "envDoc: " ++ ppShow x

instance MkDoc a => MkDoc (CEnv a) where
    mkDoc pr = \case
        MEnd a          -> mkDoc pr a
        Meta a b        -> shLam (Just "") BMeta (mkDoc pr a) (mkDoc pr b)
        Assign i (ET x _) e -> shLet i (mkDoc pr x) (mkDoc pr e)

-------------------------------------------------------------------------------- constraints env

data CEnv a
    = MEnd a
    | Meta Exp (CEnv a)
    | Assign !Int ExpType (CEnv a)       -- De Bruijn index decreasing assign reservedOp, only for metavariables (non-recursive)
  deriving (Functor)

instance (Subst Exp a, Rearrange a) => Rearrange (CEnv a) where
    rearrange l f = \case
            MEnd a -> MEnd $ rearrange l f a
            Meta a b -> Meta (rearrange l f a) (rearrange (l+1) f b)
            Assign j a b
                | l >  j -> assign j (rearrange (l-1) f a) (rearrange (l-1) f b)
                | l <= j -> assign (rearrangeFun f (j-l) + l) (rearrange l f a) (rearrange l f b)

instance (Subst Exp a, Rearrange a) => Subst Exp (CEnv a) where
    subst_ i dx x = \case
        MEnd a -> MEnd $ subst_ i dx x a
        Meta a b  -> Meta (subst_ i dx x a) (subst_ (i+1) (shiftFreeVars 1 dx) (up 1 x) b)
        Assign j a b
            | j > i, Just a' <- down i a       -> assign (j-1) a' (subst i (subst (j-1) (expr a') x) b)
            | j > i, Just x' <- down (j-1) x   -> assign (j-1) (subst i x' a) (subst i x' b)
            | j < i, Just a' <- down (i-1) a   -> assign j a' (subst (i-1) (subst j (expr a') x) b)
            | j < i, Just x' <- down j x       -> assign j (subst (i-1) x' a) (subst (i-1) x' b)
            | j == i                           -> Meta (cstr_ (ty a) x $ expr a) $ up1_ 0 b

--swapAssign :: (Int -> Exp -> CEnv Exp -> a) -> (Int -> Exp -> CEnv Exp -> a) -> Int -> Exp -> CEnv Exp -> a
swapAssign _ clet i (ET (Var j) t) b | i > j = clet j (ET (Var (i-1)) t) $ subst j (Var (i-1)) $ up1_ i b
swapAssign clet _ i a b = clet i a b

--assign :: Rearrange a => Int -> ExpType -> CEnv a -> CEnv a
assign = swapAssign Assign Assign

replaceMetas bind = \case
    Meta a t -> bind a $ replaceMetas bind t
    Assign i x t | x' <- up1_ i x -> bind (cstr_ (ty x') (Var i) $ expr x') . up 1 . up1_ i $ replaceMetas bind t
    MEnd t ->  t


-------------------------------------------------------------------------------- environments

-- SExp + Exp zipper
data Env
    = EBind1 SI Binder Env SExp2            -- zoom into first parameter of SBind
    | EBind2_ SI Binder Type Env             -- zoom into second parameter of SBind
    | EApp1 SI Visibility Env SExp2
    | EApp2 SI Visibility ExpType Env
    | ELet1 SIName Env SExp2
    | ELet2 SIName ExpType Env
    | EGlobal
    | ERHS Env
    | ELHS SIName Env

    | EAssign Int ExpType Env
    | CheckType_ SI Type Env
    | CheckIType SExp2 Env
--    | CheckSame Exp Env
    | CheckAppType SI Visibility Type Env SExp2   --pattern CheckAppType _ h t te b = EApp1 _ h (CheckType t te) b

pattern EBind2 b e env <- EBind2_ _ b e env where EBind2 b e env = EBind2_ (debugSI "6") b e env
pattern CheckType e env <- CheckType_ _ e env where CheckType e env = CheckType_ (debugSI "7") e env

parent = \case
    EAssign _ _ x        -> Right x
    EBind2 _ _ x         -> Right x
    EBind1 _ _ x _       -> Right x
    EApp1 _ _ x _        -> Right x
    EApp2 _ _ _ x        -> Right x
    ELet1 _ x _          -> Right x
    ELet2 _ _ x          -> Right x
    CheckType _ x        -> Right x
    CheckIType _ x       -> Right x
--    CheckSame _ x        -> Right x
    CheckAppType _ _ _ x _ -> Right x
    ERHS x               -> Right x
    ELHS _ x             -> Right x
    EGlobal              -> Left ()

-------------------------------------------------------------------------------- simple typing

neutType te = \case
    App_ f x        -> appTy (neutType te f) x
    Var_ i          -> snd $ varType "C" i te
    CaseFun_ s ts n -> appTy (foldl appTy (nType s) $ makeCaseFunPars te n ++ ts) (Neut n)
    TyCaseFun_ s [m, t, f] n -> foldl appTy (nType s) [m, t, Neut n, f]
    Fun s a _      -> foldlrev appTy (nType s) a

neutType' te = \case
    App_ f x        -> appTy (neutType' te f) x
    Var_ i          -> varType' i te
    CaseFun_ s ts n -> appTy (foldl appTy (nType s) $ makeCaseFunPars' te n ++ ts) (Neut n)
    TyCaseFun_ s [m, t, f] n -> foldl appTy (nType s) [m, t, Neut n, f]
    Fun s a _      -> foldlrev appTy (nType s) a

makeCaseFunPars te n = case neutType te n of
    (hnf -> TyCon (TyConName _ _ _ _ (CaseFunName _ _ pars)) xs) -> take pars $ reverse xs
    x -> error $ "makeCaseFunPars: " ++ ppShow x

makeCaseFunPars' te n = case neutType' te n of
    (hnf -> TyCon (TyConName _ _ _ _ (CaseFunName _ _ pars)) xs) -> take pars $ reverse xs

-------------------------------------------------------------------------------- inference

type ExpType' = CEnv ExpType

inferN :: forall m . Monad m => Env -> SExp2 -> IM m ExpType'
inferN e s = do
    b <- asks $ (TraceTypeCheck `elem`) . fst
    mapExceptT (mapReaderT $ mapWriterT $ fmap filt) $ inferN_ (if b then \s x m -> tell [ITrace s x] >> m else \_ _ m -> m) e s
  where
    filt (e@Right{}, is) = (e, filter f is)
    filt x = x

    f ITrace{} = False
    f _ = True

substTo i x = subst i x . up1_ (i+1)

mkLet x xt y = subst 0 x y
--mkLet x xt (ET y yt) = ET (Let (ET x xt) y) yt

ET a at `etApp` e = ET (app_ a e) (appTy at e)

inferN_ :: forall m . Monad m => (forall a . String -> String -> IM m a -> IM m a) -> Env -> SExp2 -> IM m ExpType'
inferN_ tellTrace = infer  where

    infer :: Env -> SExp2 -> IM m ExpType'
    infer te exp = tellTrace "infer" (showEnvSExp te exp) $ case exp of
        Parens x        -> infer te x
        SAnn x t        -> checkN (CheckIType x te) t TType
        SRHS x          -> infer (ERHS te) x
        SLHS n x        -> infer (ELHS n te) x
        SVar sn i       -> focusTell te exp $ ET (Var i) $ snd $ varType "C2" i te
        SLit si l       -> focusTell te exp $ ET (ELit l) (nType l)
        STyped et       -> focusTell' te exp et
        SGlobal (SIName si s) -> focusTell te exp =<< getDef te si s
        SLet le a b     -> infer (ELet1 le te b{-in-}) a{-let-} -- infer te SLamV b `SAppV` a)
        SApp_ si h a b  -> infer (EApp1 (si `validate` [sourceInfo a, sourceInfo b]) h te b) a
        SBind_ si h _ a b -> infer ((if h /= BMeta then CheckType_ (sourceInfo exp) TType else id) $ EBind1 si h te $ (if isPi h then TyType else id) b) a

    checkN :: Env -> SExp2 -> Type{-hnf-} -> IM m ExpType'
    checkN te x t = tellTrace "check" (showEnvSExpType te x t) $ checkN_ te x t

    checkN_ te (Parens e) t = checkN_ te e t
    checkN_ te e t
        | SBuiltin FprimFix `SAppV` (SLam Visible _ f) <- e = do
            pf <- getDef te mempty "primFix"
            checkN (EBind2 (BLam Visible) t $ EApp2 mempty Visible (pf `etApp` t) te) f $ up 1 t
        | x@(SGlobal (sName -> MatchName n)) `SAppV` SLamV (Wildcard _) `SAppV` a `SAppV` SVar siv v `SAppV` b <- e
            = infer te $ x `SAppV` SLam Visible SType (STyped (ET (subst (v+1) (Var 0) $ up 1 t) TType)) `SAppV` a `SAppV` SVar siv v `SAppV` b
            -- temporal hack
        | x@(SGlobal (sName -> CaseName "'Nat")) `SAppV` SLamV (Wildcard _) `SAppV` a `SAppV` b `SAppV` SVar siv v <- e
            = infer te $ x `SAppV` SLamV (STyped (ET (substTo (v+1) (Var 0) $ up 1 t) TType)) `SAppV` a `SAppV` b `SAppV` SVar siv v
            -- temporal hack
        | x@(SGlobal (sName -> CaseName "'VecS")) `SAppV` SLamV (SLamV (Wildcard _)) `SAppV` a `SAppV` b `SAppV` c `SAppV` SVar siv v <- e
        , TyConN (FTag F'VecS) [Var n', _] <- snd $ varType "xx" v te
            = infer te $ x `SAppV` SLamV (SLamV (STyped (ET (substTo (n'+2) (Var 1) $ up 2 t) TType))) `SAppV` a `SAppV` b `SAppV` c `SAppV` SVar siv v

{-
            -- temporal hack
        | x@(SGlobal (si, "'HListCase")) `SAppV` SLamV (SLamV (Wildcard _)) `SAppV` a `SAppV` b `SAppV` SVar siv v <- e
        , TVec (Var n') _ <- snd $ varType "xx" v te
            = infer te $ x `SAppV` SLamV (SLamV (STyped (subst (n'+2) (Var 1) $ up1_ (n'+3) $ up 2 t, TType))) `SAppV` a `SAppV` b `SAppV` SVar siv v
-}
        | SRHS x <- e = checkN (ERHS te) x t
{- TODO
        | SAnn v a <- e = do
            let x = t
            let same = checkSame te a x
            if same then checkN te v x else  error $ "checkSame:\n" ++ ppShow a ++ "\nwith\n" ++ showEnvExp te (ET x TType)
-}
        | SLHS n x <- e = checkN (ELHS n te) x t
        | SApp_ si h a b <- e = infer (CheckAppType si h t te b) a
        | SLam h a b <- e, Pi h' x y <- t, h == h' = do
--            tellType e t
            let same = checkSame te a x
            if same then checkN (EBind2 (BLam h) x te) b $ hnf y else error $ "checkSame:\n" ++ ppShow a ++ "\nwith\n" ++ showEnvExp te (ET x TType)
        | Pi Hidden a b <- t = do
            bb <- notHiddenLam e
            if bb then checkN (EBind2 (BLam Hidden) a te) (up1 e) $ hnf b
                 else infer (CheckType_ (sourceInfo e) t te) e
        | otherwise = infer (CheckType_ (sourceInfo e) t te) e
      where
        -- todo
        notHiddenLam = \case
            SLam Visible _ _ -> return True
            SGlobal (sName -> s) -> do
                nv <- asks snd
                case fromMaybe (error $ "infer: can't find: " ++ s) $ lookupName s nv of
                    ET (Lam _) (Pi Hidden _ _) -> return False
                    _ -> return True
            _ -> return False
{-
    -- todo
    checkSame te (Wildcard _) a = return (te, True)
    checkSame te x y = do
        (ex, _) <- checkN te x TType
        return $ ex == y
-}
    checkSame te (Wildcard _) a = True
    checkSame te SType TType = True
    checkSame te (SBind_ _ BMeta _ SType (STyped (ET (Var 0) _))) a = True
    checkSame te a b = error $ "checkSame: " ++ ppShow (a, b)

    hArgs (Pi Hidden _ b) = 1 + hArgs b
    hArgs _ = 0

    focusTell env si eet = tellType si (ty eet) >> focus_ env eet
    focusTell' env si eet = focus_ env eet

    focus_ :: Env -> ExpType -> IM m ExpType'
    focus_ env eet@(ET e et) = tellTrace "focus" (showEnvExp env eet) $ case env of
        ERHS te -> focus_ te (ET (RHS $ hnf e) et)
        ELHS n te -> focus_ te (ET (mkELet n e et te) et)
--        CheckSame x te -> focus_ (EBind2_ (debugSI "focus_ CheckSame") BMeta (cstr x e) te) $ up 1 eet
        CheckAppType si h t te b   -- App1 h (CheckType t te) b
            | Pi h' x (down 0 -> Just y) <- et, h == h' -> case t of
                Pi Hidden t1 t2 | h == Visible -> focus_ (EApp1 si h (CheckType_ (sourceInfo b) t te) b) eet  -- <<e>> b : {t1} -> {t2}
                _ -> focus_ (EBind2_ (sourceInfo b) BMeta (cstr_ TType t y) $ EApp1 si h te b) $ up 1 eet
            | otherwise -> focus_ (EApp1 si h (CheckType_ (sourceInfo b) t te) b) eet
        EApp1 si h te b
            | Pi h' x y <- et, h == h' -> checkN (EApp2 si h eet te) b $ hnf x
            | Pi Hidden x y  <- et, h == Visible -> focus_ (EApp1 mempty Hidden env $ Wildcard $ Wildcard SType) eet  --  e b --> e _ b
--            | CheckType (Pi Hidden _ _) te' <- te -> error "ok"
--            | CheckAppType Hidden _ te' _ <- te -> error "ok"
            | otherwise -> infer (CheckType_ (sourceInfo b) (Var 2) $ cstr' h (up 2 et) (Pi Visible (Var 1) (Var 1)) (up 2 e) $ EBind2_ (sourceInfo b) BMeta TType $ EBind2_ (sourceInfo b) BMeta TType te) (up 3 b)
          where
            cstr' h x y e = EApp2 mempty h (ET (evalCoe (up 1 x) (up 1 y) (Var 0) (up 1 e)) (up 1 y)) . EBind2_ (sourceInfo b) BMeta (cstr_ TType x y)
        ELet2 ln (ET x{-let-} xt) te -> focus_ te $ mkLet x{-let-} xt eet{-in-}
        CheckIType x te -> checkN te x $ hnf e
        CheckType_ si t te
            | hArgs et > hArgs t
                            -> focus_ (EApp1 mempty Hidden (CheckType_ si t te) $ Wildcard $ Wildcard SType) eet
            | hArgs et < hArgs t, Pi Hidden t1 t2 <- t
                            -> focus_ (CheckType_ si t2 $ EBind2 (BLam Hidden) t1 te) eet
            | otherwise    -> focus_ (EBind2_ si BMeta (cstr_ TType t et) te) $ up 1 eet
        EApp2 si h a te    -> focusTell te si $ a `etApp` e        --  h??
        EBind1 si h te b   -> infer (EBind2_ (sourceInfo b) h e te) b
        EBind2_ si (BLam h) a te -> focus_ te $ lamPi h a eet
        EBind2_ si (BPi h) a te -> focusTell te si $ ET (Pi h a e) TType
        _ -> focus2 env $ MEnd eet

    focus2 :: Env -> CEnv ExpType -> IM m ExpType'
    focus2 env eet = case env of
        ELet1 le te b{-in-} -> infer (ELet2 le (replaceMetas' eet{-let-}) te) b{-in-}
        EBind2_ si BMeta tt_ te
            | ERHS te'   <- te -> refocus (ERHS $ EBind2_ si BMeta tt_ te') eet
            | ELHS n te' <- te -> refocus (ELHS n $ EBind2_ si BMeta tt_ te') eet
            | Unit <- tt    -> refocus te $ subst 0 TT eet
            | Empty msg <- tt -> throwError' $ ETypeError (text msg) si
            | CW (hnf -> T2 _ x y) <- tt, let te' = EBind2_ si BMeta (up 1 $ cw y) $ EBind2_ si BMeta (cw x) te
                            -> refocus te' $ subst 2 (t2C (Var 1) (Var 0)) $ up 2 eet
            | CW (hnf -> CstrT t a b) <- tt, Just r <- cst (ET a t) b -> r
            | CW (hnf -> CstrT t a b) <- tt, Just r <- cst (ET b t) a -> r
            | CW _ <- tt, EBind2 h x te' <- te, Just x' <- down 0 tt, x == x'
                            -> refocus te $ subst 1 (Var 0) eet
            | EBind2_ si' h x te' <- te, h /= BMeta, Just b' <- down 0 tt
                            -> refocus (EBind2_ si' h (up 1 x) $ EBind2_ si BMeta b' te') $ subst 2 (Var 0) $ up 1 eet
            | ELet2 le x te' <- te, Just b' <- down 0 tt
                            -> refocus (ELet2 le (up 1 x) $ EBind2_ si BMeta b' te') $ subst 2 (Var 0) $ up 1 eet
            | EBind1 si h te' x <- te -> refocus (EBind1 si h (EBind2_ si BMeta tt_ te') $ up1_ 1 x) eet
            | ELet1 le te' x     <- te, floatLetMeta $ ty $ replaceMetas' $ Meta tt_ eet
                                    -> refocus (ELet1 le (EBind2_ si BMeta tt_ te') $ up1_ 1 x) eet
            | CheckAppType si h t te' x <- te -> refocus (CheckAppType si h (up 1 t) (EBind2_ si BMeta tt_ te') $ up1 x) eet
            | EApp1 si h te' x <- te -> refocus (EApp1 si h (EBind2_ si BMeta tt_ te') $ up1 x) eet
            | EApp2 si h x te' <- te -> refocus (EApp2 si h (up 1 x) $ EBind2_ si BMeta tt_ te') eet
            | CheckType_ si t te' <- te -> refocus (CheckType_ si (up 1 t) $ EBind2_ si BMeta tt_ te') eet
--            | CheckIType x te' <- te -> refocus (CheckType_ si (up 1 t) $ EBind2_ si BMeta tt te') eet
            | otherwise             -> focus2 te $ Meta tt_ eet
          where
            tt = hnf tt_
            refocus = refocus_ focus2
            cst :: ExpType -> Exp -> Maybe (IM m ExpType')
            cst x = \case
                Var i | fst (varType "X" i te) == BMeta
                      , Just y <- down i x
                      -> Just $ join swapAssign (\i x -> refocus $ EAssign i x te) i y $ subst 0 {-ReflCstr y-}TT $ subst (i+1) (expr $ up 1 y) eet
                _ -> Nothing

        EAssign i b te -> case te of
            ERHS te'     -> refocus' (ERHS $ EAssign i b te') eet
            ELHS n te'   -> refocus' (ELHS n $ EAssign i b te') eet
            EBind2_ si h x te' | i > 0, Just b' <- down 0 b
                              -> refocus' (EBind2_ si h (subst (i-1) (expr b') x) (EAssign (i-1) b' te')) eet
            ELet2 le x te' | i > 0, Just b' <- down 0 b
                              -> refocus' (ELet2 le (subst (i-1) (expr b') x) (EAssign (i-1) b' te')) eet
            ELet1 le te' x    -> refocus' (ELet1 le (EAssign i b te') $ subst (i+1) (up 1 b) x) eet
            EBind1 si h te' x -> refocus' (EBind1 si h (EAssign i b te') $ subst (i+1) (up 1 b) x) eet
            CheckAppType si h t te' x -> refocus' (CheckAppType si h (subst i (expr b) t) (EAssign i b te') $ subst i b x) eet
            EApp1 si h te' x  -> refocus' (EApp1 si h (EAssign i b te') $ subst i b x) eet
            EApp2 si h x te'  -> refocus' (EApp2 si h (subst i (expr b) x) $ EAssign i b te') eet
            CheckType_ si t te'   -> refocus' (CheckType_ si (subst i (expr b) t) $ EAssign i b te') eet
            EAssign j a te' | i < j
                              -> refocus' (EAssign (j-1) (subst i (expr b) a) $ EAssign i (up1_ (j-1) b) te') eet
            t  | Just te' <- pull1 i b te -> refocus' te' eet
               | otherwise -> swapAssign (\i x -> focus2 te . Assign i x) (\i x -> refocus' $ EAssign i x te) i b eet
            -- todo: CheckSame Exp Env
          where
            refocus' = fix refocus_

            pull1 i b = \case
                EBind2_ si h x te | i > 0, Just b' <- down 0 b
                    -> EBind2_ si h (subst (i-1) (expr b') x) <$> pull1 (i-1) b' te
                EAssign j a te
                    | i < j  -> EAssign (j-1) (subst i (expr b) a) <$> pull1 i (up1_ (j-1) b) te
                    | j <= i -> EAssign j (subst i (expr b) a) <$> pull1 (i+1) (up1_ j b) te
                te  -> pull i te

            pull i = \case
                EBind2 BMeta _ te | i == 0 -> Just te
                EBind2_ si h x te | i > 0 -> EBind2_ si h <$> down (i-1) x <*> pull (i-1) te
                EAssign j a te  -> EAssign (if j <= i then j else j-1) <$> down i a <*> pull (if j <= i then i+1 else i) te
                _               -> Nothing

        EGlobal{} -> return eet
        _ -> case eet of
            MEnd x -> throwError' $ ErrorMsg $ "focus todo:" <+> pShow x
            _ -> case env of
                _ -> throwError' $ ErrorMsg $ "focus checkMetas:" <$$> pShow env <$$> "---" <$$> pShow eet
      where
        refocus_ :: (Env -> CEnv ExpType -> IM m ExpType') -> Env -> CEnv ExpType -> IM m ExpType'
        refocus_ _ e (MEnd at) = focus_ e at
        refocus_ f e (Meta x at) = f (EBind2 BMeta x e) at
        refocus_ _ e (Assign i x at) = focus2 (EAssign i x e) at

        replaceMetas' = replaceMetas $ lamPi Hidden

-------------------------------------------------------------------------------- re-checking

type Message = String

recheck :: Monad m => SIName -> Env -> ExpType -> m ExpType
recheck sn e = return . recheck' sn e

-- todo: check type also
recheck' :: SIName -> Env -> ExpType -> ExpType
recheck' sn e (ET x xt) = ET (recheck_ "main" (checkEnv e) (ET x xt)) xt
  where
    err s = error $ "At " ++ ppShow sn ++ "\n" ++ s

    checkEnv = \case
        e@EGlobal{} -> e
        EBind1 si h e b -> EBind1 si h (checkEnv e) b
        EBind2_ si h t e -> EBind2_ si h (checkType e t) $ checkEnv e            --  E [\(x :: t) -> e]    -> check  E [t]
        ELet1 le e b -> ELet1 le (checkEnv e) b
        ELet2 le x e -> ELet2 le (recheck'' "env" e x) $ checkEnv e
        EApp1 si h e b -> EApp1 si h (checkEnv e) b
        EApp2 si h a e -> EApp2 si h (recheck'' "env" e a) $ checkEnv e    --  E [a x]  ->  check
        EAssign i x e -> EAssign i (recheck'' "env" e $ up1_ i x) $ checkEnv e                -- __ <i := x>
        CheckType_ si x e -> CheckType_ si (checkType e x) $ checkEnv e
--        CheckSame x e -> CheckSame (recheck'' "env" e x) $ checkEnv e
        CheckAppType si h x e y -> CheckAppType si h (checkType e x) (checkEnv e) y

    recheck'' msg te a@(ET x xt) = ET (recheck_ msg te a) xt
    checkType te e = recheck_ "check" te (ET e TType)

    recheck_ msg te = \case
        ET (Var k) zt -> Var k    -- todo: check var type
        ET (Lam_ md b) (Pi h a bt) -> Lam_ md $ recheck_ "9" (EBind2 (BLam h) a te) (ET b bt)
        ET (Pi_ md h a b) TType -> Pi_ md h (checkType te a) $ checkType (EBind2 (BPi h) a te) b
        ET (ELit l) zt -> ELit l  -- todo: check literal type
        ET TType TType -> TType
        ET (Neut (App__ md a b)) zt
            | ET (Neut a') at <- recheck'' "app1" te $ ET (Neut a) (neutType te a)
            -> checkApps "a" [] zt (Neut . App__ md a' . head) te at [b]
        ET (Con_ md s n as) zt      -> checkApps (ppShow s) [] zt (Con_ md s n . reverse . drop (conParams s)) te (conType zt s) $ mkConPars n zt ++ reverse as
        ET (TyCon_ md s as) zt      -> checkApps (ppShow s) [] zt (TyCon_ md s . reverse) te (nType s) $ reverse as
        ET (Neut (CaseFun__ fs s@(CaseFunName _ t pars) as n)) zt -> checkApps (ppShow s) [] zt (\xs -> evalCaseFun fs s (init $ drop pars xs) (last xs)) te (nType s) (makeCaseFunPars te n ++ as ++ [Neut n])
        ET (Neut (TyCaseFun__ fs s [m, t, f] n)) zt  -> checkApps (ppShow s) [] zt (\[m, t, n, f] -> evalTyCaseFun_ fs s [m, t, f] n) te (nType s) [m, t, Neut n, f]
        ET (Neut (Fun_ md f a x)) zt -> checkApps ("lab-" ++ show f ++ ppShow a ++ "\n" ++ ppShow (nType f)) [] zt (\xs -> Neut $ Fun_ md f (reverse xs) x) te (nType f) $ reverse a   -- TODO: recheck x
        ET (Let (ET x xt) y) zt -> Let (recheck'' "let" te $ ET x xt) $ recheck_ "let_in" te $ ET y (Pi Visible xt $ up1 zt)  -- TODO
        ET (RHS x) zt -> RHS $ recheck_ msg te (ET x zt)
      where
        checkApps s acc zt f _ t []
            | t == zt = f $ reverse acc
            | otherwise = 
                     err $ "checkApps' " ++ s ++ " " ++ msg ++ "\n" ++ showEnvExp te{-todo-} (ET t TType) ++ "\n\n" ++ showEnvExp te (ET zt TType)
        checkApps s acc zt f te t@(hnf -> Pi h x y) (b_: xs) = checkApps (s++"+") (b: acc) zt f te (appTy t b) xs where b = recheck_ "checkApps" te (ET b_ x)
        checkApps s acc zt f te t _ =
             err $ "checkApps " ++ s ++ " " ++ msg ++ "\n" ++ showEnvExp te{-todo-} (ET t TType) ++ "\n\n" ++ showEnvExp e (ET x xt)

-------------------------------------------------------------------------------- inference for statements

inference :: MonadFix m => [Stmt] -> IM m [GlobalEnv]
inference [] = return []
inference (x:xs) = do
    y <- handleStmt x
    (y:) <$> withEnv y (inference xs)

handleStmt :: MonadFix m => Stmt -> IM m GlobalEnv
handleStmt = \case
  Primitive n t_ -> do
        t <- inferType n $ trSExp' t_
        tellType (sourceInfo n) t
        let fn = mkFunDef (FName n) t
        addToEnv n $ flip ET t $ lamify' t $ \xs -> Neut $ Fun fn xs delta
  StLet n mt t_ -> do
        let t__ = maybe id (flip SAnn) mt t_
        ET x t <- inferTerm n $ trSExp' t__
        tellType (sourceInfo n) t
        addToEnv n (ET x t)
{-        -- hack
        when (snd (getParams t) == TType) $ do
            let ps' = fst $ getParams t
                t'' =   (TType :~> TType)
                  :~> addParams ps' (Var (length ps') `app_` DFun (FunName (snd n) t) (downTo 0 $ length ps'))
                  :~>  TType
                  :~> Var 2 `app_` Var 0
                  :~> Var 3 `app_` Var 1
            addToEnv (fst n, MatchName (snd n)) (lamify t'' $ \[m, tr, n', f] -> evalTyCaseFun (TyCaseFunName (snd n) t) [m, tr, f] n', t'')
-}
  PrecDef{} -> return mempty
  Data s (map (second trSExp') -> ps) (trSExp' -> t_@(UncurryS tl_ _)) cs -> do
    vty <- inferType s $ UncurryS ps t_
    tellType (sourceInfo s) vty
    let
        sint = FName s
        pnum' = length $ filter ((== Visible) . fst) ps
        inum = arity vty - length ps

        mkConstr j (cn, trSExp' -> ct@(UncurryS ctl (AppsS c (map snd -> xs))))
            | c == SGlobal s && take pnum' xs == downToS "a3" (length ctl) pnum'
            = do
                cty <- removeHiddenUnit <$> inferType cn (UncurryS [(Hidden, x) | (Visible, x) <- ps] ct)
--                tellType (sourceInfo cn) cty
                let     pars = zipWith (\x -> second $ STyped . flip ET TType . rUp (1+j) x) [0..] $ drop (length ps) $ fst $ getParams cty
                        act = length . fst . getParams $ cty
                        acts = map fst . fst . getParams $ cty
                        conn = ConName (FName cn) j cty
                e <- addToEnv cn $ ET (Con conn 0 []) cty
                return (e, ((conn, cty)
                       , UncurryS pars
                       $ (foldl SAppV (sVar ".cs" $ j + length pars) $ drop pnum' xs ++ [AppsS (SGlobal cn) (zip acts $ downToS ("a4 " ++ sName cn ++ " " ++ show (length ps)) (j+1+length pars) (length ps) ++ downToS "a5" 0 (act- length ps))]
                       :: SExp2)))
            | otherwise = throwError' $ ErrorMsg "illegal data definition (parameters are not uniform)" -- ++ show (c, cn, take pnum' xs, act)

        motive = UncurryS (replicate inum (Visible, Wildcard SType)) $
           SPi Visible (AppsS (SGlobal s) $ zip (map fst ps) (downToS "a6" inum $ length ps) ++ zip (map fst tl_) (downToS "a7" 0 inum)) SType

    (e1, es, tcn, cfn@(CaseFunName _ ct _), _, _) <- mfix $ \ ~(_, _, _, _, ct', cons') -> do
        let cfn = CaseFunName sint ct' $ length ps
        let tcn = TyConName sint inum vty (map fst cons') cfn
        e1 <- addToEnv s (ET (TyCon tcn []) vty)
        (unzip -> (mconcat -> es, cons)) <- withEnv e1 $ zipWithM mkConstr [0..] cs
        ct <- withEnv (e1 <> es) $ inferType s
            ( UncurryS
                ( [(Hidden, x) | (_, x) <- ps]
                ++ (Visible, motive)
                : map ((,) Visible . snd) cons
                ++ replicate inum (Hidden, Wildcard SType)
                ++ [(Visible, AppsS (SGlobal s) $ zip (map fst ps) (downToS "a8" (inum + length cs + 1) $ length ps) ++ zip (map fst tl_) (downToS "a9" 0 inum))]
                )
            $ foldl SAppV (sVar ".ct" $ length cs + inum + 1) $ downToS "a10" 1 inum ++ [sVar ".24" 0]
            )
        return (e1, es, tcn, cfn, ct, cons)

    e2 <- addToEnv (SIName (sourceInfo s) $ CaseName (sName s)) $ ET (lamify ct $ \xs -> evalCaseFun' cfn (init $ drop (length ps) xs) (last xs)) ct
    let ps' = fst $ getParams vty
        t =   (TType :~> TType)
          :~> addParams ps' (Var (length ps') `app_` TyCon tcn (downTo' 0 $ length ps'))
          :~>  TType
          :~> Var 2 `app_` Var 0
          :~> Var 3 `app_` Var 1
    e3 <- addToEnv (SIName (sourceInfo s) $ MatchName (sName s)) $ ET (lamify t $ \[m, tr, n, f] -> evalTyCaseFun (TyCaseFunName sint t) [m, tr, f] n) t
    return (e1 <> e2 <> e3 <> es)

  stmt -> error $ "handleStmt: " ++ ppShow stmt

inferTerm :: Monad m => SIName -> SExp2 -> IM m ExpType
inferTerm sn t = recheck sn EGlobal . replaceMetas (lamPi Hidden) =<< inferN EGlobal t

inferType :: Monad m => SIName -> SExp2 -> IM m Type
inferType sn t = fmap expr
    $ recheck sn EGlobal . flip ET TType . replaceMetas (Pi Hidden) . fmap expr
    =<< inferN (CheckType_ (sourceInfo sn) TType EGlobal) t