From Coq Require Import Strings.String.
From Coq Require Import Lists.List.
Import ListNotations.
From Coq Require Import Arith.EqNat.

Local Open Scope nat_scope.
Require Import Atom.
Require Import Metatheory.
Require Import FSetNotin.

Inductive type_DeBruijn : Type :=
  | ty_id : string -> type_DeBruijn
  | ty_fvar : atom -> type_DeBruijn
  | ty_bvar : nat -> type_DeBruijn
  | ty_univ : type_DeBruijn -> type_DeBruijn
  | ty_spec : type_DeBruijn -> type_DeBruijn -> type_DeBruijn
  | ty_func  : type_DeBruijn -> type_DeBruijn -> type_DeBruijn
  | ty_morph : type_DeBruijn -> type_DeBruijn -> type_DeBruijn
  | ty_ladder : type_DeBruijn -> type_DeBruijn -> type_DeBruijn
.

Inductive expr_DeBruijn : Type :=
  | ex_fvar : atom -> expr_DeBruijn
  | ex_bvar : nat -> expr_DeBruijn
  | ex_ty_abs : expr_DeBruijn -> expr_DeBruijn
  | ex_ty_app : expr_DeBruijn -> type_DeBruijn -> expr_DeBruijn
  | ex_abs   : type_DeBruijn -> expr_DeBruijn -> expr_DeBruijn
  | ex_morph : type_DeBruijn -> expr_DeBruijn -> expr_DeBruijn
  | ex_app : expr_DeBruijn -> expr_DeBruijn -> expr_DeBruijn
  | ex_let : expr_DeBruijn -> expr_DeBruijn -> expr_DeBruijn
  | ex_ascend : type_DeBruijn -> expr_DeBruijn -> expr_DeBruijn
  | ex_descend : type_DeBruijn -> expr_DeBruijn -> expr_DeBruijn
.

Declare Scope ladder_type_scope.
Declare Scope ladder_expr_scope.
Declare Custom Entry ladder_type.
Declare Custom Entry ladder_expr.

Notation "[< t >]" := t
  (t custom ladder_type at level 99) : ladder_type_scope.
Notation "t" := t
  (in custom ladder_type at level 0, t ident) : ladder_type_scope.
Notation "'∀' t" := (ty_univ t)
  (t custom ladder_type at level 80, in custom ladder_type at level 80).
Notation "'<' σ τ '>'" := (ty_spec σ τ)
  (in custom ladder_type at level 80, left associativity) : ladder_type_scope.
Notation "'[' τ ']'" := (ty_spec (ty_id "Seq") τ)
  (in custom ladder_type at level 70) : ladder_type_scope.
Notation "'(' τ ')'" := τ
  (in custom ladder_type at level 5) : ladder_type_scope.
Notation "σ '->' τ" := (ty_func σ τ)
  (in custom ladder_type at level 75, right associativity) : ladder_type_scope.
Notation "σ '->morph' τ" := (ty_morph σ τ)
  (in custom ladder_type at level 75, right associativity, τ at level 80) : ladder_type_scope.
Notation "σ '~' τ" := (ty_ladder σ τ)
  (in custom ladder_type at level 20, right associativity) : ladder_type_scope.
Notation "'$' x" := (ty_id x%string)
  (in custom ladder_type at level 10, x constr) : ladder_type_scope.
Notation "'%' x" := (ty_bvar x)
  (in custom ladder_type at level 10, x constr) : ladder_type_scope.

Notation "[{ e }]" := e
  (e custom ladder_expr at level 99) : ladder_expr_scope.
Notation "e" := e
  (in custom ladder_expr at level 0, e ident) : ladder_expr_scope.
Notation "'%' x" := (ex_bvar x)
  (in custom ladder_expr at level 10, x constr) : ladder_expr_scope.
Notation "'$' x" := (ex_fvar x)
  (in custom ladder_expr at level 10, x constr) : ladder_expr_scope.
Notation "'Λ' e" := (ex_ty_abs e)
  (in custom ladder_expr at level 10, e custom ladder_expr at level 80, right associativity) : ladder_expr_scope.
Notation "'λ' τ '↦' e" := (ex_abs τ e)
  (in custom ladder_expr at level 70, τ custom ladder_type at level 90, e custom ladder_expr at level 80, right associativity) :ladder_expr_scope.
Notation "'λ' τ '↦morph' e" := (ex_morph τ e)
  (in custom ladder_expr at level 70, τ custom ladder_type at level 90, e custom ladder_expr at level 80, right associativity) :ladder_expr_scope.
Notation "'let' e 'in' t" := (ex_let e t)
  (in custom ladder_expr at level 60, e custom ladder_expr at level 80, t custom ladder_expr at level 80, right associativity) : ladder_expr_scope.
