(*Require Import Coq.Logic.FunctionalExtensionality.*)
Require Import Coq.Logic.Eqdep.

Require Import CatSem.AXIOMS.functional_extensionality.

Require Export CatSem.CAT.retype_functor.
Require Export CatSem.CAT.ind_potype.

Set Implicit Arguments.
Unset Strict Implicit.
Unset Transparent Obligations.
Unset Automatic Introduction.

(** define an order on the retyped V, dependent on order on V *)

Section retype_ord.

Variables T U : Type.
Variable f : T -> U.

Inductive retype_ord (V : IPO T) : forall u, relation (retype f V (u)) :=
  | ctype_ord : forall t (x y : V t), x <<< y
            -> retype_ord  (ctype f x) (ctype f y).

Program Instance retype_ipo_s (V : IPO T) :
      ipo_obj_struct (retype f V) := {
  IRel := @retype_ord V
}.
Next Obligation.
Proof.
  constructor.
  unfold Reflexive.
  intro x;
  destruct x as [t x].
  constructor.
  apply IRel_refl.

  unfold Transitive.
  intros x y z H1 H2.
  generalize dependent z.
  induction H1.
  intros z H'.

  inversion H'.
  subst.
  assert (H44:=inj_pair2 _ _ _ _ _ H4).
  assert (H55:=inj_pair2 _ _ _ _ _ H5).
  subst.
  constructor.
  transitivity y; auto.
Qed.

Definition retype_ipo V : IPO U := Build_ipo_obj (retype_ipo_s V).

Section retype_po_hom.

Variables V W : IPO T.
Variable m : V ---> W.

Program Instance ret_po_mor_s : ipo_mor_struct
    (a:=retype_ipo V) (b:=retype_ipo W) (retype_map m).
Next Obligation.
Proof.
  unfold Proper;
  red.
  intros x y H.
  induction H.
  simpl.
  constructor.
  apply m.
  auto.
Qed.

Definition retype_po_mor := Build_ipo_mor ret_po_mor_s.

End retype_po_hom.

Obligation Tactic := idtac.

Program Instance retype_po_func :
    Functor_struct (Fobj := retype_ipo) retype_po_mor.
Next Obligation.
Proof.
  intros a b.
  unfold Proper;
  red.
  intros g g' H t z.
  destruct z as [t z];
  simpl.
  rewrite H;
  auto.
Qed.
Next Obligation.
Proof.
  simpl.
  intros a t z.
  destruct z as [t z];
  auto.
Qed.
Next Obligation.
Proof.
  simpl.
  intros a b c g g' t z;
  destruct z as [t z];
  auto.
Qed.

Definition RETYPE_PO := Build_Functor retype_po_func.

End retype_ord.

(** now natural transformations with SM_ipo *)
Section nattrans.

Variables T U : Type.
Variable f : T -> U.

(*
Definition bla2 c:
(forall t : U,
  (sm_ipo (T:=U) (retype f c)) t -> (retype_ipo f (sm_ipo (T:=T) c)) t) .
intros c t x.
simpl in *.
apply x.
*)

Program Instance id_ccs c :
ipo_mor_struct
 (a:=sm_ipo (T:=U) (retype f c))
 (b:=retype_ipo f (sm_ipo (T:=T) c)) (fun t x => x).
Next Obligation.
Proof.
  unfold Proper.
  red.
  intros x y H.
  induction H.
  apply retype_ipo_s.
Qed.

Definition id_cc c := Build_ipo_mor (id_ccs c).

Program Instance RT_NTs : NT_struct
    (F:=IDelta U O RETYPE f)
    (G:=RETYPE_PO f O IDelta T)
    (fun c => id_cc c).

Definition RT_NT := Build_NT RT_NTs.

Program Instance id_dds c :
ipo_mor_struct
 (b:=sm_ipo (T:=U) (retype f c))
 (a:=retype_ipo f (sm_ipo (T:=T) c)) (fun t x => x).
Next Obligation.
Proof.
  unfold Proper;
  red.
  intros x y H.
  induction H.
  simpl.
  induction H.
  constructor.
