From MetaCoq.Template Require Import utils config.
From MetaCoq.PCUIC Require Import PCUICAst PCUICAstUtils PCUICInduction
    PCUICLiftSubst.

Require Import ssreflect.
From Equations Require Import Equations.

Local Set Keyed Unification.

Set Default Goal Selector "!".

#[global]
Hint Rewrite @app_context_length : wf.

Generalizable Variables Σ Γ t T.

Inductive context_subst : context → list term → list term → Type :=
| context_subst_nil : context_subst [] [] []
| context_subst_ass Γ args s na t a :
    context_subst Γ args s →
    context_subst (vass na t :: Γ) (args ++ [a]) (a :: s)
| context_subst_def Γ args s na b t :
    context_subst Γ args s →
    context_subst (vdef na b t :: Γ) args (subst s 0 b :: s).
Derive Signature for context_subst.

Fixpoint make_context_subst ctx args s :=
  match ctx with
  | [] ⇒ match args with
          | [] ⇒ Some s
          | a :: args ⇒ None
          end
  | d :: ctx ⇒
    match d.(decl_body) with
    | Some body ⇒ make_context_subst ctx args (subst0 s body :: s)
    | None ⇒ match args with
              | a :: args ⇒ make_context_subst ctx args (a :: s)
              | [] ⇒ None
              end
    end
  end.

Lemma context_subst_length {Γ a s} : context_subst Γ a s → #|Γ| = #|s|.
Proof. induction 1; simpl; congruence. Qed.

Lemma context_subst_assumptions_length {Γ a s} : context_subst Γ a s → context_assumptions Γ = #|a|.
Proof. induction 1; simpl; try congruence. rewrite app_length /=. lia. Qed.



Lemma make_context_subst_rec_spec ctx args s tele args' s' :
  context_subst ctx args s →
  make_context_subst tele args' s = Some s' →
  context_subst (List.rev tele ++ ctx) (args ++ args') s'.
Proof.
  induction tele in ctx, args, s, args', s' |- ×.
  - move⇒ /= Hc. case: args'.
    + move ⇒ [= <-].
      now rewrite app_nil_r.
    + move⇒ a l //.
  - move⇒ Hc /=. case: a ⇒ [na [body|] ty] /=.
    -- specialize (IHtele (vdef na body ty :: ctx) args (subst0 s body :: s) args' s').
       move⇒ /=. rewrite <- app_assoc.
       move/(IHtele _). move⇒ H /=. apply H.
       constructor. auto.
    -- case: args' ⇒ [|a args']; try congruence.
       specialize (IHtele (vass na ty :: ctx) (args ++ [a]) (a :: s) args' s').
       move⇒ /=. rewrite <- app_assoc.
       move/(IHtele _). move⇒ H /=. simpl in H. rewrite <- app_assoc in H. apply H.
       constructor. auto.
Qed.

Lemma make_context_subst_spec tele args s' :
  make_context_subst tele args [] = Some s' →
  context_subst (List.rev tele) args s'.
Proof.
  move/(make_context_subst_rec_spec [] [] [] _ _ _ context_subst_nil).
  rewrite app_nil_r /= //.
Qed.

Lemma subst_telescope_cons s k d Γ :
  subst_telescope s k (d :: Γ) =
  map_decl (subst s k) d :: subst_telescope s (S k) Γ.
Proof.
  simpl.
  unfold subst_telescope, mapi. simpl. f_equal.
  rewrite mapi_rec_Sk. apply mapi_rec_ext.
  intros. simpl. now rewrite Nat.add_succ_r.
Qed.