Notation "e 'as' τ" := (ex_ascend τ e)
  (in custom ladder_expr at level 30, e custom ladder_expr, τ custom ladder_type at level 99) : ladder_expr_scope.
Notation "e 'des' τ" := (ex_descend τ e)
  (in custom ladder_expr at level 30, e custom ladder_expr, τ custom ladder_type at level 99) : ladder_expr_scope.
Notation "e1 e2" := (ex_app e1 e2)
  (in custom ladder_expr at level 90, e2 custom ladder_expr at next level) : ladder_expr_scope.
Notation "e '#' τ" := (ex_ty_app e τ)
  (in custom ladder_expr at level 80, τ custom ladder_type at level 101, left associativity) : ladder_expr_scope.
Notation "'(' e ')'" := e
  (in custom ladder_expr, e custom ladder_expr at next level, left associativity) : ladder_expr_scope.

(* number of abstractions in a type *)
Fixpoint type_debruijn_depth (τ:type_DeBruijn) : nat :=
  match τ with
  | ty_id s => 0
  | ty_fvar s => 0
  | ty_bvar x => 0
  | ty_func s t => max (type_debruijn_depth s) (type_debruijn_depth t)
  | ty_morph s t => max (type_debruijn_depth s) (type_debruijn_depth t)
  | ty_univ t => (1 + (type_debruijn_depth t))
  | ty_spec s t => ((type_debruijn_depth s) - 1)
  | ty_ladder s t => max (type_debruijn_depth s) (type_debruijn_depth t)
  end.

(* get the list of all free variables in a type term *)
Fixpoint type_fv (τ : type_DeBruijn) {struct τ} : atoms :=
  match τ with
  | ty_id s => {}
  | ty_fvar α => singleton α
  | ty_bvar x => {}
  | ty_univ τ => (type_fv τ)
  | ty_spec σ τ => (type_fv σ) `union` (type_fv τ)
  | ty_func σ τ => (type_fv σ) `union` (type_fv τ)
  | ty_morph σ τ => (type_fv σ) `union` (type_fv τ)
  | ty_ladder σ τ => (type_fv σ) `union` (type_fv τ)
  end.

(* substitute free variable x with type σ in τ *)
Fixpoint subst_type (x:atom) (σ:type_DeBruijn) (τ:type_DeBruijn) {struct τ} : type_DeBruijn :=
  match τ with
  | ty_id s => ty_id s
  | ty_fvar s => if x == s then σ else τ
  | ty_bvar y => ty_bvar y
  | ty_univ τ => ty_univ (subst_type x σ τ)
  | ty_spec τ1 τ2 => ty_spec (subst_type x σ τ1) (subst_type x σ τ2)
  | ty_func τ1 τ2 => ty_func (subst_type x σ τ1) (subst_type x σ τ2)
  | ty_morph τ1 τ2 => ty_morph (subst_type x σ τ1) (subst_type x σ τ2)
  | ty_ladder τ1 τ2 => ty_ladder (subst_type x σ τ1) (subst_type x σ τ2)
  end.

(* replace a free variable with a new (dangling) bound variable *)
Fixpoint type_bind_fvar (x:atom) (n:nat) (τ:type_DeBruijn) {struct τ} : type_DeBruijn :=
  match τ with
  | ty_id s => ty_id s
  | ty_fvar s => if x == s then ty_bvar n else τ
  | ty_bvar n => ty_bvar n
  | ty_univ τ1 => ty_univ (type_bind_fvar x (S n) τ1)
  | ty_spec τ1 τ2 => ty_spec (type_bind_fvar x n τ1) (type_bind_fvar x n τ2)
  | ty_func τ1 τ2 => ty_func (type_bind_fvar x n τ1) (type_bind_fvar x n τ2)
  | ty_morph τ1 τ2 => ty_morph (type_bind_fvar x n τ1) (type_bind_fvar x n τ2)
  | ty_ladder τ1 τ2 => ty_ladder (type_bind_fvar x n τ1) (type_bind_fvar x n τ2)
  end.

(* replace (dangling) index with another type *)
Fixpoint type_open_rec (k:nat) (σ:type_DeBruijn) (τ:type_DeBruijn) {struct τ} : type_DeBruijn :=
  match τ with
    | ty_id s => ty_id s
    | ty_fvar s => ty_fvar s
    | ty_bvar i => if k === i then σ else τ
    | ty_univ τ1 => ty_univ (type_open_rec (S k) σ τ1)
    | ty_spec τ1 τ2 => ty_spec (type_open_rec k σ τ1) (type_open_rec k σ τ2)
    | ty_func τ1 τ2 => ty_func (type_open_rec k σ τ1) (type_open_rec k σ τ2)
    | ty_morph τ1 τ2 => ty_morph (type_open_rec k σ τ1) (type_open_rec k σ τ2)
    | ty_ladder τ1 τ2 => ty_ladder (type_open_rec k σ τ1) (type_open_rec k σ τ2)
  end.

