module Core.Declarative

import Core.Environment
import Core.Environment.Extension
import Core.Term
import Core.Term.Substitution
import Core.Term.Thinned
import Core.Thinning

-- Definition ------------------------------------------------------------------

public export
data EnvWf : Env n -> Type
public export
data TypeWf : Env n -> Term n -> Type
public export
data TypeConv : Env n -> Term n -> Term n -> Type
public export
data TermWf : Env n -> Term n -> Term n -> Type
public export
data TermConv : Env n -> Term n -> Term n -> Term n -> Type

data EnvWf where
  Lin : EnvWf [<]
  (:<) : EnvWf env -> TypeWf env (expand a) -> EnvWf (env :< a)

data TypeWf where
  SetTyWf :
    EnvWf env ->
    ---
    TypeWf env Set
  PiTyWf :
    TypeWf env a ->
    TypeWf (env :< pure a) b ->
    ---
    TypeWf env (Pi a b)
  LiftWf :
    TermWf env a Set ->
    ---
    TypeWf env a

data TypeConv where
  ReflTy :
    TypeWf env a ->
    ---
    TypeConv env a a
  SymTy :
    TypeConv env a b ->
    ---
    TypeConv env b a
  TransTy :
    TypeConv env a b ->
    TypeConv env b c ->
    ---
    TypeConv env a c
  PiConv :
    TypeWf env a ->
    TypeConv env a c ->
    TypeConv (env :< pure a) b d ->
    ---
    TypeConv env (Pi a b) (Pi c d)
  LiftConv :
    TermConv env a b Set ->
    ---
    TypeConv env a b

data TermWf where
  PiTmWf :
    TermWf env a Set ->
    TermWf (env :< pure a) b Set ->
    ---
    TermWf env (Pi a b) Set
  VarWf :
    EnvWf env ->
    ---
    TermWf env (Var i) (expand $ index env i)
  AbsWf :
    TypeWf env a ->
    TermWf (env :< pure a) t b ->
    ---
    TermWf env (Abs t) (Pi a b)
  AppWf :
    TermWf env t (Pi a b) ->
    TermWf env u a ->
    ---
    TermWf env (App t u) (subst1 b u)
  ConvWf :
    TermWf env t a ->
    TypeConv env a b ->
    ---
    TermWf env t b

data TermConv where
  ReflTm :
    TermWf env t a ->
    ---
    TermConv env t t a
  SymTm :
    TermConv env t u a ->
    ---
    TermConv env u t a
  TransTm :
    TermConv env t u a ->
    TermConv env u v a ->
    ---
    TermConv env t v a
  AppConv :
    TermConv env f g (Pi a b) ->
    TermConv env t u a ->
    ---
    TermConv env (App f t) (App g u) (subst1 b t)
  PiTmConv :
    TypeWf env a ->
    TermConv env a c Set ->
    TermConv (env :< pure a) b d Set ->
    ---
    TermConv env (Pi a b) (Pi c d) Set
  PiEta :
    TypeWf env a ->
    TermWf env t (Pi a b) ->
    TermWf env u (Pi a b) ->
    TermConv (env :< pure a)
      (App (wkn t $ drop $ id _) (Var FZ))
      (App (wkn u $ drop $ id _) (Var FZ))
      b ->
    ---
    TermConv env t u (Pi a b)
  PiBeta :
    TypeWf env a ->
    TermWf (env :< pure a) t b ->
    TermWf env u a ->
    ---
    TermConv env (App (Abs t) u) (subst1 t u) (subst1 b u)
  ConvConv :
    TermConv env t u a ->
    TypeConv env a b ->
    ---
    TermConv env t u b

%name EnvWf envWf
%name TypeWf tyWf
%name TypeConv tyConv
%name TermWf tmWf
%name TermConv tmConv

-- Respects Environment Quotient -----------------------------------------------

export
envWfRespEq : EnvWf env1 -> env1 =~= env2 -> EnvWf env2
export
typeWfRespEq : TypeWf env1 a -> env1 =~= env2 -> TypeWf env2 a
export
typeConvRespEq : TypeConv env1 a b -> env1 =~= env2 -> TypeConv env2 a b
export
termWfRespEq : TermWf env1 t a -> env1 =~= env2 -> TermWf env2 t a
export
termConvRespEq : TermConv env1 t u a -> env1 =~= env2 -> TermConv env2 t u a

