Require Import PeanoNat.
Import Nat.

Require Import Lists.List.
Import ListNotations.

Inductive type : Set :=
| tyvar : nat -> type
| arrow : type -> type -> type.

Scheme Equality for type.

Inductive term : Set :=
| var : nat -> term
| free : nat -> term
| app : term -> term -> term
| abs : type -> term -> term.

Declare Custom Entry lambda.
Notation "<{ e }>" := e (e custom lambda at level 99).
Notation "( x )" := x (in custom lambda, x at level 99).
Notation "{ x }" := x (x at level 99).
Notation "x" := x (in custom lambda at level 0, x constr at level 0).
Notation "S -> T" := (arrow S T) (in custom lambda at level 50, right associativity).
Notation "x y" := (app x y) (in custom lambda at level 1, left associativity).
Notation "'λ' t , x" :=
  (abs t x) (in custom lambda at level 90,
  t custom lambda at level 99,
  x custom lambda at level 99,
  left associativity).
Coercion var : nat >-> term.

Reserved Notation "'[' x ':=' s ']' t" (in custom lambda at level 20).
Fixpoint substi (x : nat) (s t : term) : term :=
  match t with
  | var y => if x =? y then s else var y
  | free y => free y
  | <{ m n }> => <{ ([x := s] m) ([x := s] n) }>
  | <{ λ α, m }> => <{ λ α, [x := s] m }>
  end

where "'[' x ':=' s ']' t" := (substi x s t) (in custom lambda).

Fixpoint β_reduce (m : term) : term :=
  match m with
  | var x => var x
  | free x => free x
  | <{ (λ α , m) n }> => <{ [0 := n] m }>
  | <{ m n }> => <{ {β_reduce m} {β_reduce n} }>
  | <{ λ α, m }> => <{ λ α , {β_reduce m} }>
  end.

Fixpoint β_nf (m : term) : bool :=
  match m with
  | var _ => true
  | free _ => true
  | <{ (λ α, m) n }> => false
  | <{ m n }> => β_nf m && β_nf n
  | <{ λ α, m }> => β_nf m
  end.

Reserved Notation "E ';' Γ '|-' t '::' T"
  (at level 101, t custom lambda, T custom lambda at level 0).

Definition context := list type.
Definition environment := list (nat * type).