Notation "'[' z '~>' u ']' e" := (subst_type z u e) (at level 68).
Notation "'{' k '~>' σ '}' τ" := (type_open_rec k σ τ) (at level 67).
Definition type_open σ τ := type_open_rec 0 σ τ.

(* is the type locally closed ? *)
Inductive type_lc : type_DeBruijn -> Prop :=
  | Tlc_Id : forall s, type_lc (ty_id s)
  | Tlc_Var : forall s, type_lc (ty_fvar s)
  | Tlc_Univ : forall τ1 L,
      (forall x, (x `notin` L) -> type_lc (type_open (ty_fvar x) τ1)) ->
      type_lc (ty_univ τ1)
  | Tlc_Spec : forall τ1 τ2, type_lc τ1 -> type_lc τ2 -> type_lc (ty_spec τ1 τ2)
  | Tlc_Func : forall τ1 τ2, type_lc τ1 -> type_lc τ2 -> type_lc (ty_func τ1 τ2)
  | Tlc_Morph : forall τ1 τ2, type_lc τ1 -> type_lc τ2 -> type_lc (ty_morph τ1 τ2)
  | Tlc_Ladder : forall τ1 τ2, type_lc τ1 -> type_lc τ2 -> type_lc (ty_ladder τ1 τ2)
.

(** Substitution in Expressions *)


(* get the list of all free variables in an expression *)
Fixpoint expr_fv (e : expr_DeBruijn) {struct e} : atoms :=
  match e with
  | ex_fvar n => singleton n
  | ex_bvar y => {}
  | ex_ty_abs t' => (expr_fv t')
  | ex_ty_app t' σ => (expr_fv t')
  | ex_morph σ t' => (expr_fv t')
  | ex_abs σ t' => (expr_fv t')
  | ex_app t1 t2 => (expr_fv t1) `union` (expr_fv t2)
  | ex_let t1 t2 => (expr_fv t1) `union` (expr_fv t2)
  | ex_ascend τ t' => (expr_fv t')
  | ex_descend τ t' => (expr_fv t')
  end.

(* substitute free variable x with expression s in t *)
Fixpoint ex_subst (x:atom) (s:expr_DeBruijn) (t:expr_DeBruijn) {struct t} : expr_DeBruijn :=
  match t with
  | ex_fvar n => if x == n then s else t
  | ex_bvar y => ex_bvar y
  | ex_ty_abs t' => ex_ty_abs (ex_subst x s t')
  | ex_ty_app t' σ => ex_ty_app (ex_subst x s t') σ
  | ex_morph σ t' => ex_morph σ (ex_subst x s t')
  | ex_abs σ t' => ex_abs σ (ex_subst x s t')
  | ex_app t1 t2 => ex_app (ex_subst x s t1) (ex_subst x s t2)
  | ex_let t1 t2 => ex_let (ex_subst x s t1) (ex_subst x s t2)
  | ex_ascend τ t' => ex_ascend τ (ex_subst x s t')
  | ex_descend τ t' => ex_descend τ (ex_subst x s t')
  end.

(* substitute free type-variable α with type τ in e *)
Fixpoint ex_subst_type (α:atom) (τ:type_DeBruijn) (e:expr_DeBruijn) {struct e} : expr_DeBruijn :=
  match e with
  | ex_fvar n => ex_fvar n
  | ex_bvar y => ex_bvar y
  | ex_ty_abs e' => ex_ty_abs (ex_subst_type α τ e')
  | ex_ty_app e' σ => ex_ty_app (ex_subst_type α τ e') (subst_type α τ σ)
  | ex_morph σ e' => ex_morph (subst_type α τ σ) (ex_subst_type α τ e')
  | ex_abs σ e' => ex_abs (subst_type α τ σ) (ex_subst_type α τ e')
  | ex_app e1 e2 => ex_app (ex_subst_type α τ e1) (ex_subst_type α τ e2)
  | ex_let e1 e2 => ex_let (ex_subst_type α τ e1) (ex_subst_type α τ e2)
  | ex_ascend σ e' => ex_ascend (subst_type α τ σ) (ex_subst_type α τ e')
  | ex_descend σ e' => ex_descend (subst_type α τ σ) (ex_subst_type α τ e')
  end.