envWfRespEq envWf Refl = envWf
envWfRespEq (envWf :< tyWf) (prf :< prf') =
  envWfRespEq envWf prf :<
  rewrite sym prf' in typeWfRespEq tyWf prf

typeWfRespEq (SetTyWf envWf) prf = SetTyWf (envWfRespEq envWf prf)
typeWfRespEq (PiTyWf tyWf tyWf1) prf =
  PiTyWf
    (typeWfRespEq tyWf prf)
    (typeWfRespEq tyWf1 $ prf :< Refl)
typeWfRespEq (LiftWf tmWf) prf = LiftWf (termWfRespEq tmWf prf)

typeConvRespEq (ReflTy tyWf) prf = ReflTy (typeWfRespEq tyWf prf)
typeConvRespEq (SymTy tyConv) prf = SymTy (typeConvRespEq tyConv prf)
typeConvRespEq (TransTy tyConv tyConv1) prf =
  TransTy
    (typeConvRespEq tyConv prf)
    (typeConvRespEq tyConv1 prf)
typeConvRespEq (PiConv tyWf tyConv tyConv1) prf =
  PiConv
    (typeWfRespEq tyWf prf)
    (typeConvRespEq tyConv prf)
    (typeConvRespEq tyConv1 $ prf :< Refl)
typeConvRespEq (LiftConv tmConv) prf = LiftConv (termConvRespEq tmConv prf)

termWfRespEq (PiTmWf tmWf tmWf1) prf =
  PiTmWf
    (termWfRespEq tmWf prf)
    (termWfRespEq tmWf1 $ prf :< Refl)
termWfRespEq (VarWf {i} envWf) prf =
  rewrite indexCong prf i in
  VarWf (envWfRespEq envWf prf)
termWfRespEq (AbsWf tyWf tmWf) prf =
  AbsWf
    (typeWfRespEq tyWf prf)
    (termWfRespEq tmWf $ prf :< Refl)
termWfRespEq (AppWf tmWf tmWf1) prf =
  AppWf
    (termWfRespEq tmWf prf)
    (termWfRespEq tmWf1 prf)
termWfRespEq (ConvWf tmWf tyConv) prf =
  ConvWf
    (termWfRespEq tmWf prf)
    (typeConvRespEq tyConv prf)

termConvRespEq (ReflTm tmWf) prf = ReflTm (termWfRespEq tmWf prf)
termConvRespEq (SymTm tmConv) prf = SymTm (termConvRespEq tmConv prf)
termConvRespEq (TransTm tmConv tmConv1) prf =
  TransTm
    (termConvRespEq tmConv prf)
    (termConvRespEq tmConv1 prf)
termConvRespEq (AppConv tmConv tmConv1) prf =
  AppConv
    (termConvRespEq tmConv prf)
    (termConvRespEq tmConv1 prf)
termConvRespEq (PiTmConv tyWf tmConv tmConv1) prf =
  PiTmConv
    (typeWfRespEq tyWf prf)
    (termConvRespEq tmConv prf)
    (termConvRespEq tmConv1 $ prf :< Refl)
termConvRespEq (PiEta tyWf tmWf tmWf1 tmConv) prf =
  PiEta
    (typeWfRespEq tyWf prf)
    (termWfRespEq tmWf prf)
    (termWfRespEq tmWf1 prf)
    (termConvRespEq tmConv $ prf :< Refl)
termConvRespEq (PiBeta tyWf tmWf tmWf1) prf =
  PiBeta
    (typeWfRespEq tyWf prf)
    (termWfRespEq tmWf $ prf :< Refl)
    (termWfRespEq tmWf1 prf)
termConvRespEq (ConvConv tmConv tyConv) prf =
  ConvConv
    (termConvRespEq tmConv prf)
    (typeConvRespEq tyConv prf)

-- Well-Formed Environment -----------------------------------------------------

export
typeWfImpliesEnvWf : TypeWf env a -> EnvWf env
export
typeConvImpliesEnvWf : TypeConv env a b -> EnvWf env
export
termWfImpliesEnvWf : TermWf env t a -> EnvWf env
export
termConvImpliesEnvWf : TermConv env t u a -> EnvWf env

typeWfImpliesEnvWf (SetTyWf envWf) = envWf
typeWfImpliesEnvWf (PiTyWf tyWf tyWf1) = typeWfImpliesEnvWf tyWf
typeWfImpliesEnvWf (LiftWf tmWf) = termWfImpliesEnvWf tmWf

typeConvImpliesEnvWf (ReflTy tyWf) = typeWfImpliesEnvWf tyWf
typeConvImpliesEnvWf (SymTy tyConv) = typeConvImpliesEnvWf tyConv
typeConvImpliesEnvWf (TransTy tyConv tyConv1) = typeConvImpliesEnvWf tyConv
typeConvImpliesEnvWf (PiConv tyWf tyConv tyConv1) = typeConvImpliesEnvWf tyConv
typeConvImpliesEnvWf (LiftConv tmConv) = termConvImpliesEnvWf tmConv

termWfImpliesEnvWf (PiTmWf tmWf tmWf1) = termWfImpliesEnvWf tmWf
termWfImpliesEnvWf (VarWf envWf) = envWf
termWfImpliesEnvWf (AbsWf tyWf tmWf) = typeWfImpliesEnvWf tyWf
termWfImpliesEnvWf (AppWf tmWf tmWf1) = termWfImpliesEnvWf tmWf
termWfImpliesEnvWf (ConvWf tmWf tyConv) = termWfImpliesEnvWf tmWf

termConvImpliesEnvWf (ReflTm tmWf) = termWfImpliesEnvWf tmWf
termConvImpliesEnvWf (SymTm tmConv) = termConvImpliesEnvWf tmConv
termConvImpliesEnvWf (TransTm tmConv tmConv1) = termConvImpliesEnvWf tmConv
termConvImpliesEnvWf (AppConv tmConv tmConv1) = termConvImpliesEnvWf tmConv
termConvImpliesEnvWf (PiTmConv tyWf tmConv tmConv1) = termConvImpliesEnvWf tmConv
termConvImpliesEnvWf (PiEta tyWf tmWf tmWf1 tmConv) = termWfImpliesEnvWf tmWf
termConvImpliesEnvWf (PiBeta tyWf tmWf tmWf1) = typeWfImpliesEnvWf tyWf
termConvImpliesEnvWf (ConvConv tmConv tyConv) = termConvImpliesEnvWf tmConv

-- Weakening -------------------------------------------------------------------

public export
record ExtendsWf (0 thin : m `Thins` n) (0 env2 : Env n) (0 env1 : Env m) where
  constructor MkExtWf
  ext : Extends thin env2 env1
  wf : EnvWf env2

%name ExtendsWf extWf

wknPiEta :
  (t : Term m) ->
  (thin : m `Thins` n) ->
  wkn (App (wkn t (wkn1 m)) (Var 0)) (keep thin) =
  App (wkn (wkn t thin) (wkn1 n)) (Var 0)
wknPiEta t thin = cong2 App (sym $ wkn1Comm t thin) (cong Var $ wknKeepFZ thin)

export
weakenTypeWf :
  TypeWf env1 a ->
  ExtendsWf thin env2 env1 ->
  TypeWf env2 (wkn a thin)
export
weakenTypeConv :
  TypeConv env1 a b ->
  ExtendsWf thin env2 env1 ->
  TypeConv env2 (wkn a thin) (wkn b thin)
export
weakenTermWf :
  TermWf env1 t a ->
  ExtendsWf thin env2 env1 ->
  TermWf env2 (wkn t thin) (wkn a thin)
export
weakenTermConv :
  TermConv env1 t u a ->
  ExtendsWf thin env2 env1 ->
  TermConv env2 (wkn t thin) (wkn u thin) (wkn a thin)

weakenTypeWf (SetTyWf envWf) extWf = SetTyWf extWf.wf
weakenTypeWf (PiTyWf {a} tyWf tyWf1) extWf =
  let tyWf = weakenTypeWf tyWf extWf in
  PiTyWf tyWf
    (typeWfRespEq
      (weakenTypeWf tyWf1 $ MkExtWf (Keep' extWf.ext) (extWf.wf :< tyWf))
      (Refl :< sym (wknId $ wkn a thin)))
weakenTypeWf (LiftWf tmWf) extWf = LiftWf (weakenTermWf tmWf extWf)

weakenTypeConv (ReflTy tyWf) extWf = ReflTy (weakenTypeWf tyWf extWf)
weakenTypeConv (SymTy tyConv) extWf = SymTy (weakenTypeConv tyConv extWf)
weakenTypeConv (TransTy tyConv tyConv1) extWf =
  TransTy
    (weakenTypeConv tyConv extWf)
    (weakenTypeConv tyConv1 extWf)
weakenTypeConv (PiConv {a} tyWf tyConv tyConv1) extWf =
  let tyWf = weakenTypeWf tyWf extWf in
  PiConv tyWf
    (weakenTypeConv tyConv extWf)
    (typeConvRespEq
      (weakenTypeConv tyConv1 $ MkExtWf (Keep' extWf.ext) (extWf.wf :< tyWf))
      (Refl :< sym (wknId $ wkn a thin)))
weakenTypeConv (LiftConv tmConv) extWf = LiftConv (weakenTermConv tmConv extWf)

weakenTermWf (PiTmWf {a} tmWf tmWf1) extWf =
  let tmWf = weakenTermWf tmWf extWf in
  PiTmWf tmWf
    (termWfRespEq
      (weakenTermWf tmWf1 $ MkExtWf (Keep' extWf.ext) (extWf.wf :< LiftWf tmWf))
      (Refl :< sym (wknId $ wkn a thin)))
weakenTermWf (VarWf {i} envWf) extWf =
  rewrite sym $ expandHomo (index env1 i) thin in
  rewrite sym $ indexHomo extWf.ext i in
  VarWf extWf.wf
weakenTermWf (AbsWf {a} tyWf tmWf) extWf =
  let tyWf = weakenTypeWf tyWf extWf in
  AbsWf tyWf
    (termWfRespEq
      (weakenTermWf tmWf $ MkExtWf (Keep' extWf.ext) (extWf.wf :< tyWf))
      (Refl :< sym (wknId $ wkn a thin)))
weakenTermWf (AppWf {b, u} tmWf tmWf1) extWf =
  rewrite wknSubst1 b u thin in
  AppWf
    (weakenTermWf tmWf extWf)
    (weakenTermWf tmWf1 extWf)
weakenTermWf (ConvWf tmWf tyConv) extWf =
  ConvWf
    (weakenTermWf tmWf extWf)
    (weakenTypeConv tyConv extWf)

weakenTermConv (ReflTm tmWf) extWf = ReflTm (weakenTermWf tmWf extWf)
weakenTermConv (SymTm tmConv) extWf = SymTm (weakenTermConv tmConv extWf)
weakenTermConv (TransTm tmConv tmConv1) extWf =
  TransTm
    (weakenTermConv tmConv extWf)
    (weakenTermConv tmConv1 extWf)
weakenTermConv (AppConv {b, t} tmConv tmConv1) extWf =
  rewrite wknSubst1 b t thin in
  AppConv
    (weakenTermConv tmConv extWf)
    (weakenTermConv tmConv1 extWf)
weakenTermConv (PiTmConv {a} tyWf tmConv tmConv1) extWf =
  let tyWf = weakenTypeWf tyWf extWf in
  PiTmConv tyWf
    (weakenTermConv tmConv extWf)
    (termConvRespEq
      (weakenTermConv tmConv1 $ MkExtWf (Keep' extWf.ext) (extWf.wf :< tyWf))
      (Refl :< sym (wknId $ wkn a thin)))
weakenTermConv (PiEta {a, t, u} tyWf tmWf tmWf1 tmConv) extWf =
  let tyWf = weakenTypeWf tyWf extWf in
  PiEta tyWf
    (weakenTermWf tmWf extWf)
    (weakenTermWf tmWf1 extWf)
    (termConvRespEq
      (rewrite sym $ wknPiEta t thin in
       rewrite sym $ wknPiEta u thin in
       weakenTermConv tmConv $ MkExtWf (Keep' extWf.ext) (extWf.wf :< tyWf))
      (Refl :< sym (wknId $ wkn a thin)))
weakenTermConv (PiBeta {a, b, t, u} tyWf tmWf tmWf1) extWf =
  let tyWf = weakenTypeWf tyWf extWf in
  rewrite wknSubst1 t u thin in
  rewrite wknSubst1 b u thin in
  PiBeta tyWf
    (termWfRespEq
      (weakenTermWf tmWf $ MkExtWf (Keep' extWf.ext) (extWf.wf :< tyWf))
      (Refl :< sym (wknId $ wkn a thin)))
    (weakenTermWf tmWf1 extWf)
weakenTermConv (ConvConv tmConv tyConv) extWf =
  ConvConv
    (weakenTermConv tmConv extWf)
    (weakenTypeConv tyConv extWf)