content: 
  {-# language LambdaCase, Strict, OverloadedStrings, DerivingVia #-}
  {-# options_ghc -Wincomplete-patterns #-}
  
  {-|
  Simple universe polymorphism. Features:
  
    - Non-first-class levels: there is a distinct Pi, lambda and application for
      levels. Levels are distinct from terms in syntax/values.
    - The surface language only allows abstraction over finite levels. Internally,
      the type of level-polymorphic functions is (U omega), but that's not
      accessible to programmers.
    - Bidirectional elaboration.
    - We have zero, suc, maximum and bound variables in finite levels.
    - Pi returns in (U (max i j)).
    - Conversion checking for finite levels is complete w.r.t. the obvious
      semantics for the supported operations. For example, (suc (max x (max y z)))
      is convertible to (max (suc y) (max (suc z) (suc x))).
  
  This implementation uses a simplification, where it does not thread levels in infer/check,
  and uses a heterogeneous conversion function where sides can differ in levels; concretely,
  if sides are types, their levels may differ. This works because whenever heterogeneous
  conversion checking succeeds, it is implied that the levels of sides are equal.
  
  This trick does not work anymore if we add general metavariables to the system.
  Here, we should pass a level to `check` and return one from `infer`, and in 
  the conversion checking case in `check` we should first unify expected & inferred levels,
  and then unify expected & inferred types.
  
  -}
  
  import qualified Data.IntMap.Strict as M
  
  import Data.Foldable
  import Data.Maybe
  import Data.String
  import Debug.Trace
  import Data.Coerce
  
  --------------------------------------------------------------------------------
  
  -- | De Bruijn indices.
  newtype Ix  = Ix Int deriving (Eq, Show, Ord, Num) via Int
  
  -- | De Bruijn levels.
  newtype Lvl = Lvl Int deriving (Eq, Show, Ord, Num) via Int
  
  lvlToIx :: Lvl -> Lvl -> Ix
  lvlToIx l x = coerce (l - x - 1)
  
  -- Universe levels
  --------------------------------------------------------------------------------
  
  data FinLevel
    = FLVar Ix
    | FLSuc FinLevel
    | FLZero
    | FLMax FinLevel FinLevel
    deriving Show
  
  data Level
    = Omega
    | FinLevel FinLevel
    deriving Show
  
  -- | Finite levels are of the form `maximum [n₀, x₁ + n₁, x₂ + n₂, ... xᵢ + nᵢ]`,
  -- | where the `n₀` is the first `Int` field, and the `M.Map` maps zero or more xᵢ rigid
  --   level variables to the nᵢ values.
  data VFinLevel = VFL Int (M.IntMap Int)
    deriving (Eq, Show)
  
  -- | Universe levels.
  data VLevel
    -- | The limit of all finite levels.
    = VOmega
    | VFinLevel {-# unpack #-} VFinLevel
    deriving (Eq, Show)
  
  instance Semigroup VFinLevel where
    VFL n xs <> VFL n' xs' = VFL (max n n') (M.unionWith max xs xs')
  
  instance Monoid VFinLevel where
    mempty = VFL 0 mempty
  
  instance Semigroup VLevel where
    VFinLevel i <> VFinLevel j = VFinLevel (i <> j)
    _           <> _           = VOmega
  
  instance Monoid VLevel where
    mempty = VFinLevel mempty
  
  vFinLevelVar :: Lvl -> VFinLevel
  vFinLevelVar x = VFL 0 (M.singleton (coerce x) 0)
  
  addToVFinLevel :: Int -> VFinLevel -> VFinLevel
  addToVFinLevel n (VFL m xs) = VFL (n + m) ((+n) <$> xs)
  
  addToFinLevel :: Int -> FinLevel -> FinLevel
  addToFinLevel n i = go n i where
    go 0 i = i
    go n i = go (n - 1) (FLSuc i)
  
  vSuc :: VFinLevel -> VFinLevel
  vSuc = addToVFinLevel 1
  
  -- Raw syntax
  --------------------------------------------------------------------------------
  
  type Name = String
  type RTy  = RTm
  
  data RLevel
    = RLVar Name
    | RLSuc RLevel          -- ^ suc i
    | RLZero                -- ^ 0
    | RLMax RLevel RLevel   -- ^ i ⊔ j
    deriving Show
  
  data RTm
    = RVar Name
    | RApp RTm RTm          -- ^ t u
    | RLvlApp RTm RLevel    -- ^ t @i
    | RLam Name RTm         -- ^ λ x. t
    | RLvlLam Name RTm      -- ^ λ (i : Level). t
    | RPi Name RTm RTm      -- ^ (x : A)     → B
    | RLvlPi Name RTm       -- ^ (i : Level) → B
    | RLet Name RTy RTm RTm -- ^ let x : A = t; u
    | RU RLevel             -- ^ U i
    deriving Show
  
  --------------------------------------------------------------------------------
  
  type Ty = Tm
  
  data Tm
    = Var Ix
    | App Tm Tm
    | LvlApp Tm FinLevel
    | Lam Name Tm
    | LvlLam Name Tm
    | Pi Name Tm Tm
    | LvlPi Name Tm
    | Let Name Ty Tm Tm
    | U Level
    deriving Show
  
  type VTy = Val
  
  data Val
    = VVar Lvl
    | VApp Val ~Val
    | VLvlApp Val VFinLevel
    | VLam Name (Val -> Val)
    | VLvlLam Name (VFinLevel -> Val)
    | VPi Name Val (Val -> Val)
    | VLvlPi Name (VFinLevel -> Val)
    | VU VLevel
  
  data Env = ENil | EVal Env ~Val | EFinLevel Env {-# unpack #-} VFinLevel
  
  levelVar :: Env -> Ix -> VFinLevel
  levelVar (EFinLevel _ i) 0 = i
  levelVar (EFinLevel e i) x = levelVar e (x - 1)
  levelVar (EVal _ _     ) 0 = undefined
  levelVar (EVal e _     ) x = levelVar e (x - 1)
  levelVar ENil            _ = undefined
  
  var :: Env -> Ix -> Val
  var (EFinLevel _ i) 0 = undefined
  var (EFinLevel e i) x = var e (x - 1)
  var (EVal _ t     ) 0 = t
  var (EVal e _     ) x = var e (x - 1)
  var ENil            _ = undefined
  
  finLevel :: Env -> FinLevel -> VFinLevel
  finLevel e = \case
    FLZero    -> mempty
    FLSuc i   -> addToVFinLevel 1 (finLevel e i)
    FLMax i j -> finLevel e i <> finLevel e j
    FLVar x   -> levelVar e x
  
  level :: Env -> Level -> VLevel
  level e = \case
    Omega       -> VOmega
    FinLevel i  -> VFinLevel (finLevel e i)
  
  eval :: Env -> Tm -> Val
  eval e = \case
    Var x       -> var e x
    App t u     -> case (eval e t, eval e u) of
                     (VLam _ t, u) -> t u
                     (t, u)        -> VApp t u
    LvlApp t u  -> case (eval e t, finLevel e u) of
                     (VLvlLam _ t, u) -> t u
                     (t, u)           -> VLvlApp t u
    Lam x t     -> VLam x (\u -> eval (EVal e u) t)
    LvlLam x t  -> VLvlLam x (\i -> eval (EFinLevel e i) t)
    Pi x a b    -> VPi x (eval e a) (\u -> eval (EVal e u) b)
    LvlPi x b   -> VLvlPi x (\i -> eval (EFinLevel e i) b)
    Let x a t u -> eval (EVal e (eval e t)) u
    U i         -> VU (level e i)
  
  conv :: Lvl -> Val -> Val -> Bool
  conv l t t' = case (t, t') of
    (VVar x      , VVar x'     ) -> x == x'
    (VApp t u    , VApp t' u'  ) -> conv l t t' && conv l u u'
    (VPi x a b   , VPi _ a' b' ) -> conv l a a' && conv (l + 1) (b (VVar l)) (b' (VVar l))
    (VLvlPi x b  , VLvlPi _ b' ) -> conv (l + 1) (b (vFinLevelVar l)) (b' (vFinLevelVar l))
    (VU i        , VU i'       ) -> i == i'
    (VLam x t    , VLam _ t'   ) -> conv (l + 1) (t (VVar l)) (t' (VVar l))
    (VLam x t    , t'          ) -> conv (l + 1) (t (VVar l)) (VApp t' (VVar l))
    (t           , VLam x t'   ) -> conv (l + 1) (VApp t (VVar l)) (t' (VVar l))
    (VLvlLam x t , VLvlLam _ t') -> conv (l + 1) (t (vFinLevelVar l)) (t' (vFinLevelVar l))
    (VLvlLam x t , t'          ) -> conv (l + 1) (t (vFinLevelVar l)) (VLvlApp t' (vFinLevelVar l))
    (t           , VLvlLam x t') -> conv (l + 1) (VLvlApp t (vFinLevelVar l)) (t' (vFinLevelVar l))
    _                            -> False
  
  quoteFinLevel :: Lvl -> VFinLevel -> FinLevel
  quoteFinLevel l (VFL n xs) =
    M.foldlWithKey
      (\i x n -> FLMax i (addToFinLevel n (FLVar (lvlToIx l (Lvl x)))))
      (addToFinLevel n FLZero)
      xs
  
  quoteLevel :: Lvl -> VLevel -> Level
  quoteLevel l = \case
    VOmega      -> Omega
    VFinLevel i -> FinLevel (quoteFinLevel l i)
  
  quote :: Lvl -> Val -> Tm
  quote l = \case
    VVar x      -> Var (lvlToIx l x)
    VApp t u    -> App (quote l t) (quote l u)
    VLvlApp t i -> LvlApp (quote l t) (quoteFinLevel l i)
    VLam x t    -> Lam x (quote (l + 1) (t (VVar l)))
    VLvlLam x t -> LvlLam x (quote (l + 1) (t (vFinLevelVar l)))
    VPi x a b   -> Pi x (quote l a) (quote (l + 1) (b (VVar l)))
    VLvlPi x b  -> LvlPi x (quote (l + 1) (b (vFinLevelVar l)))
    VU i        -> U (quoteLevel l i)
  
  nf :: Tm -> Tm
  nf = quote 0 . eval ENil
  
  --------------------------------------------------------------------------------
  
  type M = Either String
  
  data Binders = BNil | BType Binders Name ~VTy | BLevel Binders Name
  
  data Cxt = Cxt {env :: Env, binders :: Binders, lvl :: Lvl}
  
  define :: Name -> VTy -> Val -> Cxt -> Cxt
  define x a ~t (Cxt e bs l) = Cxt (EVal e t) (BType bs x a) (l + 1)
  
  bind :: Name -> VTy -> Cxt -> Cxt
  bind x a (Cxt e bs l) = Cxt (EVal e (VVar l)) (BType bs x a) (l + 1)
  
  bindLevel :: Name -> Cxt -> Cxt
  bindLevel x (Cxt e bs l) =
    Cxt (EFinLevel e (vFinLevelVar l)) (BLevel bs x) (l + 1)
  
  check :: Cxt -> RTm -> VTy -> M Tm
  check cxt t a = case (t, a) of
  
    (RLam x t, VPi x' a b) -> do
      Lam x <$> check (bind x a cxt) t (b (VVar (lvl cxt)))
  
    (RLvlLam x t, VLvlPi x' b) -> do
      LvlLam x <$> check (bindLevel x cxt) t (b (vFinLevelVar (lvl cxt)))
  
    (RLet x a t u, b) -> do
      (a, _) <- checkTy cxt a
      let ~va = eval (env cxt) a
      u <- check cxt t va
      t <- check (define x va (eval (env cxt) u) cxt) t b
      pure (Let x a t u)
  
    (t, a) -> do
      (t, a') <- infer cxt t
      if conv (lvl cxt) a a'
        then pure t
        else Left $ "Type mismatch, expected\n\n    "
              ++ show (quote (lvl cxt) a)
              ++ "\n\ninferred\n\n    " ++ show (quote (lvl cxt) a')
  
  checkTy :: Cxt -> RTm -> M (Tm, VLevel)
  checkTy cxt t = do
    (t, a) <- infer cxt t
    case a of
      VU l -> pure (t, l)
      _    -> Left "expected a type"
  
  inferLevelVar :: Cxt -> Name -> M Lvl
  inferLevelVar (Cxt _ bs l) x = go l bs where
    go l BNil                       = Left $ "Name not in scope: " ++ x
    go l (BType bs _ _)             = go (l - 1) bs
    go l (BLevel bs x') | x == x'   = pure $! l - 1
                        | otherwise = go (l - 1) bs
  
  inferVar :: Cxt -> Name -> M (Lvl, VTy)
  inferVar (Cxt _ bs l) x = go l bs where
    go l BNil                        = Left $ "Name not in scope: " ++ x
    go l (BLevel vs _)               = go (l - 1) bs
    go l (BType bs x' a) | x == x'   = pure (l - 1, a)
                         | otherwise = go (l - 1) bs
  
  checkFinLevel :: Cxt -> RLevel -> M FinLevel
  checkFinLevel cxt = \case
    RLVar x   -> do x <- inferLevelVar cxt x
                    pure (FLVar (lvlToIx (lvl cxt) x))
    RLSuc i   -> FLSuc <$> checkFinLevel cxt i
    RLZero    -> pure FLZero
    RLMax i j -> FLMax <$> checkFinLevel cxt i <*> checkFinLevel cxt j
  
  infer :: Cxt -> RTm -> M (Tm, VTy)
  infer cxt = \case
  
    RVar x -> do
      (x, a) <- inferVar cxt x
      pure (Var (lvlToIx (lvl cxt) x), a)
  
    RApp t u -> do
      (t, a) <- infer cxt t
      case a of
        VPi x a b -> do
          u <- check cxt u a
          pure (App t u, b (eval (env cxt) u))
        _ ->
          Left "expected a function"
  
    RLvlApp t i -> do
      (t, a) <- infer cxt t
      case a of
        VLvlPi x b -> do
          i <- checkFinLevel cxt i
          pure (LvlApp t i, b (finLevel (env cxt) i))
        _ ->
          Left "expected a level-polymorphic function"
  
    RLvlPi x b -> do
      (b, bl) <- checkTy (bindLevel x cxt) b
      pure (LvlPi x b, VU VOmega)
  
    RLam{} ->
      Left "can't infer type for lambda"
  
    RLvlLam{} ->
      Left "can't infer type for lambda"
  
    RPi x a b -> do
      (a, al) <- checkTy cxt a
      (b, bl) <- checkTy (bind x (eval (env cxt) a) cxt) b
      pure (Pi x a b, VU (al <> bl))
  
    RLet x a t u -> do
      (a, al) <- checkTy cxt a
      let ~va = eval (env cxt) a
      t <- check cxt t va
      (u, b) <- infer (define x va (eval (env cxt) t) cxt) u
      pure (Let x a t u, b)
  
    RU i -> do
      i <- checkFinLevel cxt i
      pure (U (FinLevel i), VU (VFinLevel (vSuc (finLevel (env cxt) i))))
  
  
  elab :: RTm -> IO ()
  elab t = do
    case infer (Cxt ENil BNil 0) t of
      Left err -> putStrLn err
      Right (t, a) -> do
        putStrLn "---- term ----"
        print t
        putStrLn "---- nf   ----"
        print $ nf t
        putStrLn "---- type ----"
        print $ quote 0 a
  
  
  instance IsString RTm where fromString = RVar
  instance IsString RLevel where fromString = RLVar
  
  infixl 7 $$
  infixl 7 $$$
  infixr 4 ==>
  ($$)  = RApp
  ($$$) = RLvlApp
  (==>) = RPi "_"
  
  --------------------------------------------------------------------------------
  
  -- elab p1
  p1 =
    RLet "id" (RLvlPi "i" $ RPi "A" (RU "i") $ "A" ==> "A")
              (RLvlLam "i" $ RLam "A" $ RLam "x" "x") $
    "id" $$$ RLSuc (RLSuc RLZero) $$ RU (RLSuc RLZero) $$ RU RLZero
  
  p2 = RLvlPi "i" $ RPi "A" (RU "i") $ "A" ==> "A"
  
  
  u0 = VU (VFinLevel (VFL 0 M.empty))
  u1 = VU (VFinLevel (VFL 1 M.empty))
  
  foo =  check (Cxt ENil BNil 0)
               (RLam "x" (RVar "x"))
               (VPi "x" (VPi "y" u1 (const u0)) (const (VPi "y" u0 (const u1))))
  