content: {-# LANGUAGE TupleSections #-}
  
  -- | Convert the concrete syntax into the syntax of cubical TT.
  module Resolver where
  
  import Control.Applicative
  import Control.Monad
  import Control.Monad.Reader
  import Control.Monad.Except
  import Control.Monad.Identity
  import Data.Maybe
  import Data.List
  import Data.Map (Map,(!))
  import qualified Data.Map as Map
  
  import Exp.Abs
  import CTT (Ter,Ident,Loc(..),mkApps,mkWheres)
  import qualified CTT
  import Connections (negFormula,andFormula,orFormula)
  import qualified Connections as C
  
  -- | Useful auxiliary functions
  
  -- Applicative cons
  (<:>) :: Applicative f => f a -> f [a] -> f [a]
  a <:> b = (:) <$> a <*> b
  
  -- Un-something functions
  unVar :: Exp -> Maybe Ident
  unVar (Var (AIdent (_,x))) = Just x
  unVar _                    = Nothing
  
  unWhere :: ExpWhere -> Exp
  unWhere (Where e ds) = Let ds e
  unWhere (NoWhere e)  = e
  
  -- Tail recursive form to transform a sequence of applications
  -- App (App (App u v) ...) w  into (u, [v, …, w])
  -- (cleaner than the previous version of unApps)
  unApps :: Exp -> [Exp] -> (Exp, [Exp])
  unApps (App u v) ws = unApps u (v : ws)
  unApps u         ws = (u, ws)
  
  -- Turns an expression of the form App (... (App id1 id2) ... idn)
  -- into a list of idents
  appsToIdents :: Exp -> Maybe [Ident]
  appsToIdents = mapM unVar . uncurry (:) . flip unApps []
  
  -- Transform a sequence of applications
  -- (((u v1) .. vn) phi1) .. phim into (u,[v1,..,vn],[phi1,..,phim])
  unAppsFormulas :: Exp -> [Formula]-> (Exp,[Exp],[Formula])
  unAppsFormulas (AppFormula u phi) phis = unAppsFormulas u (phi:phis)
  unAppsFormulas u phis = (x,xs,phis)
    where (x,xs) = unApps u []
  
  -- Flatten a tele
  flattenTele :: [Tele] -> [(Ident,Exp)]
  flattenTele tele =
    [ (unAIdent i,typ) | Tele id ids typ <- tele, i <- id:ids ]
  
  -- Flatten a PTele
  flattenPTele :: [PTele] -> Resolver [(Ident,Exp)]
  flattenPTele []                   = return []
  flattenPTele (PTele exp typ : xs) = case appsToIdents exp of
    Just ids -> do
      pt <- flattenPTele xs
      return $ map (,typ) ids ++ pt
    Nothing -> throwError "malformed ptele"
  
  -------------------------------------------------------------------------------
  -- | Resolver and environment
  
  data SymKind = Variable | Constructor | PConstructor | Name
    deriving (Eq,Show)
  
  -- local environment for constructors
  data Env = Env { envModule :: String,
                   variables :: [(Ident,SymKind)] }
    deriving (Eq,Show)
  
  type Resolver a = ReaderT Env (ExceptT String Identity) a
  
  emptyEnv :: Env
  emptyEnv = Env "" []
  
  runResolver :: Resolver a -> Either String a
  runResolver x = runIdentity $ runExceptT $ runReaderT x emptyEnv
  
  updateModule :: String -> Env -> Env
  updateModule mod e = e{envModule = mod}
  
  insertIdent :: (Ident,SymKind) -> Env -> Env
  insertIdent (n,var) e
    | n == "_"  = e
    | otherwise = e{variables = (n,var) : variables e}
  
  insertIdents :: [(Ident,SymKind)] -> Env -> Env
  insertIdents = flip $ foldr insertIdent
  
  insertName :: AIdent -> Env -> Env
  insertName (AIdent (_,x)) = insertIdent (x,Name)
  
  insertNames :: [AIdent] -> Env -> Env
  insertNames = flip $ foldr insertName
  
  insertVar :: Ident -> Env -> Env
  insertVar x = insertIdent (x,Variable)
  
  insertVars :: [Ident] -> Env -> Env
  insertVars = flip $ foldr insertVar
  
  insertAIdent :: AIdent -> Env -> Env
  insertAIdent (AIdent (_,x)) = insertIdent (x,Variable)
  
  insertAIdents :: [AIdent] -> Env -> Env
  insertAIdents  = flip $ foldr insertAIdent
  
  getLoc :: (Int,Int) -> Resolver Loc
  getLoc l = Loc <$> asks envModule <*> pure l
  
  unAIdent :: AIdent -> Ident
  unAIdent (AIdent (_,x)) = x
  
  resolveName :: AIdent -> Resolver C.Name
  resolveName (AIdent (l,x)) = do
    modName <- asks envModule
    vars <- asks variables
    case lookup x vars of
      Just Name -> return $ C.Name x
      _ -> throwError $ "Cannot resolve name " ++ x ++ " at position " ++
                        show l ++ " in module " ++ modName
  
  resolveVar :: AIdent -> Resolver Ter
  resolveVar (AIdent (l,x)) = do
    modName <- asks envModule
    vars    <- asks variables
    case lookup x vars of
      Just Variable    -> return $ CTT.Var x
      Just Constructor -> return $ CTT.Con x []
      Just PConstructor ->
        throwError $ "The path constructor " ++ x ++ " is used as a" ++
                     " variable at " ++ show l ++ " in " ++ modName ++
                     " (path constructors should have their type in" ++
                     " curly braces as first argument)"
      Just Name        ->
        throwError $ "Name " ++ x ++ " used as a variable at position " ++
                     show l ++ " in module " ++ modName
      _ -> throwError $ "Cannot resolve variable " ++ x ++ " at position " ++
                        show l ++ " in module " ++ modName
  
  lam :: (Ident,Exp) -> Resolver Ter -> Resolver Ter
  lam (a,t) e = CTT.Lam a <$> resolveExp t <*> local (insertVar a) e
  
  lams :: [(Ident,Exp)] -> Resolver Ter -> Resolver Ter
  lams = flip $ foldr lam
  
  plam :: AIdent -> Resolver Ter -> Resolver Ter
  plam i e = CTT.PLam (C.Name (unAIdent i)) <$> local (insertName i) e
  
  plams :: [AIdent] -> Resolver Ter -> Resolver Ter
  plams [] _ = throwError "Empty plam abstraction"
  plams xs e = foldr plam e xs
  
  bind :: (Ter -> Ter) -> (Ident,Exp) -> Resolver Ter -> Resolver Ter
  bind f (x,t) e = f <$> lam (x,t) e
  
  binds :: (Ter -> Ter) -> [(Ident,Exp)] -> Resolver Ter -> Resolver Ter
  binds f = flip $ foldr $ bind f
  
  resolveApps :: Exp -> [Exp] -> Resolver Ter
  resolveApps x xs = mkApps <$> resolveExp x <*> mapM resolveExp xs
  
  resolveExp :: Exp -> Resolver Ter
  resolveExp e = case e of
    U             -> return CTT.U
    Var x         -> resolveVar x
    App t s       -> resolveApps x xs
      where (x,xs) = unApps t [s]
    Sigma ptele b -> do
      tele <- flattenPTele ptele
      binds CTT.Sigma tele (resolveExp b)
    Pi ptele b    -> do
      tele <- flattenPTele ptele
      binds CTT.Pi tele (resolveExp b)
    Fun a b       -> bind CTT.Pi ("_",a) (resolveExp b)
    Lam ptele t   -> do
      tele <- flattenPTele ptele
      lams tele (resolveExp t)
    Fst t         -> CTT.Fst <$> resolveExp t
    Snd t         -> CTT.Snd <$> resolveExp t
    Pair t0 ts    -> do
      e  <- resolveExp t0
      es <- mapM resolveExp ts
      return $ foldr1 CTT.Pair (e:es)
    Split t brs -> do
      t'   <- resolveExp t
      brs' <- mapM resolveBranch brs
      l@(Loc n (i,j)) <- getLoc (case brs of
                                    OBranch (AIdent (l,_)) _ _:_ -> l
                                    PBranch (AIdent (l,_)) _ _ _:_ -> l
                                    _ -> (0,0))
      return $ CTT.Split (n ++ "_L" ++ show i ++ "_C" ++ show j) l t' brs'
    Let decls e   -> do
      (rdecls,names) <- resolveDecls decls
      mkWheres rdecls <$> local (insertIdents names) (resolveExp e)
    PLam is e     -> plams is (resolveExp e)
    Hole (HoleIdent (l,_)) -> CTT.Hole <$> getLoc l
    AppFormula t phi ->
      let (x,xs,phis) = unAppsFormulas e []
      in case x of
        PCon n a ->
          CTT.PCon (unAIdent n) <$> resolveExp a <*> mapM resolveExp xs
                                <*> mapM resolveFormula phis
        _ -> CTT.AppFormula <$> resolveExp t <*> resolveFormula phi
    PathP a u v   -> CTT.PathP <$> resolveExp a <*> resolveExp u <*> resolveExp v
    Comp u v ts   -> CTT.Comp <$> resolveExp u <*> resolveExp v <*> resolveSystem ts
    HComp u v ts  -> CTT.HComp <$> resolveExp u <*> resolveExp v <*> resolveSystem ts
    Fill u v ts   -> CTT.Fill <$> resolveExp u <*> resolveExp v <*> resolveSystem ts
    Trans u v     -> CTT.Comp <$> resolveExp u <*> resolveExp v <*> pure Map.empty
    Glue u ts     -> CTT.Glue <$> resolveExp u <*> resolveSystem ts
    GlueElem u ts -> CTT.GlueElem <$> resolveExp u <*> resolveSystem ts
    UnGlueElem u ts -> CTT.UnGlueElem <$> resolveExp u <*> resolveSystem ts
    Id a u v      -> CTT.Id <$> resolveExp a <*> resolveExp u <*> resolveExp v
    IdPair u ts   -> CTT.IdPair <$> resolveExp u <*> resolveSystem ts
    IdJ a t c d x p -> CTT.IdJ <$> resolveExp a <*> resolveExp t <*> resolveExp c
                               <*> resolveExp d <*> resolveExp x <*> resolveExp p
    _ -> do
      modName <- asks envModule
      throwError ("Could not resolve " ++ show e ++ " in module " ++ modName)
  
  resolveWhere :: ExpWhere -> Resolver Ter
  resolveWhere = resolveExp . unWhere
  
  resolveSystem :: System -> Resolver (C.System Ter)
  resolveSystem (System ts) = do
    ts' <- sequence [ (,) <$> resolveFace alpha <*> resolveExp u
                    | Side alpha u <- ts ]
    let alphas = map fst ts'
    unless (nub alphas == alphas) $
      throwError $ "system contains same face multiple times: " ++
                   C.showListSystem ts'
    -- Note: the symbols in alpha are in scope in u, but they mean 0 or 1
    return $ Map.fromList ts'
  
  resolveFace :: [Face] -> Resolver C.Face
  resolveFace alpha =
    Map.fromList <$> sequence [ (,) <$> resolveName i <*> resolveDir d
                              | Face i d <- alpha ]
  
  resolveDir :: Dir -> Resolver C.Dir
  resolveDir Dir0 = return 0
  resolveDir Dir1 = return 1
  
  resolveFormula :: Formula -> Resolver C.Formula
  resolveFormula (Dir d)          = C.Dir <$> resolveDir d
  resolveFormula (Atom i)         = C.Atom <$> resolveName i
  resolveFormula (Neg phi)        = negFormula <$> resolveFormula phi
  resolveFormula (Conj phi _ psi) =
      andFormula <$> resolveFormula phi <*> resolveFormula psi
  resolveFormula (Disj phi psi)   =
      orFormula <$> resolveFormula phi <*> resolveFormula psi
  
  resolveBranch :: Branch -> Resolver CTT.Branch
  resolveBranch (OBranch (AIdent (_,lbl)) args e) = do
    re <- local (insertAIdents args) $ resolveWhere e
    return $ CTT.OBranch lbl (map unAIdent args) re
  resolveBranch (PBranch (AIdent (_,lbl)) args is e) = do
    re <- local (insertNames is . insertAIdents args) $ resolveWhere e
    let names = map (C.Name . unAIdent) is
    return $ CTT.PBranch lbl (map unAIdent args) names re
  
  resolveTele :: [(Ident,Exp)] -> Resolver CTT.Tele
  resolveTele []        = return []
  resolveTele ((i,d):t) =
    ((i,) <$> resolveExp d) <:> local (insertVar i) (resolveTele t)
  
  resolveLabel :: [(Ident,SymKind)] -> Label -> Resolver CTT.Label
  resolveLabel _ (OLabel n vdecl) =
    CTT.OLabel (unAIdent n) <$> resolveTele (flattenTele vdecl)
  resolveLabel cs (PLabel n vdecl is sys) = do
    let tele' = flattenTele vdecl
        ts    = map fst tele'
        names = map (C.Name . unAIdent) is
        n'    = unAIdent n
        cs'   = delete (n',PConstructor) cs
    CTT.PLabel n' <$> resolveTele tele' <*> pure names
                  <*> local (insertNames is . insertIdents cs' . insertVars ts)
                        (resolveSystem sys)
  
  -- Resolve a non-mutual declaration; returns resolver for type and
  -- body separately
  resolveNonMutualDecl :: Decl -> (Ident,Resolver CTT.Ter
                                  ,Resolver CTT.Ter,[(Ident,SymKind)])
  resolveNonMutualDecl d = case d of
    DeclDef (AIdent (_,f)) tele t body ->
      let tele' = flattenTele tele
          a     = binds CTT.Pi tele' (resolveExp t)
          d     = lams tele' (local (insertVar f) $ resolveWhere body)
      in (f,a,d,[(f,Variable)])
    DeclData x tele sums  -> resolveDeclData x tele sums null
    DeclHData x tele sums ->
      resolveDeclData x tele sums (const False) -- always pick HSum
    DeclSplit (AIdent (l,f)) tele t brs ->
      let tele' = flattenTele tele
          vars  = map fst tele'
          a     = binds CTT.Pi tele' (resolveExp t)
          d     = do
                    loc <- getLoc l
                    ty  <- local (insertVars vars) $ resolveExp t
                    brs' <- local (insertVars (f:vars)) (mapM resolveBranch brs)
                    lams tele' (return $ CTT.Split f loc ty brs')
      in (f,a,d,[(f,Variable)])
    DeclUndef (AIdent (l,f)) tele t ->
      let tele' = flattenTele tele
          a     = binds CTT.Pi tele' (resolveExp t)
          d     = CTT.Undef <$> getLoc l <*> a
      in (f,a,d,[(f,Variable)])
  
  -- Helper function to resolve data declarations. The predicate p is
  -- used to decide if we should use Sum or HSum.
  resolveDeclData :: AIdent -> [Tele] -> [Label] -> ([(Ident,SymKind)] -> Bool) ->
                     (Ident, Resolver Ter, Resolver Ter, [(Ident, SymKind)])
  resolveDeclData (AIdent (l,f)) tele sums p =
    let tele' = flattenTele tele
        a     = binds CTT.Pi tele' (return CTT.U)
        cs    = [ (unAIdent lbl,Constructor) | OLabel lbl _ <- sums ]
        pcs   = [ (unAIdent lbl,PConstructor) | PLabel lbl _ _ _ <- sums ]
        sum   = if p pcs then CTT.Sum else CTT.HSum
        d = lams tele' $ local (insertVar f) $
              sum <$> getLoc l <*> pure f
                  <*> mapM (resolveLabel (cs ++ pcs)) sums
    in (f,a,d,(f,Variable):cs ++ pcs)
  
  resolveRTele :: [Ident] -> [Resolver CTT.Ter] -> Resolver CTT.Tele
  resolveRTele [] _ = return []
  resolveRTele (i:is) (t:ts) = do
    a  <- t
    as <- local (insertVar i) (resolveRTele is ts)
    return ((i,a):as)
  
  -- Best effort to find the location of a declaration. This implementation
  -- returns the location of the first identifier it contains.
  findDeclLoc :: Decl -> Resolver Loc
  findDeclLoc d = getLoc loc
      where loc = fromMaybe (-1, 0) $ mloc d
            mloc d = case d of
              DeclDef (AIdent (l, _)) _ _ _   -> Just l
              DeclData (AIdent (l, _)) _ _    -> Just l
              DeclHData (AIdent (l, _)) _ _   -> Just l
              DeclSplit (AIdent (l, _)) _ _ _ -> Just l
              DeclUndef (AIdent (l, _)) _ _   -> Just l
              DeclMutual ds                   -> listToMaybe $ mapMaybe mloc ds
              DeclOpaque (AIdent (l, _))      -> Just l
              DeclTransparent (AIdent (l, _)) -> Just l
              DeclTransparentAll              -> Nothing
  
  -- Resolve a declaration
  resolveDecl :: Decl -> Resolver (CTT.Decls,[(Ident,SymKind)])
  resolveDecl d = case d of
    DeclMutual decls -> do
      let (fs,ts,bs,nss) = unzip4 $ map resolveNonMutualDecl decls
          ns = concat nss -- TODO: some sanity checks? Duplicates!?
      when (nub (map fst ns) /= concatMap (map fst) nss) $
        throwError ("Duplicated constructor or ident: " ++ show nss)
      as <- resolveRTele fs ts
      -- The bodies know about all the names and constructors in the
      -- mutual block
      ds <- sequence $ map (local (insertIdents ns)) bs
      let ads = zipWith (\ (x,y) z -> (x,(y,z))) as ds
      l <- findDeclLoc d
      return (CTT.MutualDecls l ads,ns)
    DeclOpaque i  -> do
      resolveVar i
      return (CTT.OpaqueDecl (unAIdent i), [])
    DeclTransparent i -> do
      resolveVar i
      return (CTT.TransparentDecl (unAIdent i), [])
    DeclTransparentAll -> return (CTT.TransparentAllDecl, [])
    _ -> do let (f,typ,body,ns) = resolveNonMutualDecl d
            l <- findDeclLoc d
            a <- typ
            d <- body
            return (CTT.MutualDecls l [(f,(a,d))],ns)
  
  resolveDecls :: [Decl] -> Resolver ([CTT.Decls],[(Ident,SymKind)])
  resolveDecls []     = return ([],[])
  resolveDecls (d:ds) = do
    (rtd,names)  <- resolveDecl d
    (rds,names') <- local (insertIdents names) $ resolveDecls ds
    return (rtd : rds, names' ++ names)
  
  resolveModule :: Module -> Resolver ([CTT.Decls],[(Ident,SymKind)])
  resolveModule (Module (AIdent (_,n)) _ decls) =
    local (updateModule n) $ resolveDecls decls
  
  resolveModules :: [Module] -> Resolver ([CTT.Decls],[(Ident,SymKind)])
  resolveModules []         = return ([],[])
  resolveModules (mod:mods) = do
    (rmod, names)  <- resolveModule mod
    (rmods,names') <- local (insertIdents names) $ resolveModules mods
    return (rmod ++ rmods, names' ++ names)
  