Qed.

Definition id_dd c := Build_ipo_mor (id_dds c).

Program Instance NNNT2s : NT_struct
    (G:=IDelta U O RETYPE f)
    (F:=RETYPE_PO f O IDelta T)
    (fun c => id_dd c).

Definition NNNT2 := Build_NT NNNT2s.

End nattrans.

Section Transp_po.

Variables U U': Type.
Variables f g : U -> U'.
Hypothesis H : forall t, g t = f t.

Obligation Tactic := idtac.

Program Instance transp_po_s (V : IPO U) :
  ipo_mor_struct (a:=RETYPE_PO (fun t : U => f t) V)
                 (b:=RETYPE_PO (fun t : U => g t) V)
          (transp H (V:=V)).
Next Obligation.
Proof.
  intros V t.
  assert (Ha : f = g) by
    (apply functional_extensionality; auto).
  unfold Proper;
  red.
  unfold transp.
  subst.
  intros x y H'.
  induction H';
  simpl.
  rewrite (UIP_refl _ _ (H t)).
  simpl.
  constructor;
  auto.
Qed.

Definition transp_po : forall V : IPO U,
     (RETYPE_PO (fun t : U => f t)) V --->
     (RETYPE_PO (fun t : U => g t)) V :=
     fun V => Build_ipo_mor (transp_po_s V).

Program Instance transp_po_NT :
    NT_struct (F:=RETYPE_PO (fun t => f t))
           (G := RETYPE_PO (fun t => g t)) transp_po.
Next Obligation.
Proof.
  simpl.
  intros V W f' t y.
  induction y.
  simpl.
  rewrite <- H.
  simpl.
  auto.
Qed.

Definition Transp_ord : NT (RETYPE_PO (fun t => f t))
         (RETYPE_PO (fun t => g t)) := Build_NT transp_po_NT.

End Transp_po.

Section transp_po_id.

(** retyping with the identity function is the identity *)

Variables U U' : Type.
Variable f : U -> U'.
Variable H : forall t, f t = f t.

Lemma transp_po_id : forall V, transp_po H (V) == id _.
Proof.
  simpl.
  intros V t y.
  rewrite transp_id.
  auto.
Qed.

End transp_po_id.

Section id_retype_po.

Variable T : Type.
Variable V : IPO T.

Program Instance id_retype_po_s : ipo_mor_struct (a:=V)
  (b:=retype_ipo (fun t => t) V) (id_retype(V:=V)).
