From Undecidability.Shared.Libs.PSL Require Import Vectors.
Require Import Undecidability.SOL.SOL.
From Undecidability.SOL.Util Require Import Subst Syntax.
Require Import Arith Lia Vector.

Import SubstNotations.
Unset Implicit Arguments.

Arguments eval_function {_ _ _ _ _}.
Arguments eval_predicate {_ _ _ _ _}.
Arguments eval {_ _ _ _}.
Arguments sat {_ _ _ _}.
Arguments new_env {_} _ _ _.
Arguments get_indi {_} _ _.
Arguments get_func {_} _ _.
Arguments get_pred {_} _ _.

Notation "⟨ a , b , c ⟩" := (new_env a b c).

Class Ent X Y `{funcs_signature, preds_signature} := ent : X -> Y -> Prop.
Notation "X ⊨ phi" := (ent X phi) (at level 20).
Class Ent' X `{funcs_signature, preds_signature} := ent' : forall M : Model, env (M_domain M) -> X -> Prop.
Notation "( M , rho ) ⊨ phi" := (ent' M rho phi) (at level 0).

#[export] Instance ent_env `{funcs_signature, preds_signature} domain I : Ent (env domain) form :=
  @sat _ _ domain I.
#[export] Instance ent'_form `{funcs_signature, preds_signature} : Ent' form :=
  fun M rho phi => @sat _ _ (M_domain M) (M_interp M) rho phi.
#[export] Instance ent_model `{funcs_signature, preds_signature} : Ent Model form :=
  fun M phi => forall rho, @sat _ _ (M_domain M) (M_interp M) rho phi.
#[export] Instance ent_model_theory `{funcs_signature, preds_signature} : Ent Model (form -> Prop) :=
  fun M T => forall phi, T phi -> M phi.
#[export] Instance ent_theory `{funcs_signature, preds_signature} : Ent (form -> Prop) form :=
  fun T phi => forall (M : Model) rho, (forall psi, T psi -> (M, rho) psi) -> (M, rho) phi.
#[export] Instance ent'_theory `{funcs_signature, preds_signature} : Ent' (form -> Prop) :=
  fun M rho T => forall phi, T phi -> (M, rho) phi.
#[export] Instance ent'_form' `{funcs_signature, preds_signature} : Ent' form :=
  fun M rho phi => @sat _ _ (M_domain M) (M_interp M) rho phi.

Section Environment.

  Variable domain : Type.

  Definition env_equiv (rho1 : env domain) rho2 := forall n,
        get_indi rho1 n = get_indi rho2 n
    /\ forall ar v, get_func rho1 n ar v = get_func rho2 n ar v
                /\ (get_pred rho1 n ar v <-> get_pred rho2 n ar v).
  Notation "rho1 ≡ rho2" := (env_equiv rho1 rho2) (at level 30).

  Lemma env_equiv_symm rho1 rho2 :
    rho1 rho2 -> rho2 rho1.
  Proof.
    intros H. intros n. specialize (H n). split. easy.
    intros ar v. destruct H as [_ H]. specialize (H ar v). easy.
  Qed.

  Lemma env_equiv_cons_i rho1 rho2 x :
  rho1 rho2 -> x .: get_indi rho1, get_func rho1, get_pred rho1 x .: get_indi rho2, get_func rho2, get_pred rho2 .
  Proof.
    intros H n. split. 2:split. destruct n. all: firstorder.
  Qed.

  Lemma env_equiv_cons_f rho1 rho2 ar (f : vec domain ar -> domain) f' :
    rho1 rho2 -> (forall v, f v = f' v)
    -> get_indi rho1, f .: get_func rho1, get_pred rho1 get_indi rho2, f' .: get_func rho2, get_pred rho2 .
  Proof.
    intros H Hf n. split. 2:split.
    - apply H.
    - destruct n; cbn; destruct (Nat.eq_dec ar ar0) as [->|]; firstorder.
    - apply H.
  Qed.

  Lemma env_equiv_cons_p rho1 rho2 ar (P : vec domain ar -> Prop) :
    rho1 rho2 -> get_indi rho1, get_func rho1, P .: get_pred rho1 get_indi rho2, get_func rho2, P .: get_pred rho2 .
  Proof.
    intros H n. split. 2:split.
    - apply H.
    - apply H.
    - destruct n; cbn; destruct (Nat.eq_dec ar ar0) as [->|]. all: firstorder.
  Qed.

End Environment.

Notation "rho1 ≡ rho2" := (env_equiv _ rho1 rho2) (at level 30).

Section SatExt.

  Context {Σ_funcs : funcs_signature}.
  Context {Σ_preds : preds_signature}.
  Variable domain : Type.
  Context {I : interp domain}.

  Lemma eval_function_ext rho1 rho2 ar (f : function ar) :
    (forall n ar v, get_func rho1 n ar v = get_func rho2 n ar v)
      -> forall v, eval_function rho1 f v = eval_function rho2 f v.
  Proof.
    intros H v. destruct f; cbn. apply H. reflexivity.
  Qed.

  Lemma eval_predicate_ext rho1 rho2 ar (P : predicate ar) :
    (forall n ar v, get_pred rho1 n ar v <-> get_pred rho2 n ar v)
      -> forall v, eval_predicate rho1 P v <-> eval_predicate rho2 P v.
  Proof.
    intros H v. destruct P; cbn. apply H. reflexivity.
  Qed.

  Lemma eval_ext rho1 rho2 t :
    (forall n, get_indi rho1 n = get_indi rho2 n)
    -> (forall n ar v, get_func rho1 n ar v = get_func rho2 n ar v)
    -> eval rho1 t = eval rho2 t.
  Proof.
    intros H1 H2. induction t.
    - apply H1.
    - cbn. enough (map (eval rho1) v = map (eval rho2) v) as -> by apply H2.
      now apply map_ext_forall.
    - cbn. f_equal. now apply map_ext_forall.
  Qed.

  Lemma sat_ext rho1 rho2 phi :
    rho1 rho2 -> sat rho1 phi <-> sat rho2 phi.
  Proof.
    revert rho1 rho2. induction phi; cbn; intros rho1 rho2 H.
    - easy.
    - destruct p.
      + rename t into v. enough (map (eval rho1) v = map (eval rho2) v) as <- by apply H.
        apply map_ext. induction v; firstorder. apply eval_ext; apply H.
      + rename t into v. enough (map (eval rho1) v = map (eval rho2) v) as <- by easy.
        apply map_ext. induction v; firstorder. apply eval_ext; apply H.
    - specialize (IHphi1 rho1 rho2); specialize (IHphi2 rho1 rho2).
      destruct b; cbn; tauto.
    - destruct q; split; cbn.
      + intros H1 x. eapply IHphi. 2: apply H1. now apply env_equiv_symm, env_equiv_cons_i.
      + intros H1 x. eapply IHphi. 2: apply H1. now apply env_equiv_cons_i.
      + intros [d H1]. exists d. eapply IHphi. 2: apply H1. now apply env_equiv_symm, env_equiv_cons_i.
      + intros [d H1]. exists d. eapply IHphi. 2: apply H1. now apply env_equiv_cons_i.
    - destruct q; split; cbn.
      + intros H1 f. eapply IHphi. 2: apply H1. now apply env_equiv_symm, env_equiv_cons_f.
      + intros H1 f. eapply IHphi. 2: apply H1. now apply env_equiv_cons_f.
      + intros [f H1]. exists f. eapply IHphi. 2: apply H1. now apply env_equiv_symm, env_equiv_cons_f.
      + intros [f H1]. exists f. eapply IHphi. 2: apply H1. now apply env_equiv_cons_f.
    - destruct q; split; cbn.
      + intros H1 P. eapply IHphi. 2: apply H1. now apply env_equiv_symm, env_equiv_cons_p.
      + intros H1 P. eapply IHphi. 2: apply H1. now apply env_equiv_cons_p.
      + intros [P H1]. exists P. eapply IHphi. 2: apply H1. now apply env_equiv_symm, env_equiv_cons_p.
      + intros [P H1]. exists P. eapply IHphi. 2: apply H1. now apply env_equiv_cons_p.
  Qed.

End SatExt.

Section BoundedSat.

  Context {Σ_funcs : funcs_signature}.
  Context {Σ_preds : preds_signature}.
  Variable domain : Type.
  Context {I : interp domain}.

  Lemma sat_ext_bounded_term t rho sigma :
    (forall x, ~ bounded_indi_term x t -> get_indi rho x = get_indi sigma x)
    -> (forall x ar, ~ bounded_func_term ar x t -> get_func rho x ar = get_func sigma x ar)
    -> eval rho t = eval sigma t.
  Proof.
    intros H1 H2. induction t; cbn.
    - apply H1. cbn. lia.
    - rewrite H2. f_equal. apply map_ext_in. intros t H. eapply Forall_in in IH.
      apply IH. intros x H3. apply H1. cbn. intros H4. apply H3. eapply Forall_in in H4.
      apply H4. trivial. intros x ar' H3. apply H2. cbn. intros [H4 H5]. apply H3. eapply Forall_in in H5.
      apply H5. easy. easy. cbn. lia.
    - f_equal. apply map_ext_in. intros t H. eapply Forall_in in IH.
      apply IH. intros x H3. apply H1. cbn. intros H4. apply H3. eapply Forall_in in H4.
      apply H4. trivial. intros x ar' H3. apply H2. cbn. intros H4. apply H3. eapply Forall_in in H4.
      apply H4. easy. easy.
  Qed.

  Lemma sat_ext_bounded phi rho sigma :
    (forall x, ~ bounded_indi x phi -> get_indi rho x = get_indi sigma x)
    -> (forall x ar, ~ bounded_func ar x phi -> get_func rho x ar = get_func sigma x ar)
    -> (forall x ar, ~ bounded_pred ar x phi -> get_pred rho x ar = get_pred sigma x ar)
    -> sat rho phi <-> sat sigma phi.
  Proof.
    revert rho sigma. induction phi; cbn; intros rho sigma H1 H2 H3.
    - reflexivity.
    - erewrite map_ext_in with (g := eval sigma); revgoals.
      intros ? H. apply sat_ext_bounded_term. intros x H4. apply H1. intros H5. apply H4.
      eapply Forall_in in H5. apply H5. easy. intros x ar' H4. apply H2. intros H5. apply H4.
      eapply Forall_in in H5. apply H5. easy. destruct p; cbn.
      + rewrite H3. reflexivity. cbn. lia.
      + reflexivity.
    - specialize (IHphi1 rho sigma). specialize (IHphi2 rho sigma).
      destruct b; cbn; setoid_rewrite IHphi1; try setoid_rewrite IHphi2; try reflexivity; clear IHphi1 IHphi2. all: firstorder.
    - destruct q; split.
      + intros H d. eapply IHphi; clear IHphi. 4: apply (H d). intros []. all: firstorder.
      + intros H d. eapply IHphi; clear IHphi. 4: apply (H d). intros []. all: firstorder.
      + intros [d H]. exists d. eapply IHphi; clear IHphi. 4: apply H. intros []. all: firstorder.
      + intros [d H]. exists d. eapply IHphi; clear IHphi. 4: apply H. intros []. all: firstorder.
    - destruct q; split.
      + intros H f. eapply IHphi; clear IHphi. 4: eapply (H f); trivial. 2: intros [] ar H4; cbn; destruct Nat.eq_dec as [->|]. all: firstorder.
      + intros H f. eapply IHphi; clear IHphi. 4: eapply (H f); trivial. 2: intros [] ar H4; cbn; destruct Nat.eq_dec as [->|]. all: firstorder.
      + intros [f H]. exists f. eapply IHphi; clear IHphi. 4: eapply H; trivial. 2: intros [] ar H4; cbn; destruct Nat.eq_dec as [->|]. all: firstorder.
      + intros [f H]. exists f. eapply IHphi; clear IHphi. 4: eapply H; trivial. 2: intros [] ar H4; cbn; destruct Nat.eq_dec as [->|]. all: firstorder.
    - destruct q; split.
      + intros H P. eapply IHphi; clear IHphi. 4: eapply (H P); trivial. 3: intros [] ar H4; cbn; destruct Nat.eq_dec as [->|]. all: firstorder.
      + intros H P. eapply IHphi; clear IHphi. 4: eapply (H P); trivial. 3: intros [] ar H4; cbn; destruct Nat.eq_dec as [->|]. all: firstorder.
      + intros [P H]. exists P. eapply IHphi; clear IHphi. 4: eapply H; trivial. 3: intros [] ar H4; cbn; destruct Nat.eq_dec as [->|]. all: firstorder.
      + intros [P H]. exists P. eapply IHphi; clear IHphi. 4: eapply H; trivial. 3: intros [] ar H4; cbn; destruct Nat.eq_dec as [->|]. all: firstorder.
  Qed.

  Lemma sat_ext_closed phi rho sigma :
    bounded_indi 0 phi -> (forall ar, bounded_func ar 0 phi) -> (forall ar, bounded_pred ar 0 phi)
    -> (sat rho phi <-> sat sigma phi).
  Proof.
    intros Bi Bf Bp. apply sat_ext_bounded.
    - intros x H. exfalso. apply H. eapply bounded_indi_up. 2: apply Bi. lia.
    - intros x ar H. exfalso. apply H. eapply bounded_func_up. 2: apply Bf. lia.
    - intros x ar H. exfalso. apply H. eapply bounded_pred_up. 2: apply Bp. lia.
  Qed.

  Lemma sat_ext_closed_funcfree phi rho sigma :
    funcfree phi -> bounded_indi 0 phi -> (forall ar, bounded_pred ar 0 phi)
    -> (sat rho phi <-> sat sigma phi).
  Proof.
    intros F Bi Bp. apply sat_ext_bounded.
    - intros x H. exfalso. apply H. eapply bounded_indi_up. 2: apply Bi. lia.
    - intros x ar H. exfalso. apply H. apply funcfree_bounded_func, F.
    - intros x ar H. exfalso. apply H. eapply bounded_pred_up. 2: apply Bp. lia.
  Qed.

  Lemma sat_ext_closed_FO phi rho sigma :
    first_order phi -> bounded_indi 0 phi -> (sat rho phi <-> sat sigma phi).
  Proof.
    intros F B. apply sat_ext_bounded.
    - intros x H. exfalso. apply H. eapply bounded_indi_up. 2: apply B. lia.
    - intros x ar H. exfalso. apply H. apply funcfree_bounded_func, firstorder_funcfree, F.
    - intros x ar H. exfalso. apply H. apply firstorder_bounded_pred, F.
  Qed.

End BoundedSat.

Section Subst.

  Context {Σ_funcs : funcs_signature}.
  Context {Σ_preds : preds_signature}.
  Variable domain : Type.
  Context {I : interp domain}.

  Lemma eval_function_subst_cons_shift_f (rho : env domain) ar (f : function ar) ar' (g : vec domain ar' -> domain) :
    eval_function rho f = eval_function get_indi rho, g .: get_func rho, get_pred rho (f[ ar']f).
  Proof.
    destruct f.
    - unfold econs, econs_func, econs_ar, shift, shift_f; cbn.
      destruct Nat.eq_dec as [->|]; cbn. now destruct Nat.eq_dec.
      destruct n. destruct Nat.eq_dec; try easy. now destruct Nat.eq_dec.
    - reflexivity.
  Qed.

  Lemma eval_predicate_subst_cons_shift_p (rho : env domain) ar (P : predicate ar) ar' (Q : vec domain ar' -> Prop) :
    eval_predicate rho P = eval_predicate get_indi rho, get_func rho, Q .: get_pred rho (P[ ar']p).
  Proof.
    destruct P.
    - unfold econs, econs_pred, econs_ar, shift, shift_p; cbn.
      destruct Nat.eq_dec as [->|]; cbn. now destruct Nat.eq_dec.
      destruct n. destruct Nat.eq_dec; try easy. now destruct Nat.eq_dec.
    - reflexivity.
  Qed.

  Lemma eval_subst_cons_shift_f (rho : env domain) ar (f : vec domain ar -> domain) t :
    eval rho t = eval get_indi rho, f .: get_func rho, get_pred rho t[ ar]f.
  Proof.
    induction t; cbn [eval].
    - reflexivity.
    - rewrite eval_function_subst_cons_shift_f with (g := f).
      cbn. f_equal. rewrite map_map. apply map_ext_forall, IH.
    - cbn. f_equal. rewrite map_map. apply map_ext_forall, IH.
  Qed.

  Lemma eval_comp_i (rho : env domain) σ t :
    eval rho (t[σ]i) = eval σ >> eval rho, get_func rho, get_pred rho t.
  Proof.
    induction t; cbn. reflexivity. all: f_equal; rewrite map_map; apply map_ext_forall, IH.
  Qed.

  Lemma eval_comp_f (rho : env domain) σ t :
    eval rho (t[σ]f) = eval get_indi rho, σ >>> eval_function rho, get_pred rho t.
  Proof.
    induction t; cbn. reflexivity. all: f_equal; rewrite map_map; apply map_ext_forall, IH.
  Qed.

  Lemma sat_comp_i rho σ phi :
    sat rho (phi[σ]i) <-> sat σ >> eval rho, get_func rho, get_pred rho phi.
  Proof.
    induction phi in rho, σ |- *; cbn.
    - reflexivity.
    - destruct p; cbn; erewrite map_map, map_ext; try reflexivity;
      induction t; firstorder using eval_comp_i.
    - specialize (IHphi1 rho σ); specialize (IHphi2 rho σ).
      destruct b; cbn; intuition easy.
    - destruct q.
      + setoid_rewrite IHphi. split; intros H d; eapply sat_ext.
        2, 4: apply (H d). all: clear; intros []; split; try easy;
        cbn; erewrite eval_comp_i; now destruct rho.
      + setoid_rewrite IHphi. split; intros [d H]; exists d; eapply sat_ext.
        2, 4: apply H. all: clear; intros []; split; try easy;
        cbn; erewrite eval_comp_i; now destruct rho.
    - destruct q.
      + setoid_rewrite IHphi; split; intros H f; eapply sat_ext.
        2, 4: apply (H f). all: clear; split; try easy. 2: symmetry.
        all: apply eval_subst_cons_shift_f.
      + setoid_rewrite IHphi; split; intros [f H]; exists f; eapply sat_ext.
        2, 4: apply H. all: clear; split; try easy. 2: symmetry.
        all: apply eval_subst_cons_shift_f.
    - destruct q.
      + setoid_rewrite IHphi; split; intros H P; eapply sat_ext.
        2, 4: apply (H P). all: clear; split; try easy; now apply eval_ext.
      + setoid_rewrite IHphi; split; intros [P H]; exists P; eapply sat_ext.
        2, 4: apply H. all: clear; split; try easy; now apply eval_ext.
  Qed.

  Lemma sat_comp_f rho σ phi :
    sat rho (phi[σ]f) <-> sat get_indi rho, σ >>> eval_function rho, get_pred rho phi.
  Proof.
    induction phi in rho, σ |- *; cbn.
    - reflexivity.
    - destruct p; cbn; erewrite map_map, map_ext; try reflexivity;
      induction t; firstorder using eval_comp_f.
    - specialize (IHphi1 rho σ); specialize (IHphi2 rho σ).
      destruct b; cbn; intuition easy.
    - destruct q.
      + setoid_rewrite IHphi. split; intros H d; eapply sat_ext.
        2, 4: apply (H d). all: easy.
      + setoid_rewrite IHphi. split; intros [d H]; exists d; eapply sat_ext.
        2, 4: apply H. all: easy.
    - destruct q.
      + setoid_rewrite IHphi; split; intros H f; eapply sat_ext.
        2, 4: apply (H f).
        all: clear; intros []; repeat split; try easy; cbn; destruct Nat.eq_dec as [->|]; cbn.
        1, 5: destruct Nat.eq_dec; try easy; rewrite Eqdep_dec.UIP_dec with (p1 := e) (p2 := eq_refl); try easy; decide equality.
        1-3: now rewrite eval_function_subst_cons_shift_f with (g := f).
        all: now rewrite <- eval_function_subst_cons_shift_f with (g := f).
      + setoid_rewrite IHphi; split; intros [f H]; exists f; eapply sat_ext.
        2, 4: apply H.
        all: clear; intros []; repeat split; try easy; cbn; destruct Nat.eq_dec as [->|]; cbn.
        1, 5: destruct Nat.eq_dec; try easy; rewrite Eqdep_dec.UIP_dec with (p1 := e) (p2 := eq_refl); try easy; decide equality.
        1-3: now rewrite eval_function_subst_cons_shift_f with (g := f).
        all: now rewrite <- eval_function_subst_cons_shift_f with (g := f).
    - destruct q.
      + setoid_rewrite IHphi; split; intros H P; eapply sat_ext.
        2, 4: apply (H P). all: clear; intros; split; try easy; now apply eval_ext.
      + setoid_rewrite IHphi; split; intros [P H]; exists P; eapply sat_ext.
        2, 4: apply H. all: clear; intros; split; try easy; now apply eval_ext.
  Qed.

  Lemma sat_comp_p rho σ phi :
    sat rho (phi[σ]p) <-> sat get_indi rho, get_func rho, σ >>> eval_predicate rho phi.
  Proof.
    induction phi in rho, σ |- *; cbn.
    - reflexivity.
    - destruct p; cbn; rename t into v;
      enough (map (eval rho) v = map (eval get_indi rho, get_func rho, σ >>> eval_predicate rho) v) as -> by reflexivity;
      apply map_ext_forall; induction v; firstorder; now apply eval_ext.
    - specialize (IHphi1 rho σ); specialize (IHphi2 rho σ).
      destruct b; cbn; intuition easy.
    - destruct q.
      + setoid_rewrite IHphi. split; intros H d; eapply sat_ext.
        2, 4: apply (H d). all: easy.
      + setoid_rewrite IHphi. split; intros [d H]; exists d; eapply sat_ext.
        2, 4: apply H. all: easy.
    - destruct q.
      + setoid_rewrite IHphi. split; intros H f; eapply sat_ext.
        2, 4: apply (H f). all: easy.
      + setoid_rewrite IHphi. split; intros [f H]; exists f; eapply sat_ext.
        2, 4: apply H. all: easy.
    - destruct q.
      + setoid_rewrite IHphi; split; intros H P; eapply sat_ext.
        2, 4: apply (H P).
        all: clear; intros []; split; try easy; split; try easy; cbn; destruct Nat.eq_dec as [->|]; cbn.
        1, 5: destruct Nat.eq_dec; try easy; rewrite Eqdep_dec.UIP_dec with (p1 := e) (p2 := eq_refl); try easy; decide equality.
        1-3: now rewrite eval_predicate_subst_cons_shift_p with (Q := P).
        all: now rewrite <- eval_predicate_subst_cons_shift_p with (Q := P).
      + setoid_rewrite IHphi; split; intros [P H]; exists P; eapply sat_ext.
        2, 4: apply H.
        all: clear; intros []; split; try easy; split; try easy; cbn; destruct Nat.eq_dec as [->|]; cbn.
        1, 5: destruct Nat.eq_dec; try easy; rewrite Eqdep_dec.UIP_dec with (p1 := e) (p2 := eq_refl); try easy; decide equality.
        1-3: now rewrite eval_predicate_subst_cons_shift_p with (Q := P).
        all: now rewrite <- eval_predicate_subst_cons_shift_p with (Q := P).
  Qed.

End Subst.