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

#[export] Instance econs_subst_indi `{funcs_signature} : Econs _ _ := econs_indi term.
#[export] Instance econs_subst_func `{funcs_signature} ar : Econs _ _ := econs_ar ar function.
#[export] Instance econs_subst_pred `{preds_signature} ar : Econs _ _ := econs_ar ar predicate.


Class IdSubst X := ids : nat -> X.
Class Shift X := shift : X.

#[export] Instance var_indi' `{funcs_signature} : IdSubst term := var_indi.
#[export] Instance var_func' `{funcs_signature} : IdSubst (forall ar, function ar) := var_func.
#[export] Instance var_pred' `{preds_signature} : IdSubst (forall ar, predicate ar) := var_pred.

#[export] Instance shift_i `{funcs_signature} : Shift (nat -> term) :=
  fun n => var_indi (S n).

#[export] Instance shift_f `{funcs_signature} : Shift (nat -> nat -> forall ar, function ar) :=
  fun ar n ar' => if Nat.eq_dec ar ar' then @var_func _ (S n) ar' else @var_func _ n ar'.

#[export] Instance shift_p `{preds_signature} : Shift (nat -> nat -> forall ar, predicate ar) :=
  fun ar n ar' => if Nat.eq_dec ar ar' then @var_pred _ (S n) ar' else @var_pred _ n ar'.


Class Subst_i `{funcs_signature} X := subst_i : (nat -> term) -> X -> X.
Class Subst_f `{funcs_signature} X := subst_f : (nat -> (forall ar, function ar)) -> X -> X.
Class Subst_p `{preds_signature} X := subst_p : (nat -> (forall ar, predicate ar)) -> X -> X.

Local Notation "X [ σ ]i" := (subst_i σ X) (at level 7, left associativity, format "X '/' [ σ ]i").
Local Notation "X [ σ ]f" := (subst_f σ X) (at level 7, left associativity, format "X '/' [ σ ]f").
Local Notation "X [ σ ]p" := (subst_p σ X) (at level 7, left associativity, format "X '/' [ σ ]p").

#[export] Instance subst_function `{funcs_signature} {ar} : Subst_f (function ar) := fun σf f => match f with
  | var_func f => σf f ar
  | f => f
end.

#[export] Instance subst_term_i `{funcs_signature} : Subst_i term := fix subst_term_i σi t := match t with
  | var_indi x => σi x
  | func f v => func f (Vector.map (subst_term_i σi) v)
end.

#[export] Instance subst_term_f `{funcs_signature} : Subst_f term := fix subst_term_f σf t := match t with
  | var_indi x => var_indi x
  | func f v => func (subst_function σf f) (Vector.map (subst_term_f σf) v)
end.

#[export] Instance subst_predicate `{preds_signature} {ar} : Subst_p (predicate ar) := fun σp P => match P with
  | var_pred P => σp P ar
  | P => P
end.

Definition up_i `{funcs_signature} (σi : nat -> term) := econs (var_indi 0) (fun x => (σi x)[shift]i).
Definition up_f `{funcs_signature} (σf : nat -> forall ar, function ar) ar := econs (@var_func _ 0 ar) (fun x ar' => (σf x ar')[shift ar]f).
Definition up_p `{preds_signature} (σp : nat -> forall ar, predicate ar) ar := econs (@var_pred _ 0 ar) (fun x ar' => (σp x ar')[shift ar]p).

#[export] Instance subst_form_i `{funcs_signature, preds_signature, operators} : Subst_i form := fix subst_form_i σi phi := match phi with
  | fal => fal
  | atom P v => atom P (Vector.map (subst_term_i σi) v)
  | bin op phi psi => bin op (subst_form_i σi phi) (subst_form_i σi psi)
  | quant_indi op phi => quant_indi op (subst_form_i (up_i σi) phi)
  | quant_func op ar phi => quant_func op ar (subst_form_i (fun n => (σi n)[shift ar]f ) phi)
  | quant_pred op ar phi => quant_pred op ar (subst_form_i σi phi)
end.

#[export] Instance subst_form_f `{funcs_signature, preds_signature, operators} : Subst_f form := fix subst_form_f σf phi := match phi with
  | fal => fal
  | atom P v => atom P (Vector.map (subst_term_f σf) v)
  | bin op phi psi => bin op (subst_form_f σf phi) (subst_form_f σf psi)
  | quant_indi op phi => quant_indi op (subst_form_f σf phi)
  | quant_func op ar phi => quant_func op ar (subst_form_f (up_f σf ar) phi)
  | quant_pred op ar phi => quant_pred op ar (subst_form_f σf phi)
end.

#[export] Instance subst_form_p `{funcs_signature, preds_signature, operators} : Subst_p form := fix subst_form_p σp phi := match phi with
  | fal => fal
  | atom P v => atom (subst_predicate σp P) v
  | bin op phi psi => bin op (subst_form_p σp phi) (subst_form_p σp psi)
  | quant_indi op phi => quant_indi op (subst_form_p σp phi)
  | quant_func op ar phi => quant_func op ar (subst_form_p σp phi)
  | quant_pred op ar phi => quant_pred op ar (subst_form_p (up_p σp ar) phi)
end.

Declare Scope subst_scope.
Open Scope subst_scope.

Module SubstNotations.

  Notation "x .: sigma" := (econs x sigma) (at level 70, right associativity) : subst_scope.
  Notation "↑" := (shift) : subst_scope.
  Notation "f >> g" := (fun x => g (f x)) (at level 50) : subst_scope.
  Notation "f >>> g" := (fun x y => g (f x y)) (at level 50) : subst_scope.
  Notation "x '..'" := (econs x ids) (at level 1, format "x ..") : subst_scope.

  Notation "X [ σ ]i" := (subst_i σ X) (at level 7, left associativity, format "X '/' [ σ ]i").
  Notation "X [ σ ]f" := (subst_f σ X) (at level 7, left associativity, format "X '/' [ σ ]f").
  Notation "X [ σ ]p" := (subst_p σ X) (at level 7, left associativity, format "X '/' [ σ ]p").

End SubstNotations.