Next Obligation.
Proof.
  unfold Proper;
  red.
  intros x y H.
  unfold id_retype.
  assert (H':=ctype_ord (fun t => t) H).
  auto.
Qed.

End id_retype_po.


From iris Require Import program_logic.weakestpre.
From iris.proofmode Require Import tactics.
From st.STLCmuVS Require Import lang typing wkpre generic.lift contexts logrel.definitions logrel.compat_lemmas contexts scopedness.
From st.STLCmu Require Import types.
From st.prelude Require Import big_op_three.
From st Require Import resources.

Canonical Structure typeO := leibnizO type.

Section definition.

  Context `{Σ : !gFunctors} `{semΣ_inst : !semΣ Σ}.

  (* We define a stuckness-sensistive and insensitive one *)
  Context (s : stuckness).

  Lemma auto_related_typed_expr Γ (e : expr) τ (pe : typed Γ e τ) :
    open_exprel_typed s Γ e e τ.
  Proof.
    induction pe.
    - by apply compat_Var.
    - by apply compat_Unit.
    - by apply compat_Bool.
    - by apply compat_Int.
    - by apply compat_BinOp.
    - by apply compat_Seq.
    - by apply compat_Pair.
    - by eapply compat_Fst.
    - by eapply compat_Snd.
    - by apply compat_InjL.
    - by apply compat_InjR.
    - by eapply compat_Case.
    - by apply compat_If.
    - by eapply compat_LetIn.
    - by apply compat_Lam.
    - by eapply compat_App.
    - by apply compat_Fold.
    - by apply compat_Unfold.
  Qed.

  Lemma auto_related_ctx_item_typed Γ τ Γ' τ' (Ci : ctx_item) (pCi :  typed_ctx_item Ci Γ τ Γ' τ') :
    ctx_item_rel_typed s Ci Ci Γ τ Γ' τ'.
  Proof.
    destruct pCi; intros e e' pe pe' Hee'; simpl.
    - by apply compat_Lam.
    - eapply compat_App; auto; eauto. by apply auto_related_typed_expr.
    - eapply compat_App; eauto. by apply auto_related_typed_expr.
    - eapply compat_LetIn; eauto. by apply auto_related_typed_expr.
    - eapply compat_LetIn; eauto. by apply auto_related_typed_expr.
    - apply compat_Pair; auto. by apply auto_related_typed_expr.
    - apply compat_Pair; auto. by apply auto_related_typed_expr.
    - by eapply compat_Fst.
    - by eapply compat_Snd.
    - by eapply compat_InjL.
    - by eapply compat_InjR.
    - eapply compat_Case; eauto; by apply auto_related_typed_expr.
    - eapply compat_Case; eauto; by apply auto_related_typed_expr.
    - eapply compat_Case; eauto; by apply auto_related_typed_expr.
    - apply compat_If; eauto; by apply auto_related_typed_expr.
    - apply compat_If; auto; by apply auto_related_typed_expr.
    - apply compat_If; auto; by apply auto_related_typed_expr.
    - eapply compat_BinOp; eauto. by apply auto_related_typed_expr.
    - eapply compat_BinOp; eauto. by apply auto_related_typed_expr.
    - by apply compat_Fold.
    - by apply compat_Unfold.
  Qed.

  Lemma auto_related_ctx_typed Γ τ Γ' τ' (C : ctx) (pC : typed_ctx C Γ τ Γ' τ') :
    ctx_rel_typed s C C Γ τ Γ' τ'.
  Proof.
    induction pC.
    - by intros e e' Hee'.
    - intros e e' pe pe' Hee'. simpl. apply (auto_related_ctx_item_typed _ _ _ _ k H).
      eapply scoped_ctx_fill. eauto. by eapply ctx_typed_scoped.
      eapply scoped_ctx_fill. eauto. by eapply ctx_typed_scoped.
      by apply IHpC.
  Qed.

End definition.


Require Import Declarative.BasicProperties.
Require Import Declarative.KindProperties.
Require Import Declarative.OccurenceReasoning.
Require Import Declarative.FSetReasoning.
Require Import Extraction.

Require Import Program.Tactics.

Lemma binds_subst : forall Γ1 Γ2 x x' A B e,
    ⊢ Γ1, x' : B,, Γ2 ->
    x : A ∈ Γ1, x' : B,, Γ2 -> x <> x' -> x : [e / x'] A ∈ Γ1 ,, [e // x'] Γ2.
Proof.
  induction Γ2; simpl; intros.
  - destruct H0.
    + inversion H0. subst. contradiction.
    + inversion H. subst.
      assert (x' # A) by (eapply notin_dom_bind_fresh; eauto).
      rewrite fresh_subst_eq; auto.
  - destruct a. destruct H0.
    + inversion H0. subst. auto.
    + right. inversion H. apply IHΓ2 with B; auto.
Qed.

Lemma subst_ctx_distr_cons : forall Γ x A y e,
    [e // y] Γ , x : [e / y] A = [e // y] (Γ , x : A).
Proof. reflexivity. Qed.

Lemma ctx_app_cons_assoc : forall (Γ1 Γ2 : context) x A,
    Γ1 ,, Γ2 , x : A = Γ1 ,, (Γ2 , x : A).
Proof. reflexivity. Qed.

Hint Rewrite ctx_app_cons_assoc : substitution.
Hint Rewrite subst_ctx_distr_cons : substitution.
Hint Rewrite subst_open_var_assoc : substitution.
Hint Rewrite subst_extract : substitution.

Ltac crush_lc :=
  repeat
    match goal with
    | |- forall x, x `notin` _ -> lc_expr _ => intros
    | _ => first
            [ progress auto
            | progress eauto using lc_subst, lc_open_preserve
            | progress autorewrite with assoc; auto ]
    | H : lc_expr ?T |- lc_expr (?e ^^ _) =>
      match T with
      | context [e] => inversion H; subst
      end
    | |- lc_expr ([_ / _] _) => apply lc_subst
    end
.

Hint Extern 1 (lc_expr _) => crush_lc : lc.
Hint Extern 1 (forall x, x `notin` _ -> lc_expr _) => crush_lc : lc.

Lemma value_subst : forall e x v,
    lc_expr v -> value e -> value ([v / x] e).
Proof.
  intros * Lc V. induction V; simpl; eauto 4 with lc.
Qed.

Hint Resolve mono_lc : mono.
Hint Resolve subst_mono : mono.
Hint Resolve value_subst : value.

Lemma reduction_substitution : forall A B,
    A ⟶ B -> forall x e, mono_type e -> ([e / x] A) ⟶ ([e / x] B).
Proof with try rewrite subst_open_distr; eauto 4 with lc mono value.
  intros A B Hr.
  induction Hr; simpl; intros...
  - simpl. constructor. crush_lc.
    inversion H0; subst.
    apply lc_e_mu with (add x L); crush_lc.
Qed.

Lemma dom_subst_equal : forall Γ x v,
    dom ([v // x] Γ) = dom Γ.
Proof.
  induction Γ; intros.
  + reflexivity.
  + destruct a. simpl. now rewrite (IHΓ x v).
Qed.

Ltac subst_gather_helper e :=
  match type of e with
  | mono_type ?t => constr:(fv_eexpr (extract t))
  end
.

Ltac gather_for_substitution :=
  let L1 := gather_atoms_with (fun L : atoms => L) in
  let L2 := gather_atoms_with (fun x : atom => singleton x) in
  let L3 := gather_atoms_with_tactic subst_gather_helper in
  constr:(L1 `union` L2 `union` L3).

Ltac substitution_strategy :=
  match goal with
  | _ => solve [auto 3]
  | _ => progress autorewrite with substitution
  | |- _ ⊢ e_app _ _ <: e_app _ _ : [_ / _] _ ^^ _ =>
    (* [v / x] A ^^ e = [v / x] A ^^ [v / x] e *)
    rewrite -> subst_open_distr
  | |- _ ⊢ ([?v / ?x] ?e1) ^^ _ <: _ : _ =>
    rewrite <- subst_open_distr
  | |- _ `notin` fv_eexpr _ => solve [eauto 4 using fv_subst_inclusion]
  | |- _ ⟶ _ => solve [eauto using reduction_substitution]
  | _ => try_constructors
  | _ => solve [eauto]
  end
.

Ltac apply_substitution_strategy :=
  repeat (intros; substitution_strategy).

Ltac solve_subst :=
  solve [
      adjust_cofinites_for gather_for_substitution;
      apply_substitution_strategy].

Theorem substitution : forall Γ1 Γ2 x A B e1 e2 e3,
  Γ1 , x : B ,, Γ2 ⊢ e1 <: e2 : A ->
  Γ1 ⊢ e3 : B -> mono_type e3 ->
  Γ1 ,, [e3 // x] Γ2 ⊢ [e3 / x] e1 <: [e3 / x] e2 : [e3 / x] A.
Proof.
  intros until e3. intros Hsub Hsub3 Mono.
  remember (Γ1 , x : B ,, Γ2) as Γ.
  generalize dependent HeqΓ.
  generalize x Γ2. clear x Γ2.
  apply sub_mut with
      (P := fun Γ e1 e2 A => fun (_ : Γ ⊢ e1 <: e2 : A) =>
         forall x Γ2, Γ = Γ1, x : B,, Γ2 ->
         Γ1 ,, [e3 // x] Γ2 ⊢ [e3 / x] e1 <: [e3 / x] e2 : [e3 / x] A)
      (P0 := fun Γ => fun (_ : ⊢ Γ) =>
         forall x Γ2, Γ = Γ1 , x : B,, Γ2 -> ⊢ Γ1 ,, [e3 // x] Γ2);
      simpl; intros; subst.
  - destruct (x == x0).
    + subst. assert (A0 = B) by (eapply binds_type_equal; eauto). subst.
      rewrite fresh_subst_eq. apply weakening_app; auto.
      apply prefix_wf in w. inversion w. subst.
      eapply fresh_ctx_fresh_expr; eauto.
    + apply s_var. auto. apply binds_subst with B; auto.
  - solve_subst.
  - solve_subst.
  - solve_subst.
  - solve_subst.
  - solve_subst.
  - solve_subst.
  - solve_subst.
  - adjust_cofinites_for gather_for_substitution.
    eapply s_mu.
    + inversion m. subst.
      pick fresh x' and apply mono_mu; autorewrite with assoc; eauto.
    + apply_substitution_strategy.
    + apply_substitution_strategy.
  - solve_subst.
  - solve_subst.
  - apply s_forall_l with (add x L) ([e3 / x] t) k;
      apply_substitution_strategy.
  - solve_subst.
  - solve_subst.
  - solve_subst.

  - destruct Γ2; inversion H.
  - destruct Γ2; inversion H1; subst; simpl in *.
    + auto.
    + apply wf_cons with k; auto.
      rewrite dom_app. rewrite dom_subst_equal.
      eauto using dom_insert_subset.
  - assumption.
Qed.

Corollary substitution_cons : forall Γ x A B e1 e2 e3,
    Γ, x : B ⊢ e1 <: e2 : A ->
    Γ ⊢ e3 : B -> mono_type e3 ->
    Γ ⊢ [e3 / x] e1 <: [e3 / x] e2 : [e3 / x] A.
Proof.
  intros *.
  replace (Γ , x : B) with (Γ ,, one (pair x B) ,, nil) by reflexivity.
  intros.
  replace Γ with (Γ ,, [e3 // x] nil) by reflexivity.
  eapply substitution; simpl in *; eauto.
Qed.

Lemma ctx_type_correct : forall Γ x A,
    ⊢ Γ -> x : A ∈ Γ -> exists k, Γ ⊢ A : e_kind k.
Proof.
  intros * Wf.
  induction Wf; simpl; intros.
  - inversion H.
  - destruct H1.
    + inversion H1. subst.
      eauto.
    + destruct (IHWf H1) as [k0 IH].
      eauto.
Qed.

Theorem type_correctness : forall Γ e1 e2 A,
    Γ ⊢ e1 <: e2 : A -> A = BOX \/ exists k, Γ ⊢ A : e_kind k.
Proof.
  intros * Sub.
  induction Sub; eauto with wf.
  - eauto 3 using ctx_type_correct.
  - pick fresh x.
    destruct (H2 x Fr) as [H' | [k H']].
    + auto.
    + apply kind_sub_inversion_l in H' as (E1 & E2 & E3). subst. eauto with wf.
  - destruct IHSub2 as [H0 | [k H0]]. inversion H0.
    dependent induction H0.
    + clear H1 H3 IHusub IHusub.
      pick fresh x for (L `union` fv_expr B).
      assert (Fr2 : x `notin` L) by auto.
      specialize (H2 x Fr2).
      right. exists k.
      rewrite open_subst_eq with (x := x); auto.
      replace (e_kind k) with ([t / x] e_kind k) by reflexivity.
      eapply substitution_cons; eauto.
    + apply IHusub1 with A k; auto.
      apply kind_sub_inversion_l in H0_0 as (E1 & E2 & E3). now subst.
Qed.

Ltac conclude_type_refl H :=
  match type of H with
  | ?G ⊢ _ <: _ : ?A =>
    let k := fresh "k" in
    let H := fresh "H" in
    let EBox := fresh "Ebox" in
    assert (A = BOX \/ exists k, G ⊢ A : e_kind k) as [Ebox | [k H]]
        by (eapply type_correctness; eassumption);
    [inversion EBox; subst |]
  end
.