content: -- Pure-type system-like example
  -- Includes Pi/Sigma, untyped equivalence
  module PTS where
  
  import Rebound
  
  import qualified Rebound.Bind.Pat as Pat
  import Rebound.Bind.PatN as PatN
  import Rebound.Context
  import Control.Monad.Except (ExceptT, MonadError (..), runExceptT)
  import Data.Fin(Fin(..), f0,f1,f2)
  import Data.Fin qualified as Fin
  import Data.Vec qualified as Vec
  import GHC.Generics (Generic1)
  
  -- In a pure type system, terms and types are combined
  -- into the same syntactic class.
  
  data Exp (n :: Nat) where
    -- | sort
    Star :: Exp n
    -- | dependent type `Pi x : A . B`
    Pi :: Exp n -> Bind1 Exp Exp n -> Exp n
    -- | variable
    Var :: Fin n -> Exp n
    -- | lambda expression, with type annotation  `lambda x:A.B`
    Lam :: Exp n -> Bind1 Exp Exp n -> Exp n
    -- | application
    App :: Exp n -> Exp n -> Exp n
    -- | dependent pair `Sigma x:A . B`
    Sigma :: Exp n -> Bind1 Exp Exp n -> Exp n
    -- | construct a pair, third argument is type annotation
    Pair :: Exp n -> Exp n -> Exp n -> Exp n
    -- | elimination form for pairs. `split e1 as (x,y) in e2`
    -- Binds two variables to
    -- the two components of the pair
    Split :: Exp n -> Bind2 Exp Exp n -> Exp n
  
  deriving instance (Generic1 Exp)
  ----------------------------------------------
  
  instance SubstVar Exp where
    var :: Fin n -> Exp n
    var = Var
  
  instance Shiftable Exp where
    shift = shiftFromApplyE @Exp
  
  instance Subst Exp Exp where
    isVar (Var x) = Just (Refl, x)
    isVar _ = Nothing
  ----------------------------------------------
  
  t00 :: Exp N2
  t00 = App (Var f0) (Var f0)
  
  t01 :: Exp N2
  t01 = App (Var f0) (Var f1)
  
  -- Does a variable appear free in a term?
  
  -- >>> appearsFree f0 t00
  -- True
  
  -- >>> appearsFree f1 t00
  -- False
  instance FV Exp where
  
  -- >>> :t weaken' s1 t00
  -- weaken' s1 t00 :: Exp ('S ('S N1))
  
  -- >>> weaken' s1 t00
  -- 0 0
  
  weaken' :: SNat m -> Exp n -> Exp (m + n)
  weaken' m = applyE @Exp (weakenE' m)
  
  -- >>> strengthenRec s1 s1 snat t00
  -- Just (0 0)
  
  -- >>> strengthenRec s1 s1 snat t01
  -- Nothing
  
  instance Strengthen Exp where
  
  ----------------------------------------------
  -- Examples
  
  -- The identity function "λ x. x". With de Bruijn indices
  -- we write it as "λ. 0"
  t0 :: Exp Z
  t0 = Lam Star (bind1 (Var f0))
  
  -- A larger term "λ x. λy. x (λ z. z z)"
  -- λ. λ. 1 (λ. 0 0)
  t1 :: Exp Z
  t1 =
    Lam
      Star
      ( bind1
          ( Lam
              Star
              ( bind1
                  ( Var f1
                      `App` (Lam Star (bind1 (Var f0)) `App` Var f0)
                  )
              )
          )
      )
  
  -- To show lambda terms, we can write a simple recursive instance of
  -- Haskell's `Show` type class. In the case of a binder, we use the `unbind1`
  -- operation to access the body of the lambda expression.
  
  -- >>> t0
  -- λ *. 0
  
  -- >>> t1
  -- λ *. λ *. 1 ((λ *. 0) 0)
  
  -- Polymorphic identity function and its type
  
  tyid = Pi Star (bind1 (Pi (Var f0) (bind1 (Var f1))))
  
  tmid = Lam Star (bind1 (Lam (Var f0) (bind1 (Var f0))))
  
  -- >>> tyid
  -- Pi *. 0 -> 1
  
  -- >>> tmid
  -- λ *. λ 0. 0
  
  instance Show (Exp n) where
    showsPrec :: Int -> Exp n -> String -> String
    showsPrec _ Star = showString "*"
    showsPrec d (Pi a b)
      | appearsFree FZ (getBody1 b) =
          showParen (d > 9) $
            showString "Pi "
              . shows a
              . showString ". "
              . shows (getBody1 b)
      | otherwise =
          showParen (d > 9) $
            showsPrec 11 a
              . showString " -> "
              . showsPrec 9 (getBody1 b)
    showsPrec d (Sigma a b)
      | appearsFree FZ (getBody1 b) =
          showParen (d > 9) $
            showString "Sigma "
              . shows a
              . showString ". "
              . shows (getBody1 b)
      | otherwise =
          showParen (d > 9) $
            showsPrec 11 a
              . showString " * "
              . showsPrec 9 (getBody1 b)
    showsPrec _ (Var x) = shows x
    showsPrec d (App e1 e2) =
      showParen (d > 0) $
        showsPrec 10 e1
          . showString " "
          . showsPrec 11 e2
    showsPrec d (Lam t b) =
      showParen (d > 9) $
        showString "λ "
          . shows t
          . showString ". "
          . shows (getBody1 b)
    showsPrec d (Pair e1 e2 t) =
      showParen (d > 0) $
        showsPrec 10 e1
          . showString ", "
          . showsPrec 11 e2
    showsPrec d (Split t b) =
      showParen (d > 10) $
        showString "split"
          . showsPrec 10 t
          . showString " in "
          . shows (getBody2 b)
  
  -- With the instance above the derivable equality instance
  -- is alpha-equivalence
  deriving instance (Eq (Exp n))
  
  --------------------------------------------------------
  
  {- We can write the usual operations for evaluating
     lambda terms to values -}
  
  -- big-step evaluation
  
  -- >>> eval t1
  -- λ *. λ *. 1 ((λ *. 0) 0)
  
  -- >>> eval (t1 `App` t0)
  -- λ *. (λ *. 0) ((λ *. 0) 0)
  
  eval :: Exp n -> Exp n
  eval (Var x) = Var x
  eval (Lam a b) = Lam a b
  eval (App e1 e2) =
    let v = eval e2
     in case eval e1 of
          Lam a b -> eval (instantiate1 b v)
          t -> App t v
  eval Star = Star
  eval (Pi a b) = Pi a b
  eval (Sigma a b) = Sigma a b
  eval (Pair a b t) = Pair a b t
  eval (Split a b) =
    case eval a of
      Pair a1 a2 _ ->
        eval (instantiate2 b (eval a1) (eval a2))
      t -> Split t b
  
  -- small-step evaluation
  
  -- >>> step (t1 `App` t0)
  -- Just (λ *. λ *. 0 (λ *. 0 0))
  
  step :: Exp n -> Maybe (Exp n)
  step (Var x) = Nothing
  step (Lam a b) = Nothing
  step (App (Lam a b) e2) = Just (instantiate1 b e2)
  step (App e1 e2)
    | Just e1' <- step e1 = Just (App e1' e2)
    | Just e2' <- step e2 = Just (App e1 e2')
    | otherwise = Nothing
  step Star = Nothing
  step (Pi a b) = Nothing
  step (Sigma a b) = Nothing
  step (Pair a b _) = Nothing
  step (Split (Pair a1 a2 _) b) = Just (PatN.instantiate2 b a1 a2)
  step (Split a b)
    | Just a' <- step a = Just (Split a' b)
    | otherwise = Nothing
  
  eval' :: Exp n -> Exp n
  eval' e
    | Just e' <- step e = eval' e'
    | otherwise = e
  
  -- normalization
  -- to normalize under a lambda expression, we must first unbind
  -- it and then rebind it when finished
  
  -- >>> nf t1
  -- λ. λ. 1 0
  
  -- >>> nf (t1 `App` t0)
  -- λ *. 0
  
  -- reduce the term everywhere, as much as possible
  nf :: Exp n -> Exp n
  nf (Var x) = Var x
  nf (Lam a b) = Lam a (bind1 (nf (getBody1 b)))
  nf (App e1 e2) =
    case nf e1 of
      Lam a b -> nf (instantiate1 b e2)
      t -> App t (nf e2)
  nf Star = Star
  nf (Pi a b) = Pi (nf a) (bind1 (nf (getBody1 b)))
  nf (Sigma a b) = Sigma (nf a) (bind1 (nf (getBody1 b)))
  nf (Pair a b t) = Pair (nf a) (nf b) (nf t)
  nf (Split a b) =
    case nf a of
      Pair a1 a2 _ -> nf (PatN.instantiate2 b a1 a2)
      t -> Split t (PatN.bind2 (nf (getBody2 b)))
  
  -- first find the head form
  whnf :: Exp n -> Exp n
  whnf (App a1 a2) = case whnf a1 of
    Lam a b -> whnf (instantiate1 b a1)
    t -> App t a2
  whnf (Split a b) = case whnf a of
    Pair a1 a2 _ -> whnf (PatN.instantiate2 b a1 a2)
    t -> Split t b
  -- all other terms are already in head form
  whnf a = a
  
  norm :: Exp n -> Exp n
  norm a = case whnf a of
    Lam a b -> Lam (norm a) (bind1 (norm (getBody1 b)))
    Pi a b -> Pi (norm a) (bind1 (norm (getBody1 b)))
    Sigma a b -> Sigma (norm a) (bind1 (norm (getBody1 b)))
    Pair a b t -> Pair (norm a) (norm b) (norm t)
    Star -> Star
    App a b -> App a (norm b)
    Split a b -> Split a (PatN.bind2 (norm (getBody2 b)))
    Var x -> Var x
  
  --------------------------------------------------------
  -- We can also write functions that manipulate the
  -- environment explicitly
  
  -- >>> evalEnv idE t1
  -- λ *. λ *. 1 ((λ *. 0) 0)
  
  -- Below, if n is 0, then this function acts like an
  -- "environment-based" bigstep evaluator. The result of
  -- evaluating a lambda expression is a closure --- the body
  -- of the lambda paired with its environment. That is exactly
  -- what the implementation of bind1 does.
  
  -- In the case of beta-reduction, the `unBindWith` operation
  -- applies its argument to the environment and subterm in the
  -- closure. In other words, this function calls `evalEnv`
  -- recursively with the saved environment and body of the lambda term.
  
  evalEnv :: Env Exp m n -> Exp m -> Exp n
  evalEnv r (Var x) = applyEnv r x
  evalEnv r (Lam a b) = applyE r (Lam a b)
  evalEnv r (App e1 e2) =
    let v = evalEnv r e2
     in case evalEnv r e1 of
          Lam a b ->
            unbindWith1 b (\r' e' -> evalEnv (v .: r') e')
          t -> App t v
  evalEnv r Star = Star
  evalEnv r (Pi a b) = applyE r (Pi a b)
  evalEnv r (Sigma a b) = applyE r (Sigma a b)
  evalEnv r (Pair a b t) = applyE r (Pair a b t)
  evalEnv r (Split a b) =
    case evalEnv r a of
      Pair a1 a2 _ ->
        PatN.unbindWith2 b ( \r' e' -> evalEnv (a1 .: (a2 .: (r' .>> r))) e')
      t -> Split t (applyE r b)
  
  ----------------------------------------------------------------
  data Err where
    Equate :: Exp n -> Exp n -> Err
    PiExpected :: Exp n -> Err
    SigmaExpected :: Exp n -> Err
    VarEscapes :: Exp n -> Err
  
  deriving instance (Show Err)
  
  equate :: (MonadError Err m) => Exp n -> Exp n -> m ()
  equate t1 t2 = do
    let n1 = whnf t1
        n2 = whnf t2
    equateWHNF n1 n2
  
  equateWHNF :: (MonadError Err m) => Exp n -> Exp n -> m ()
  equateWHNF n1 n2 =
    case (n1, n2) of
      (Star, Star) -> pure ()
      (Var x, Var y) | x == y -> pure ()
      (Lam _ b1, Lam _ b2) -> equate (getBody1 b1) (getBody1 b2)
      (App a1 a2, App b1 b2) -> do
        equateWHNF a1 b1
        equate a2 b2
      (Pi tyA1 b1, Pi tyA2 b2) -> do
        equate tyA1 tyA2
        equate (getBody1 b1) (getBody1 b2)
      (Pair a1 a2 _, Pair b1 b2 _) -> do
        equate a1 b1
        equate a2 b2
      (Split a1 b1, Split a2 b2) -> do
        equateWHNF a1 a2
        equate (getBody2 b1) (getBody2 b2)
      (Sigma tyA1 b1, Sigma tyA2 b2) -> do
        equate tyA1 tyA2
        equate (getBody1 b1) (getBody1 b2)
      (_, _) -> throwError (Equate n1 n2)
  
  ----------------------------------------------------------------
  
  checkType ::
    forall n m.
    (MonadError Err m, SNatI n) =>
    Ctx Exp n ->
    Exp n ->
    Exp n ->
    m ()
  checkType g e t1 = do
    t2 <- inferType g e
    equate (whnf t2) t1
  
  inferType ::
    forall n m.
    (MonadError Err m, SNatI n) =>
    Ctx Exp n ->
    Exp n ->
    m (Exp n)
  inferType g (Var x) = pure (applyEnv g x)
  inferType g Star = pure Star
  inferType g (Pi a b) = do
    checkType g a Star
    checkType (g +++ a) (getBody1 b) Star
    pure Star
  inferType g (Lam tyA b) = do
    checkType g tyA Star
    tyB <- inferType (g +++ tyA) (getBody1 b)
    return $ Pi tyA (bind1 tyB)
  inferType g (App a b) = do
    tyA <- inferType g a
    case whnf tyA of
      Pi tyA1 tyB1 -> do
        checkType g b tyA1
        pure $ instantiate1 tyB1 b
      t -> throwError (PiExpected t)
  inferType g (Sigma a b) = do
    checkType g a Star
    checkType (g +++ a) (getBody1 b) Star
    pure Star
  inferType g (Pair a b ty) = do
    tyA <- inferType g a
    tyB <- inferType g b
    case ty of
      (Sigma tyA tyB) -> do
        checkType g a tyA
        checkType g b (instantiate1 tyB a)
        pure ty
      _ -> throwError (SigmaExpected ty)
  inferType g (Split a b) = do
    tyA <- inferType g a
    case whnf tyA of
      Sigma tyA' tyB' -> do
        let g' :: Ctx Exp (S (S n))
            g' = g +++ tyA' +++ getBody1 tyB'
        ty <- inferType g' (getBody2 b)
        let ty' = whnf ty
        case strengthenN s2 ty' of
          Nothing -> throwError (VarEscapes ty)
          Just ty'' -> pure ty''
      _ -> throwError (SigmaExpected tyA)
  
  
  -- >>> inferType zeroE tmid :: Either Err (Exp N0)
  -- Right (Pi *. 0 -> 1)
  
  
  -- >>> inferType zeroE (App tmid tyid) :: Either Err (Exp N0)
  -- Right ((Pi *. 0 -> 1) -> Pi *. 0 -> 1)
  
  