(* replace a free variable with a new (dangling) bound variable *)
Fixpoint expr_bind_fvar (x:atom) (n:nat) (e:expr_DeBruijn) {struct e} : expr_DeBruijn :=
  match e with
  | ex_fvar s => if x == s then ex_bvar n else e
  | ex_bvar n => ex_bvar n
  | ex_ty_abs e' => ex_ty_abs (expr_bind_fvar x n e')
  | ex_ty_app e' σ => ex_ty_app (expr_bind_fvar x n e') σ
  | ex_morph σ e' => ex_morph σ (expr_bind_fvar x (S n) e')
  | ex_abs σ e' => ex_abs σ (expr_bind_fvar x (S n) e')
  | ex_app e1 e2 => ex_app (expr_bind_fvar x n e1) (expr_bind_fvar x n e2)
  | ex_let e1 e2 => ex_let (expr_bind_fvar x n e1) (expr_bind_fvar x n e2)
  | ex_ascend σ e' => ex_ascend σ (expr_bind_fvar x n e')
  | ex_descend σ e' => ex_descend σ (expr_bind_fvar x n e')
  end.

(* replace (dangling) index with another expression *)
Fixpoint expr_open_rec (k:nat) (t:expr_DeBruijn) (e:expr_DeBruijn) {struct e} : expr_DeBruijn :=
  match e with
    | ex_fvar s => ex_fvar s
    | ex_bvar i => if k === i then t else e
    | ex_ty_abs e' => ex_ty_abs (expr_open_rec k t e')
    | ex_ty_app e' σ => ex_ty_app (expr_open_rec k t e') σ
    | ex_morph σ e' => ex_morph σ (expr_open_rec (S k) t e')
    | ex_abs σ e' => ex_abs σ (expr_open_rec (S k) t e')
    | ex_app e1 e2 => ex_app (expr_open_rec k t e1) (expr_open_rec k t e2)
    | ex_let e1 e2 => ex_let (expr_open_rec k t e1) (expr_open_rec k t e2)
    | ex_ascend σ e' => ex_ascend σ (expr_open_rec k t e')
    | ex_descend σ e' => ex_descend σ (expr_open_rec k t e')
  end.

Definition expr_open t e := expr_open_rec 0 t e.



(* replace (dangling) index with another expression *)
Fixpoint expr_open_type_rec (k:nat) (τ:type_DeBruijn) (e:expr_DeBruijn) {struct e} : expr_DeBruijn :=
  match e with
    | ex_fvar s => ex_fvar s
    | ex_bvar s => ex_bvar s
    | ex_ty_abs e' => ex_ty_abs (expr_open_type_rec (S k) τ e')
    | ex_ty_app e' σ => ex_ty_app (expr_open_type_rec k τ e') (type_open_rec k τ σ)
    | ex_morph σ e' => ex_morph (type_open_rec k τ σ) (expr_open_type_rec k τ e')
    | ex_abs σ e' => ex_abs (type_open_rec k τ σ) (expr_open_type_rec k τ e')
    | ex_app e1 e2 => ex_app (expr_open_type_rec k τ e1) (expr_open_type_rec k τ e2)
    | ex_let e1 e2 => ex_let (expr_open_type_rec k τ e1) (expr_open_type_rec k τ e2)
    | ex_ascend σ e' => ex_ascend σ (expr_open_type_rec k τ e')
    | ex_descend σ e' => ex_descend σ (expr_open_type_rec k τ e')
  end.

Definition expr_open_type τ e := expr_open_type_rec 0 τ e.


(* is the expression locally closed ? *)
Inductive expr_lc : expr_DeBruijn -> Prop :=
  | Elc_Var : forall s, expr_lc (ex_fvar s)
  | Elc_TypAbs : forall e, expr_lc e -> expr_lc (ex_ty_abs e)
  | Elc_TypApp : forall e σ, expr_lc e -> type_lc σ -> expr_lc (ex_ty_app e σ)
  | Elc_Abs : forall σ e1 L,
      (type_lc σ) ->
      (forall x, (x `notin` L) -> expr_lc (expr_open (ex_fvar x) e1)) ->
      expr_lc (ex_abs σ e1)
  | Tlc_App : forall e1 e2, expr_lc e1 -> expr_lc e2 -> expr_lc (ex_app e1 e2)
  | Tlc_Let : forall e1 e2, expr_lc e1 -> expr_lc e2 -> expr_lc (ex_let e1 e2)
  | Tlc_Ascend : forall τ e, type_lc τ -> expr_lc e -> expr_lc (ex_ascend τ e)
  | Tlc_Descend : forall τ e, type_lc τ -> expr_lc e -> expr_lc (ex_descend τ e)
.