Lemma subst_telescope_comm_rec s k s' k' Γ:
  subst_telescope (map (subst s' k) s) k' (subst_telescope s' (#|s| + k' + k) Γ) =
  subst_telescope s' (k' + k) (subst_telescope s k' Γ).
Proof.
  induction Γ in k, k' |- *; rewrite ?subst_telescope_cons; simpl; auto.
  f_equal.
  × unfold map_decl. simpl.
    f_equal.
    + destruct a as [na [b|] ty]; simpl; auto.
      f_equal. now rewrite distr_subst_rec.
    + now rewrite distr_subst_rec.
  × specialize (IHΓ k (S k')). now rewrite Nat.add_succ_r in IHΓ.
Qed.

Lemma subst_telescope_comm s k s' Γ:
  subst_telescope (map (subst s' k) s) 0 (subst_telescope s' (#|s| + k) Γ) =
  subst_telescope s' k (subst_telescope s 0 Γ).
Proof.
  now rewrite -(subst_telescope_comm_rec _ _ _ 0) Nat.add_0_r.
Qed.

Lemma decompose_prod_n_assum_extend_ctx {ctx n t ctx' t'} ctx'' :
  decompose_prod_n_assum ctx n t = Some (ctx', t') →
  decompose_prod_n_assum (ctx ++ ctx'') n t = Some (ctx' ++ ctx'', t').
Proof.
  induction n in ctx, t, ctx', t', ctx'' |- ×.
  - simpl. intros [= → ->]. eauto.
  - simpl.
    destruct t; simpl; try congruence.
    + intros H. eapply (IHn _ _ _ _ ctx'' H).
    + intros H. eapply (IHn _ _ _ _ ctx'' H).
Qed.

Lemma context_subst_length2 {ctx args s} : context_subst ctx args s → #|args| = context_assumptions ctx.
Proof.
  induction 1; simpl; auto.
  rewrite app_length; simpl; lia.
Qed.

Lemma context_subst_fun {ctx args s s'} : context_subst ctx args s → context_subst ctx args s' → s = s'.
Proof.
  induction 1 in s' |- *; intros H'; depelim H'; auto.
  - eapply app_inj_tail in H. intuition subst.
    now specialize (IHX _ H').
  - now specialize (IHX _ H').
Qed.

Lemma context_subst_fun' {ctx args args' s s'} : context_subst ctx args s → context_subst ctx args' s' → #|args| = #|args'|.
Proof.
  induction 1 as [ | ? ? ? ? ? ? ? IHX | ? ? ? ? ? ? ? IHX ] in args', s' |- *; intros H'; depelim H'; auto.
  - now rewrite !app_length; specialize (IHX _ _ H').
  - now specialize (IHX _ _ H').
Qed.

#[global] Hint Constructors context_subst : core.

Lemma context_subst_app {ctx ctx' args s} :
  context_subst (ctx ++ ctx') args s →
  context_subst (subst_context (skipn #|ctx| s) 0 ctx) (skipn (context_assumptions ctx') args) (firstn #|ctx| s) ×
  context_subst ctx' (firstn (context_assumptions ctx') args) (skipn #|ctx| s).
Proof.
  revert ctx' args s.
  induction ctx; intros ctx' args s; simpl.
  - intros Hc.
    rewrite - !(context_subst_length2 Hc).
    now rewrite firstn_all skipn_all.
  - intros Hc.
    depelim Hc.
    × rewrite skipn_S.
      specialize (IHctx _ _ _ Hc) as [IHctx IHctx'].
      pose proof (context_subst_length2 IHctx).
      pose proof (context_subst_length2 IHctx').
      pose proof (context_subst_length2 Hc).
      rewrite context_assumptions_app in H1.
      rewrite firstn_app. rewrite (firstn_0 [a0]).
      { rewrite firstn_length_le in H0; lia. }
      rewrite app_nil_r. split; auto.
      rewrite skipn_app_le; try lia.
      rewrite subst_context_snoc.
      now constructor.

    × specialize (IHctx _ _ _ Hc).
      split; try now rewrite skipn_S.
      pose proof (context_subst_length2 Hc).
      rewrite context_assumptions_app in H.
      destruct IHctx as [IHctx _].
      pose proof (context_subst_length IHctx).
      rewrite subst_context_snoc. rewrite !skipn_S.
      rewrite /subst_decl /map_decl /= Nat.add_0_r.
      rewrite -{4}(firstn_skipn #|ctx| s0).
      rewrite subst_app_simpl. simpl.
      rewrite subst_context_length in H0. rewrite -H0.
      now constructor.
Qed.

Lemma make_context_subst_recP ctx args s tele args' s' :
  context_subst ctx args s →
  (make_context_subst tele args' s = Some s') <~>
  context_subst (List.rev tele ++ ctx) (args ++ args') s'.
Proof.
  induction tele in ctx, args, s, args', s' |- ×.
  - move⇒ /= Hc. case: args'.
    × split.
      + move ⇒ [= <-].
        now rewrite app_nil_r.
      + rewrite app_nil_r.
        move/context_subst_fun ⇒ Hs. now specialize (Hs _ Hc).
    × intros. split; try discriminate.
      intros H2. pose proof (context_subst_fun' Hc H2).
      rewrite app_length /= in H. now lia.
  - move⇒ Hc /=. case: a ⇒ [na [body|] ty] /=.
    × specialize (IHtele (vdef na body ty :: ctx) args (subst0 s body :: s) args' s').
      move⇒ /=. rewrite <- app_assoc.
      now forward IHtele by (constructor; auto).
    × destruct args' as [|a args'].
      + split; [congruence | ]. intros Hc'.
        pose proof (context_subst_length2 Hc').
        rewrite !context_assumptions_app ?app_length ?List.rev_length /= Nat.add_0_r in H.
        pose proof (context_subst_length2 Hc). lia.

      + specialize (IHtele (vass na ty :: ctx) (args ++ [a]) (a :: s) args' s').
        forward IHtele. { econstructor. auto. }
        rewrite -app_assoc. rewrite -app_comm_cons /=.
        rewrite -app_assoc in IHtele. apply IHtele.
Qed.

Lemma make_context_subst_spec_inv : ∀ (tele : list context_decl) (args s' : list term),
  context_subst (List.rev tele) args s' →
  make_context_subst tele args [] = Some s'.
Proof.
  intros. assert (H:=make_context_subst_recP [] [] [] tele args s').
  forward H by constructor.
  rewrite app_nil_r in H. destruct H.
  simpl in ×. auto.
Qed.