Require Import Coq.Logic.FunctionalExtensionality.
Require Import ChargeCore.Open.Stack.
Require Import ChargeCore.Open.Open.
Require Import ExtLib.Core.RelDec.
Require Import ExtLib.Data.List.
Require Import ExtLib.Tactics.Consider.

Set Implicit Arguments.
Unset Strict Implicit.
Set Maximal Implicit Insertion.
Set Universe Polymorphism.

Section Subst.
  Context {A val : Type} {Heq : RelDec (@eq A)}.
  Context {R : ValNull val}.

  Definition subst := A -> (A -> val) -> val.

  Definition stack_subst (s: stack A val) (sub: subst) : stack A val :=
    fun x => sub x s.

  Definition apply_subst {B}
          (e : (A -> val) -> B) (sub : subst) : (A -> val) -> B :=
    fun s => e (stack_subst s sub).

  Definition substlist := list (A * @expr A val).

  (* The identity substitution *)
  Definition subst0 : subst := fun x => stack_get x.

  (* Substitution of one variable *)
  Definition subst1 e x : subst :=
    fun x' => if x' ?[ eq ] x then e else stack_get x'.

  (* Truncating substitution of one variable *)
  Definition subst1_trunc e x : subst :=
    fun x' => if x' ?[ eq ] x then e else empty_open.

  (* Simultaneous substitution of a finite list of variables *)
  Fixpoint substl_aux (subs: substlist) : subst :=
    match subs with
      | nil => subst0
      | cons (x,e) subs' => fun x' => 
        if x' ?[ eq ] x then e else substl_aux subs' x'
    end.

(* Truncating simultaneous substitution of a finite list of variables *)
  Fixpoint substl_trunc_aux (subs: substlist) : subst :=
    match subs with
      | nil => fun _ => empty_open
      | cons (x,e) subs' => fun x' => 
        if x' ?[ eq ] x then e else substl_trunc_aux subs' x'
    end.

  Definition substl es v := substl_aux es v.
  Definition substl_trunc es v := substl_trunc_aux es v.

End Subst.