Fixpoint lookup (E : environment) (x : nat) : option type :=
  match E with
  | [] => None
  | (x', σ) :: rest => if x =? x' then Some σ else lookup rest x
  end.

Lemma lookup_in : forall E x σ,
    lookup E x = Some σ -> In (x, σ) E.
Proof.
  induction E.
  - simpl; discriminate.
  - intros.
    destruct a as [y τ].
    simpl in H.
    destruct (x =? y) eqn:Heq.
    { apply eqb_eq in Heq; subst; inversion H; simpl; left; reflexivity. }
    simpl.
    right.
    auto.
Qed.

Inductive derivation : environment -> context -> term -> type -> Prop :=
| free_rule : forall E Γ x σ,
    lookup E x = Some σ -> E ; Γ |- {free x} :: σ
| var_rule : forall E Γ x σ,
    nth_error Γ x = Some σ -> E ; Γ |- x :: σ
| app_rule : forall E Γ t1 t2 σ τ,
    (E ; Γ |- t1 :: (σ -> τ)) ->
    (E ; Γ |- t2 :: σ) ->
    (E ; Γ |- t1 t2 :: τ)
| abs_rule : forall E Γ σ τ m,
    (E ; σ :: Γ |- m :: τ) ->
    (E ; Γ |- λ σ, m :: (σ -> τ))

where "E ; Γ '|-' t '::' T" := (derivation E Γ t T).

Hint Constructors derivation : core.

Theorem uniqueness_of_types : forall E Γ t σ τ,
    E ; Γ |- t :: σ -> E ; Γ |- t :: τ -> σ = τ.
Proof.
  intros E Γ t.
  generalize dependent Γ.
  induction t; intros; inversion H; inversion H0; subst;
    try (rewrite H4 in H9; inversion H9; reflexivity).
  - apply (IHt1 Γ <{σ0 -> σ}> <{σ1 -> τ}> H5) in H12.
    inversion H12.
    reflexivity.
  - assert (G : τ0 = τ1) ; eauto.
    subst.
    reflexivity.
Qed.

Fixpoint find_type (E : environment) (Γ : context) (m : term) : option type :=
  match m with
  | var x => nth_error Γ x
  | free x => lookup E x
  | app t1 t2 => match find_type E Γ t1, find_type E Γ t2 with
                | Some (<{ σ -> τ }>), Some σ' =>
                    if type_eq_dec σ σ'
                    then Some τ
                    else None
                | _, _ => None
                end
  | abs σ t => option_map (arrow σ) (find_type E (σ :: Γ) t)
  end.

Theorem find_type_correct : forall E Γ t σ,
    find_type E Γ t = Some σ -> E ; Γ |- t :: σ.
Proof.
  intros E Γ t.
  generalize dependent Γ.
  induction t ; auto.
  - intros.
    simpl in H.
    destruct (find_type E Γ t1) eqn:Heq ; inversion H.
    destruct t0 ; inversion H.
    destruct (find_type E Γ t2) eqn:Heq2; inversion H.
    destruct (type_eq_dec t0_1 t0); inversion H.
    subst.
    eauto.
  - intros.
    simpl in H.
    destruct (find_type E (t0 :: Γ) t1) eqn:Heq; inversion H.
    auto.
Qed.

(* Fixpoint reverse_lookup_env (E : environment) (σ : type) : option nat := *)
(*   match E with *)
(*   | [] => None *)
(*   | (x, σ') :: rest => *)
(*       if type_eq_dec σ σ' then Some x else reverse_lookup_env rest σ *)
(*   end. *)

(* Definition env_no_dup (E : environment) := forall x σ σ', *)
(*     In (x, σ) E -> In (x, σ') E -> σ = σ'. *)

(* Lemma reverse_lookup_env_correct : forall E x σ, *)
(*     env_no_dup E -> reverse_lookup_env E σ = Some x -> lookup E x = Some σ. *)
(* Proof. *)
(*   induction E. *)
(*   - simpl. *)
(*     discriminate. *)
(*   - intros. *)
(*     destruct a as [y τ]. *)
(*     simpl in *. *)
(*     destruct (eq_dec x y). *)
(*     + rewrite e. *)
(*       rewrite eqb_refl. *)
(*       unfold env_no_dup in H. *)
(*       specialize H with (x := y) (σ := σ) (σ' := τ) as H'. *)
(*       apply (f_equal Some). *)
(*       symmetry. *)
(*       apply H'; try (simpl; left; subst; reflexivity). *)
(*       destruct (type_eq_dec σ τ). *)
(*       { simpl. left. rewrite e0. reflexivity. } *)
(*       simpl. *)
(*       right. *)
(*       apply lookup_in. *)
(*       apply IHE. *)
(*       * unfold env_no_dup. *)
(*         intros. *)
(*         apply H with (x := x0); simpl; right; assumption. *)
(*       * subst. *)
(*         assumption. *)
(*     + apply eqb_neq in n as n'. *)
(*       rewrite n'. *)
(*       apply IHE. *)
(*       * unfold env_no_dup. *)
(*         intros. *)
(*         apply H with (x := x0); simpl; right; assumption. *)
(*       * destruct (type_eq_dec σ τ). *)
(*         { inversion H0. symmetry in H2. contradiction. } *)
(*         assumption. *)
(* Qed. *)

(* Fixpoint reverse_lookup_context (Γ : context) (σ : type) : option nat := *)
(*   match Γ with *)
(*   | [] => None *)
(*   | σ' :: rest => if type_eq_dec σ σ' *)
(*                 then Some 0 *)
(*                 else option_map S (reverse_lookup_context rest σ) *)
(*   end. *)

(* Lemma reverse_lookup_context_correct : forall Γ n σ, *)
(*     reverse_lookup_context Γ σ = Some n -> nth_error Γ n = Some σ. *)
(* Proof. *)
(*   induction Γ. *)
(*   - intros. *)
(*     inversion H. *)
(*   - intros. *)
(*     induction n. *)
(*     + simpl in *. *)
(*       destruct (type_eq_dec σ a). *)
(*       { rewrite e; reflexivity. } *)
(*       destruct (reverse_lookup_context Γ σ); inversion H. *)
(*     + simpl. *)
(*       apply IHΓ. *)
(*       simpl in H. *)
(*       destruct (type_eq_dec σ a) ; inversion H. *)
(*       clear H1. *)
(*       destruct (reverse_lookup_context Γ σ) ; inversion H. *)
(*       reflexivity. *)
(* Qed. *)

(* Definition or_else {A} (x1 x2 : option A) : option A := *)
(*   match x1, x2 with *)
(*   | Some x, _ => Some x *)
(*   | _, Some x => Some x *)
(*   | _, _ => None *)
(*   end. *)

(* Hint Unfold or_else : core. *)

(* Fixpoint find_term (E : environment) (Γ : context) (σ : type) : option term := *)
(*   match σ with *)
(*   | tyvar n => (or_else (option_map var (reverse_lookup_context Γ σ)) *)
(*                  (option_map free (reverse_lookup_env E σ))) *)
(*   | arrow σ1 σ2 => option_map (abs σ1) (find_term E (σ1 :: Γ) σ2) *)
(*   end. *)

(* Theorem find_term_correct : forall E Γ t σ, *)
(*     env_no_dup E -> *)
(*     find_term E Γ σ = Some t -> *)
(*     E ; Γ |- t :: σ. *)
(* Proof. *)
(*   intros. *)
(*   generalize dependent Γ. *)
(*   generalize dependent t0. *)
(*   induction σ. *)
(*   - intros t Γ H1. *)
(*     simpl in *. *)
(*     destruct (reverse_lookup_context Γ (tyvar n)) eqn:Hcon; *)
(*     destruct (reverse_lookup_env E (tyvar n)) eqn:Henv; *)
(*       simpl in H1; inversion H1; subst; *)
(*       auto using reverse_lookup_context_correct; *)
(*       auto using reverse_lookup_env_correct. *)
(*   - intros t Γ H1. *)
(*     simpl in H1. *)
(*     destruct (find_term E (σ1 :: Γ) σ2) eqn:Heq; inversion H1; auto. *)
(* Qed. *)