From Coq Require Import Morphisms.
From MetaCoq.Template Require Import config utils.
From MetaCoq.PCUIC Require Import PCUICAst PCUICAstUtils PCUICCases PCUICInduction
  PCUICLiftSubst PCUICTyping PCUICCumulativity
  PCUICClosed PCUICReduction
  PCUICSigmaCalculus PCUICRenameDef PCUICRenameConv PCUICOnFreeVars
  PCUICClosedConv PCUICClosedTyp.

Require Import ssreflect ssrbool.
From Equations Require Import Equations.

Implicit Types cf : checker_flags.

Set Default Goal Selector "!".
Generalizable Variables Σ Γ t T.

Ltac Equations.CoreTactics.destruct_tele_eq H ::= noconf H.

Lemma closed_ctx_lift n k ctx : closed_ctx ctx → lift_context n k ctx = ctx.
Proof.
  induction ctx in n, k |- *; auto.
  rewrite closedn_ctx_cons lift_context_snoc0 /snoc.
  move/andb_and ⇒ /= [Hctx Hd].
  rewrite IHctx // lift_decl_closed //. now apply: closed_decl_upwards.
Qed.

Lemma weaken_nth_error_ge {Γ Γ' v Γ''} : #|Γ'| ≤ v →
  nth_error (Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ') (#|Γ''| + v) =
  nth_error (Γ ,,, Γ') v.
Proof.
  intros Hv.
  rewrite → !nth_error_app_context_ge, ?lift_context_length.
  - f_equal. lia.
  - auto.
  - rewrite lift_context_length. lia.
  - rewrite lift_context_length. lia.
Qed.

Lemma weaken_nth_error_lt {Γ Γ' Γ'' v} : v < #|Γ'| →
  nth_error (Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ') v =
  option_map (lift_decl #|Γ''| (#|Γ'| - S v)) (nth_error (Γ ,,, Γ') v).
Proof.
  simpl. intros Hv.
  rewrite → !nth_error_app_context_lt.
  - rewrite nth_error_lift_context_eq.
    do 2 f_equal. lia.
  - lia.
  - now rewrite lift_context_length.
Qed.

Lemma lift_context_lift_context n k Γ : lift_context n 0 (lift_context k 0 Γ) =
  lift_context (n + k) 0 Γ.
Proof. rewrite !lift_context_alt.
  rewrite mapi_compose.
  apply mapi_ext.
  intros n' x.
  rewrite /lift_decl compose_map_decl.
  apply map_decl_ext ⇒ y.
  rewrite mapi_length; autorewrite with len.
  rewrite simpl_lift //; lia.
Qed.

Lemma weakening_renaming P Γ Γ' Γ'' :
  urenaming P (Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ') (Γ ,,, Γ')
    (lift_renaming #|Γ''| #|Γ'|).
Proof.
  intros i d hpi hnth.
  rewrite /lift_renaming.
  destruct (Nat.leb #|Γ'| i) eqn:leb; [apply Nat.leb_le in leb|eapply Nat.leb_nle in leb].
  - rewrite weaken_nth_error_ge //.
    ∃ d; split; auto.
    split; auto.
    split.
    × apply rename_ext ⇒ k. rewrite /rshiftk /lift_renaming.
      repeat nat_compare_specs.
    × destruct (decl_body d) ⇒ /= //.
      f_equal. apply rename_ext ⇒ k.
      rewrite /rshiftk; now nat_compare_specs.
  - rewrite weaken_nth_error_lt; try lia.
    rewrite hnth /=. eexists. split; [eauto|].
    simpl. rewrite !lift_rename !rename_compose /lift_renaming /rshiftk /=.
    repeat split.
    × apply rename_ext ⇒ k. now repeat nat_compare_specs.
    × destruct (decl_body d) ⇒ /= //. f_equal.
      rewrite lift_rename rename_compose /lift_renaming.
      apply rename_ext ⇒ k. simpl. now repeat nat_compare_specs.
Qed.


#[global]
Hint Rewrite rename_context_length : len.


Lemma rename_context_lift_context n k Γ :
  rename_context (lift_renaming n k) Γ = lift_context n k Γ.
Proof.
  rewrite /rename_context /lift_context.
  apply fold_context_k_ext ⇒ i t.
  now rewrite lift_rename shiftn_lift_renaming.
Qed.

Lemma smash_context_lift Δ k n Γ :
  smash_context (lift_context n (k + #|Γ|) Δ) (lift_context n k Γ) =
  lift_context n k (smash_context Δ Γ).
Proof.
  revert Δ. induction Γ as [|[na [b|] ty]]; intros Δ; simpl; auto.
  - now rewrite Nat.add_0_r.
  - rewrite -IHΓ.
    rewrite lift_context_snoc /=. f_equal.
    rewrite !subst_context_alt !lift_context_alt !mapi_compose.
    apply mapi_ext⇒ n' x.
    destruct x as [na' [b'|] ty']; simpl.
    × rewrite !mapi_length /lift_decl /subst_decl /= /map_decl /=; f_equal.
      + f_equal. rewrite Nat.add_0_r distr_lift_subst_rec /=.
        lia_f_equal.
      + rewrite Nat.add_0_r distr_lift_subst_rec; simpl. lia_f_equal.
    × rewrite !mapi_length /lift_decl /subst_decl /= /map_decl /=; f_equal.
      rewrite Nat.add_0_r distr_lift_subst_rec /=.
      repeat (lia || f_equal).
  - rewrite -IHΓ.
    rewrite lift_context_snoc /= // /lift_decl /subst_decl /map_decl /=.
    f_equal.
    rewrite lift_context_app. simpl.
    rewrite /app_context; lia_f_equal.
    rewrite /lift_context // /fold_context_k /= /map_decl /=.
    now lia_f_equal.
Qed.


Definition lift_mutual_inductive_body n k m :=
  map_mutual_inductive_body (fun k' ⇒ lift n (k' + k)) m.

Lemma lift_fix_context:
  ∀ (mfix : list (def term)) (n k : nat),
    fix_context (map (map_def (lift n k) (lift n (#|mfix| + k))) mfix) = lift_context n k (fix_context mfix).
Proof.
  intros mfix n k. unfold fix_context.
  rewrite PCUICLiftSubst.map_vass_map_def rev_mapi.
  fold (fix_context mfix).
  rewrite (lift_context_alt n k (fix_context mfix)).
  unfold lift_decl. now rewrite mapi_length fix_context_length.
Qed.

#[global]
Hint Rewrite <- lift_fix_context : lift.

Lemma lift_it_mkProd_or_LetIn n k ctx t :
  lift n k (it_mkProd_or_LetIn ctx t) =
  it_mkProd_or_LetIn (lift_context n k ctx) (lift n (length ctx + k) t).
Proof.
  induction ctx in n, k, t |- *; simpl; try congruence.
  pose (lift_context_snoc n k ctx a). unfold snoc in e. rewrite → e. clear e.
  simpl. rewrite → IHctx.
  pose (lift_context_snoc n k ctx a).
  now destruct a as [na [b|] ty].
Qed.
#[global]
Hint Rewrite lift_it_mkProd_or_LetIn : lift.

Lemma to_extended_list_map_lift:
  ∀ (n k : nat) (c : context), to_extended_list c = map (lift n (#|c| + k)) (to_extended_list c).
Proof.
  intros n k c.
  pose proof (to_extended_list_lift_above c). unf_term.
  symmetry. solve_all.
  destruct H as [x' [-> Hx]]. simpl.
  destruct (leb_spec_Set (#|c| + k) x').
  - f_equal. lia.
  - reflexivity.
Qed.

Lemma weakening_red1 `{cf:checker_flags} {Σ} Γ Γ' Γ'' M N :
  wf Σ →
  on_free_vars xpredT M →
  red1 Σ (Γ ,,, Γ') M N →
  red1 Σ (Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ') (lift #|Γ''| #|Γ'| M) (lift #|Γ''| #|Γ'| N).
Proof.
  intros.
  rewrite !lift_rename.
  eapply red1_rename; eauto.
  eapply weakening_renaming.
Qed.

Lemma weakening_red `{cf:checker_flags} {Σ:global_env_ext} {wfΣ : wf Σ} {P Γ Γ' Γ'' M N} :
  on_ctx_free_vars P (Γ ,,, Γ') →
  on_free_vars P M →
  red Σ (Γ ,,, Γ') M N →
  red Σ (Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ') (lift #|Γ''| #|Γ'| M) (lift #|Γ''| #|Γ'| N).
Proof.
  intros.
  rewrite !lift_rename.
  eapply red_rename; eauto.
  eapply weakening_renaming.
Qed.

Lemma weakening_red' `{cf:checker_flags} {Σ:global_env_ext} {wfΣ : wf Σ} {P Γ Γ' Γ'' M N} :
  on_ctx_free_vars P (Γ ,,, Γ') →
  on_free_vars P M →
  red Σ (Γ ,,, Γ') M N →
  red Σ (Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ') (lift #|Γ''| #|Γ'| M) (lift #|Γ''| #|Γ'| N).
Proof.
  now eapply weakening_red.
Qed.

Lemma weakening_red_0 {cf} {Σ:global_env_ext} {wfΣ : wf Σ} {P Γ Γ' M N n} :
  n = #|Γ'| →
  on_ctx_free_vars P Γ →
  on_free_vars P M →
  red Σ Γ M N →
  red Σ (Γ ,,, Γ') (lift0 n M) (lift0 n N).
Proof. move⇒ → onctx ont; eapply (weakening_red (Γ':=[])); tea. Qed.


Lemma weakening_cumul `{CF:checker_flags} {Σ Γ Γ' Γ'' M N} :
  wf Σ.1 →
  on_free_vars xpredT M →
  on_free_vars xpredT N →
  on_ctx_free_vars xpredT (Γ ,,, Γ') →
  Σ ;;; Γ ,,, Γ' |- M ≤ N →
  Σ ;;; Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ' |- lift #|Γ''| #|Γ'| M ≤ lift #|Γ''| #|Γ'| N.
Proof.
  intros.
  rewrite !lift_rename -rename_context_lift_context.
  eapply cumul_renameP ; tea.
  rewrite rename_context_lift_context.
  now eapply weakening_renaming.
Qed.


Lemma weakening_conv `{cf:checker_flags} :
  ∀ Σ Γ Γ' Γ'' M N,
    wf Σ.1 →
    on_free_vars xpredT M →
    on_free_vars xpredT N →
    on_ctx_free_vars xpredT (Γ ,,, Γ') →
    Σ ;;; Γ ,,, Γ' |- M = N →
    Σ ;;; Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ' |- lift #|Γ''| #|Γ'| M = lift #|Γ''| #|Γ'| N.
Proof.
  intros.
  rewrite !lift_rename -rename_context_lift_context.
  eapply conv_renameP ; tea.
  rewrite rename_context_lift_context.
  now eapply weakening_renaming.
Qed.

Lemma isType_on_free_vars {cf} {Σ : global_env_ext} {wfΣ : wf Σ} {Γ T} :
  isType Σ Γ T → on_free_vars xpredT T.
Proof.
  intros [s Hs].
  eapply subject_closed in Hs.
  rewrite closedP_on_free_vars in Hs.
  eapply on_free_vars_impl; tea ⇒ //.
Qed.

Lemma isType_on_ctx_free_vars {cf} {Σ : global_env_ext} {wfΣ : wf Σ} {Γ T} :
  isType Σ Γ T → on_ctx_free_vars xpredT Γ.
Proof.
  intros [s Hs].
  eapply typing_wf_local in Hs.
  eapply closed_wf_local in Hs; tea.
  eapply (closed_ctx_on_free_vars xpredT) in Hs.
  now eapply on_free_vars_ctx_on_ctx_free_vars_xpredT.
Qed.

Lemma weakening_conv_wt `{cf:checker_flags} {Σ : global_env_ext} {wfΣ : wf Σ} {Γ Γ' Γ'' M N} :
  isType Σ (Γ ,,, Γ') M → isType Σ (Γ ,,, Γ') N →
  Σ ;;; Γ ,,, Γ' |- M = N →
  Σ ;;; Γ ,,, Γ'' ,,, lift_context #|Γ''| 0 Γ' |- lift #|Γ''| #|Γ'| M = lift #|Γ''| #|Γ'| N.
Proof.
  intros onM onN.
  eapply weakening_conv; tea.
  1-2:now eapply isType_on_free_vars.
  now eapply isType_on_ctx_free_vars in onM.
Qed.