(* Binds tighter than function application *)
(* We have the e at precedence level one less than the precedence of the infix
   // operator to make parsing work. *)

Notation "g [{ e // x }]" := (apply_subst g (subst1 e x)) (at level 9, e at level 39,
  format "g [{ e // x }]") : open_scope.

Notation "g [{ e //! x }]" := (apply_subst g (subst1_trunc e x)) (at level 9, e at level 39,
  format "g [{ e //! x }]") : open_scope.

Notation " g '//' es " := (apply_subst g (substl es)) (at level 40) : open_scope.

Notation " g '//!' es " := (apply_subst g (substl_trunc es)) (at level 40) : open_scope.

Section Properties.
  Context {A val : Type} {Heq : RelDec (@eq A)} {HeqOk : RelDec_Correct Heq}.
  Context {R : ValNull val}.

  Lemma open_subst_stack {B} (m: @open A val B) (sub : subst) (s : @stack A val) :
    (apply_subst m sub) s = m (stack_subst s sub).
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
  Proof. reflexivity. Qed.

  Lemma subst0_id B (g: @open A val B) : (apply_subst g (@subst0 A val)) = g.
  Proof.
(* Goal: @eq (@subst A val) (subst_fresh vs (@cons A x xs)) (@substl A val Heq (@combine A (@expr A val) (@cons A x xs) (subst_fresh_l vs (@cons A x xs)))) *)
    apply functional_extensionality.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    intros s; simpl; reflexivity.
  Qed.

  Lemma subst1_stack (s: @stack A val) e x :
    stack_subst s (subst1 e x) = stack_add x (e s) s.
  Proof.
(* Goal: @eq (stack A val) (@stack_subst A val s (@subst1 A val Heq e x)) (@stack_add A val Heq x (e s) s) *)
    apply functional_extensionality; intro x'; simpl.
(* Goal: @eq val (@stack_subst A val s (@subst1 A val Heq e x) x') (@stack_add A val Heq x (e s) s x') *)
    unfold stack_subst, subst1, stack_add.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    consider (x' ?[ eq ] x); reflexivity.
  Qed.

  Lemma substl_subst1 x e :
    substl (val := val) (cons (x,e) nil) = subst1 e x.
  Proof. 
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    reflexivity.
  Qed.

  Lemma substl_subst1_trunc x e :
    substl_trunc (cons (x,e) nil) = subst1_trunc e x.
  Proof. 
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    reflexivity.
  Qed.

  Local Open Scope open_scope.
  Lemma subst1_stack_add {B} (p : open B) (e : expr) (x : A) (v : val) (s : @stack A val) : 
    p[{e//x}] (stack_add x v s) = p[{e[{V_expr v//x}]//x}] s.
  Proof.
(* Goal: @eq B (@apply_subst A val B p (@subst1 A val Heq e x) (@stack_add A val Heq x v s)) (@apply_subst A val B p (@subst1 A val Heq (@apply_subst A val val e (@subst1 A val Heq (@V_expr A val v) x)) x) s) *)
    unfold apply_subst, subst1; simpl.
(* Goal: @eq B (p (fun x0 : A => (if @rel_dec A (@eq A) Heq x0 x then @stack_get A val y else @empty_open A val R) s)) (p (@stack_add A val Heq x (s y) (@stack_empty A val R))) *)
    apply f_equal.
(* Goal: @eq (stack A val) (@stack_subst A val (@stack_add A val Heq x v s) (fun x' : A => if @rel_dec A (@eq A) Heq x' x then e else @stack_get A val x')) (@stack_subst A val s (fun x' : A => if @rel_dec A (@eq A) Heq x' x then fun s : forall _ : A, val => e (@stack_subst A val s (fun x'0 : A => if @rel_dec A (@eq A) Heq x'0 x then @V_expr A val v else @stack_get A val x'0)) else @stack_get A val x')) *)
    unfold stack_subst.
(* Goal: @eq (@subst A val) (subst_fresh vs (@cons A x xs)) (@substl A val Heq (@combine A (@expr A val) (@cons A x xs) (subst_fresh_l vs (@cons A x xs)))) *)
    apply functional_extensionality.
(* Goal: forall x0 : A, @eq val ((if @rel_dec A (@eq A) Heq x0 x then e else @stack_get A val x0) (@stack_add A val Heq x v s)) ((if @rel_dec A (@eq A) Heq x0 x then fun s : forall _ : A, val => e (fun x1 : A => (if @rel_dec A (@eq A) Heq x1 x then @V_expr A val v else @stack_get A val x1) s) else @stack_get A val x0) s) *)
    intros y.
(* Goal: @eq val ((if @rel_dec A (@eq A) Heq y x then e else @stack_get A val y) (@stack_add A val Heq x v s)) ((if @rel_dec A (@eq A) Heq y x then fun s : forall _ : A, val => e (fun x0 : A => (if @rel_dec A (@eq A) Heq x0 x then @V_expr A val v else @stack_get A val x0) s) else @stack_get A val y) s) *)
    consider (y ?[ eq ] x); intros; subst.
    + f_equal. apply functional_extensionality; intro z; simpl.
(* Goal: @eq val (@stack_add A val Heq x v s z) ((if @rel_dec A (@eq A) Heq z x then @V_expr A val v else @stack_get A val z) s) *)
      consider (z ?[ eq ] x); intros; subst.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
      * rewrite stack_lookup_add; reflexivity.
      * rewrite stack_lookup_add2; [reflexivity| intuition congruence].
    + unfold stack_get.
(* Goal: @eq val (@stack_add A val Heq x v s y) (s y) *)
      rewrite stack_lookup_add2; [reflexivity | intuition].
  Qed.
      
  Lemma stack_add_var (x : A) (s : @stack A val) (v : val) : (stack_get x) (stack_add x v s) = v.
  Proof.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    unfold stack_get. rewrite stack_lookup_add. reflexivity.
  Qed.
  
  Lemma subst1_val (v : val) (e : open val) (x : A) :
    (V_expr v)[{e//x}] = (V_expr v).
  Proof.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    reflexivity.
  Qed.

  Lemma subst_identity {B} (x : A) (s : @stack A val) (p : open B) :
    p[{(V_expr (s x))//x}] s = p s.
  Proof.
(* Goal: @eq B (@apply_subst A val B p (@subst1 A val Heq (@V_expr A val (s x)) x) s) (p s) *)
    unfold subst1, apply_subst.
(* Goal: @eq B (p (fun x0 : A => (if @rel_dec A (@eq A) Heq x0 x then @stack_get A val y else @empty_open A val R) s)) (p (@stack_add A val Heq x (s y) (@stack_empty A val R))) *)
    apply f_equal.
(* Goal: @eq (stack A val) (@stack_subst A val s (fun x' : A => if @rel_dec A (@eq A) Heq x' x then @V_expr A val (s x) else @stack_get A val x')) s *)
    unfold stack_subst; simpl.
(* Goal: @eq (@subst A val) (subst_fresh vs (@cons A x xs)) (@substl A val Heq (@combine A (@expr A val) (@cons A x xs) (subst_fresh_l vs (@cons A x xs)))) *)
    apply functional_extensionality.
(* Goal: forall x0 : A, @eq val ((if @rel_dec A (@eq A) Heq x0 x then @V_expr A val (s x) else @stack_get A val x0) s) (s x0) *)
    intros y; simpl.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    consider (y ?[ eq ] x); intro; subst; reflexivity.
  Qed.

  Lemma substl_trunc_notin x' xs es :
    ~ In x' xs ->
    substl_trunc (combine xs es) x' = @empty_open _ val _.
  Proof.
(* Goal: forall _ : not (@In A x' xs), @eq (forall _ : forall _ : A, val, val) (@substl_trunc A val Heq R (@combine A (@expr A val) xs es) x') (@empty_open A val R) *)
    intro HnotIn. revert es. induction xs as [|x]; intros.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    + simpl. destruct es; reflexivity.
    + destruct es as [|e]; [reflexivity|]. simpl. consider (x' ?[ eq ] x); intro.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
      - subst. contradiction HnotIn. simpl. auto.
      - apply IHxs. auto with datatypes.
  Qed.

  Lemma subst1_trunc_singleton_stack {B} (p : open B) (s : @stack A val) x y :
    p[{(stack_get y) //! x}] s = p (stack_add x (s y) (stack_empty A val)).
  Proof.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
    unfold apply_subst, subst1_trunc, stack_subst; simpl.
(* Goal: @eq B (p (fun x0 : A => (if @rel_dec A (@eq A) Heq x0 x then @stack_get A val y else @empty_open A val R) s)) (p (@stack_add A val Heq x (s y) (@stack_empty A val R))) *)
    apply f_equal.
(* Goal: @eq (@subst A val) (subst_fresh vs (@cons A x xs)) (@substl A val Heq (@combine A (@expr A val) (@cons A x xs) (subst_fresh_l vs (@cons A x xs)))) *)
    apply functional_extensionality.
(* Goal: forall x0 : A, @eq (forall _ : forall _ : A, val, val) (subst_fresh vs (@cons A x xs) x0) (@substl A val Heq (@combine A (@expr A val) (@cons A x xs) (subst_fresh_l vs (@cons A x xs))) x0) *)
    intro x'.
(* Goal: @eq (forall _ : forall _ : A, val, val) e (@empty_open A val R) *)
    consider (x' ?[ eq ] x); intros; subst.
    * rewrite stack_lookup_add.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
      reflexivity.
    * rewrite stack_lookup_add2; [| auto].
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
      reflexivity.
  Qed.

End Properties.

Hint Rewrite @subst0_id : open.
Hint Rewrite @substl_subst1 : open.
Hint Rewrite @substl_trunc_notin using assumption : open.

Section MoreLemmas.
  Context {A val : Type} {Heq : RelDec (@eq A)} {HeqOk : RelDec_Correct Heq}.
  Context {R : ValNull val}.

  Open Scope open_scope.

  Lemma subst_singleton {B} : 
    forall x e (P : @open A val B), P[{e//x}] = P // (cons (x, e) nil).
  Proof.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    reflexivity.
  Qed.

  Lemma subst_open_expr_nil {B} :
    forall (e : @open A val B), (e // nil) = e.
  Proof.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    reflexivity.
  Qed.

  Lemma subst_var_cons_eq : 
    forall (x:A) es (e : expr (val := val)), (stack_get x // (cons (x, e) es)) = e.
  Proof.
(* Goal: forall (x : A) (es : list (prod A (@expr A val))) (e : @expr A val), @eq (forall _ : forall _ : A, val, val) (@apply_subst A val val (@stack_get A val x) (@substl A val Heq (@cons (prod A (@expr A val)) (@pair A (@expr A val) x e) es))) e *)
    intros x es e.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
    apply functional_extensionality; intros s; simpl.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
    unfold stack_get, apply_subst, stack_subst. simpl.
(* Goal: @eq val ((if @rel_dec A (@eq A) Heq x x then e else @substl A val Heq es x) s) (e s) *)
    consider (x ?[ eq ] x); intuition congruence.
  Qed.

  Lemma subst_var_cons_ineq : forall (x:A) y es (e : expr (val := val))
    (Hneq: x <> y),
    (stack_get x // (cons (y, e) es)) = (stack_get x // es).
  Proof.
(* Goal: forall (x y : A) (es : list (prod A (@expr A val))) (e : @expr A val) (_ : not (@eq A x y)), @eq (forall _ : forall _ : A, val, val) (@apply_subst A val val (@stack_get A val x) (@substl A val Heq (@cons (prod A (@expr A val)) (@pair A (@expr A val) y e) es))) (@apply_subst A val val (@stack_get A val x) (@substl A val Heq es)) *)
    intros x y es e Hneq.
(* Goal: @eq (forall _ : forall _ : A, val, B) (@apply_subst A val B (@apply_subst A val B e (@substl A val Heq es)) (@substl A val Heq fs)) (@apply_subst A val B e (@substl A val Heq (@app (prod A (@expr A val)) (subst_substlist es fs) fs))) *)
    apply functional_extensionality; intros s.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
    unfold stack_get, apply_subst, stack_subst; simpl.
(* Goal: @eq val ((if @rel_dec A (@eq A) Heq x y then e else @substl A val Heq es x) s) (@substl A val Heq es x s) *)
    consider (x ?[ eq ] y); congruence.
  Qed.

  Lemma val_expr_substs : forall (v : val) es,
    ((V_expr v) // es) = (V_expr v).
  Proof.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    reflexivity.
  Qed.

  Definition subst_substlist (es fs : @substlist A val) : substlist :=
    map (fun x => (fst x, (snd x) // fs)) es.

  Lemma subst_combine {B} (e : @open A val B) (es fs : substlist) :
    (e // es // fs) = (e // (app (subst_substlist es fs) fs)).
  Proof.
(* Goal: @eq (forall _ : forall _ : A, val, B) (@apply_subst A val B (@apply_subst A val B e (@substl A val Heq es)) (@substl A val Heq fs)) (@apply_subst A val B e (@substl A val Heq (@app (prod A (@expr A val)) (subst_substlist es fs) fs))) *)
    apply functional_extensionality; intros s.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
    unfold apply_subst, stack_subst; simpl.
(* Goal: @eq (@subst A val) (subst_fresh vs (@cons A x xs)) (@substl A val Heq (@combine A (@expr A val) (@cons A x xs) (subst_fresh_l vs (@cons A x xs)))) *)
    f_equal. apply functional_extensionality. intros x.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
    induction es; simpl.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
    + reflexivity.
    + destruct a; simpl.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
      consider (x ?[ eq ] a); intros; subst; simpl.
(* Goal: @eq val (e0 (fun x : A => @substl A val Heq fs x s)) (@apply_subst A val val e0 (@substl A val Heq fs) s) *)
      * reflexivity.
      * apply IHes.
  Qed.


Close Scope open_scope.

End MoreLemmas.

Section SubstFresh.

  Context {A val : Type} {Heq : RelDec (@eq A)} {HeqOk : RelDec_Correct Heq}.
  Context {R : ValNull val}.
      
   Definition subst_fresh (vs: A -> val) (xs: list A) : subst :=
      fun x' => if (List.anyb (rel_dec x')) xs then V_expr (vs x') else stack_get x'.

   Fixpoint subst_fresh_l (vs: A -> val) (xs: list A) : list (@expr A val) :=
      match xs with
      | nil => nil
      | cons x xs' => cons (V_expr (vs x)) (subst_fresh_l vs xs')
      end.

    (* TODO: switch the two alt. definitions. *)
    Lemma subst_fresh_alt (vs: A -> val) (xs: list A) :
      subst_fresh vs xs = substl (combine xs (subst_fresh_l vs xs)).
    Proof.
(* Goal: @eq (@subst A val) (subst_fresh vs xs) (@substl A val Heq (@combine A (@expr A val) xs (subst_fresh_l vs xs))) *)
      induction xs as [|x]; [reflexivity|].
(* Goal: @eq (@subst A val) (subst_fresh vs (@cons A x xs)) (@substl A val Heq (@combine A (@expr A val) (@cons A x xs) (subst_fresh_l vs (@cons A x xs)))) *)
      apply functional_extensionality.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
      intro x'. simpl. 
(* Goal: @eq (forall _ : forall _ : A, val, val) (subst_fresh vs (@cons A x xs) x') (if @rel_dec A (@eq A) Heq x' x then @V_expr A val (vs x) else @substl A val Heq (@combine A (@expr A val) xs (subst_fresh_l vs xs)) x') *)
      consider (x' ?[ eq ] x); intros.
      + subst x'. unfold subst_fresh. simpl.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if if @rel_dec A (@eq A) Heq x x then true else @anyb A (@rel_dec A (@eq A) Heq x) xs then @V_expr A val (vs x) else @stack_get A val x) (@V_expr A val (vs x)) *)
        consider (x ?[ eq ] x); congruence.
      + rewrite <- IHxs. unfold subst_fresh.
(* Goal: @eq (forall _ : forall _ : A, val, val) (if @anyb A (@rel_dec A (@eq A) Heq x') (@cons A x xs) then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
        simpl. 
(* Goal: @eq (forall _ : forall _ : A, val, val) (if if @rel_dec A (@eq A) Heq x' x then true else @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
        consider (x ?[ eq ] x'); intros; [congruence|].
(* Goal: @eq (forall _ : forall _ : A, val, val) (if if @rel_dec A (@eq A) Heq x' x then true else @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') (if @anyb A (@rel_dec A (@eq A) Heq x') xs then @V_expr A val (vs x') else @stack_get A val x') *)
        consider (x' ?[ eq ]x); intros; [congruence|reflexivity].
    Qed.

End SubstFresh.
