Library MathComp.finfun

Require Import ssreflect ssrfun ssrbool eqtype ssrnat seq choice fintype tuple.


Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Section Def.

Variables (aT : finType) (rT : Type).

Inductive finfun_type : predArgType := Finfun of #|aT|.-tuple rT.

Definition finfun_of of phant (aTrT) := finfun_type.

Identity Coercion type_of_finfun : finfun_of >-> finfun_type.

Definition fgraph f := let: Finfun t := f in t.

Canonical finfun_subType := Eval hnf in [newType for fgraph].

End Def.

Notation "{ 'ffun' fT }" := (finfun_of (Phant fT))
  (at level 0, format "{ 'ffun' '[hv' fT ']' }") : type_scope.
Definition finexp_domFinType n := ordinal_finType n.
Notation "T ^ n" := (@finfun_of (finexp_domFinType n) T (Phant _)) : type_scope.

Notation Local fun_of_fin_def :=
  (fun aT rT f xtnth (@fgraph aT rT f) (enum_rank x)).

Notation Local finfun_def := (fun aT rT f ⇒ @Finfun aT rT (codom_tuple f)).

Module Type FunFinfunSig.
Parameter fun_of_fin : aT rT, finfun_type aT rTaTrT.
Parameter finfun : (aT : finType) rT, (aTrT) → {ffun aTrT}.
Axiom fun_of_finE : fun_of_fin = fun_of_fin_def.
Axiom finfunE : finfun = finfun_def.
End FunFinfunSig.

Module FunFinfun : FunFinfunSig.
Definition fun_of_fin := fun_of_fin_def.
Definition finfun := finfun_def.
Lemma fun_of_finE : fun_of_fin = fun_of_fin_def. Proof. by []. Qed.
Lemma finfunE : finfun = finfun_def. Proof. by []. Qed.
End FunFinfun.

Notation fun_of_fin := FunFinfun.fun_of_fin.
Notation finfun := FunFinfun.finfun.
Coercion fun_of_fin : finfun_type >-> Funclass.
Canonical fun_of_fin_unlock := Unlockable FunFinfun.fun_of_finE.
Canonical finfun_unlock := Unlockable FunFinfun.finfunE.

Notation "[ 'ffun' x : aT => F ]" := (finfun (fun x : aTF))
  (at level 0, x ident, only parsing) : fun_scope.

Notation "[ 'ffun' : aT => F ]" := (finfun (fun _ : aTF))
  (at level 0, only parsing) : fun_scope.

Notation "[ 'ffun' x => F ]" := [ffun x : _ F]
  (at level 0, x ident, format "[ 'ffun' x => F ]") : fun_scope.

Notation "[ 'ffun' => F ]" := [ffun : _ F]
  (at level 0, format "[ 'ffun' => F ]") : fun_scope.

Definition fmem aT rT (pT : predType rT) (f : aTpT) := [fun x mem (f x)].

Section PlainTheory.

Variables (aT : finType) (rT : Type).
Notation fT := {ffun aTrT}.
Implicit Types (f : fT) (R : pred rT).

Canonical finfun_of_subType := Eval hnf in [subType of fT].

Lemma tnth_fgraph f i : tnth (fgraph f) i = f (enum_val i).
Proof. by rewrite [@fun_of_fin]unlock enum_valK. Qed.

Lemma ffunE (g : aTrT) : finfun g =1 g.
Proof.
movex; rewrite [@finfun]unlock unlock tnth_map.
by rewrite -[tnth _ _]enum_val_nth enum_rankK.
Qed.

Lemma fgraph_codom f : fgraph f = codom_tuple f.
Proof.
apply: eq_from_tnthi; rewrite [@fun_of_fin]unlock tnth_map.
by congr tnth; rewrite -[tnth _ _]enum_val_nth enum_valK.
Qed.

Lemma codom_ffun f : codom f = val f.
Proof. by rewrite /= fgraph_codom. Qed.

Lemma ffunP f1 f2 : f1 =1 f2 f1 = f2.
Proof.
split⇒ [eq_f12 | → //]; do 2!apply: val_inj ⇒ /=.
by rewrite !fgraph_codom /= (eq_codom eq_f12).
Qed.

Lemma ffunK : cancel (@fun_of_fin aT rT) (@finfun aT rT).
Proof. by movef; apply/ffunP/ffunE. Qed.

Definition family_mem mF := [pred f : fT | [ x, in_mem (f x) (mF x)]].

Lemma familyP (pT : predType rT) (F : aTpT) f :
  reflect ( x, f x \in F x) (f \in family_mem (fmem F)).
Proof. exact: forallP. Qed.

Definition ffun_on_mem mR := family_mem (fun _mR).

Lemma ffun_onP R f : reflect ( x, f x \in R) (f \in ffun_on_mem (mem R)).
Proof. exact: forallP. Qed.

End PlainTheory.

Notation family F := (family_mem (fun_of_simpl (fmem F))).
Notation ffun_on R := (ffun_on_mem _ (mem R)).

Implicit Arguments familyP [aT rT pT F f].
Implicit Arguments ffun_onP [aT rT R f].


Lemma nth_fgraph_ord T n (x0 : T) (i : 'I_n) f : nth x0 (fgraph f) i = f i.
Proof.
by rewrite -{2}(enum_rankK i) -tnth_fgraph (tnth_nth x0) enum_rank_ord.
Qed.

Section Support.

Variables (aT : Type) (rT : eqType).

Definition support_for y (f : aTrT) := [pred x | f x != y].

Lemma supportE x y f : (x \in support_for y f) = (f x != y). Proof. by []. Qed.

End Support.

Notation "y .-support" := (support_for y)
  (at level 2, format "y .-support") : fun_scope.

Section EqTheory.

Variables (aT : finType) (rT : eqType).
Notation fT := {ffun aTrT}.
Implicit Types (y : rT) (D : pred aT) (R : pred rT) (f : fT).

Lemma supportP y D g :
  reflect ( x, x \notin Dg x = y) (y.-support g \subset D).
Proof.
by apply: (iffP subsetP) ⇒ Dg x; [apply: contraNeq | apply: contraR] ⇒ /Dg→.
Qed.

Definition finfun_eqMixin :=
  Eval hnf in [eqMixin of finfun_type aT rT by <:].
Canonical finfun_eqType := Eval hnf in EqType _ finfun_eqMixin.
Canonical finfun_of_eqType := Eval hnf in [eqType of fT].

Definition pfamily_mem y mD (mF : aTmem_pred rT) :=
  family (fun i : aTif in_mem i mD then pred_of_simpl (mF i) else pred1 y).

Lemma pfamilyP (pT : predType rT) y D (F : aTpT) f :
  reflect (y.-support f \subset D {in D, x, f x \in F x})
          (f \in pfamily_mem y (mem D) (fmem F)).
Proof.
apply: (iffP familyP) ⇒ [/= f_pfam | [/supportP f_supp f_fam] x].
  split⇒ [|x Ax]; last by have:= f_pfam x; rewrite Ax.
  by apply/subsetPx; case: ifP (f_pfam x) ⇒ //= _ fx0 /negP[].
by case: ifPnAx /=; rewrite inE /= (f_fam, f_supp).
Qed.

Definition pffun_on_mem y mD mR := pfamily_mem y mD (fun _mR).

Lemma pffun_onP y D R f :
  reflect (y.-support f \subset D {subset image f D R})
          (f \in pffun_on_mem y (mem D) (mem R)).
Proof.
apply: (iffP (pfamilyP y D (fun _R) f)) ⇒ [] [-> f_fam]; split⇒ //.
  by move_ /imageP[x Ax ->]; exact: f_fam.
by movex Ax; apply: f_fam; apply/imageP; x.
Qed.

End EqTheory.
Canonical exp_eqType (T : eqType) n := [eqType of T ^ n].

Implicit Arguments supportP [aT rT y D g].
Notation pfamily y D F := (pfamily_mem y (mem D) (fun_of_simpl (fmem F))).
Notation pffun_on y D R := (pffun_on_mem y (mem D) (mem R)).

Definition finfun_choiceMixin aT (rT : choiceType) :=
  [choiceMixin of finfun_type aT rT by <:].
Canonical finfun_choiceType aT rT :=
  Eval hnf in ChoiceType _ (finfun_choiceMixin aT rT).
Canonical finfun_of_choiceType (aT : finType) (rT : choiceType) :=
  Eval hnf in [choiceType of {ffun aTrT}].
Canonical exp_choiceType (T : choiceType) n := [choiceType of T ^ n].

Definition finfun_countMixin aT (rT : countType) :=
  [countMixin of finfun_type aT rT by <:].
Canonical finfun_countType aT (rT : countType) :=
  Eval hnf in CountType _ (finfun_countMixin aT rT).
Canonical finfun_of_countType (aT : finType) (rT : countType) :=
  Eval hnf in [countType of {ffun aTrT}].
Canonical finfun_subCountType aT (rT : countType) :=
  Eval hnf in [subCountType of finfun_type aT rT].
Canonical finfun_of_subCountType (aT : finType) (rT : countType) :=
  Eval hnf in [subCountType of {ffun aTrT}].


Section FinTheory.

Variables aT rT : finType.
Notation fT := {ffun aTrT}.
Notation ffT := (finfun_type aT rT).
Implicit Types (D : pred aT) (R : pred rT) (F : aTpred rT).

Definition finfun_finMixin := [finMixin of ffT by <:].
Canonical finfun_finType := Eval hnf in FinType ffT finfun_finMixin.
Canonical finfun_subFinType := Eval hnf in [subFinType of ffT].
Canonical finfun_of_finType := Eval hnf in [finType of fT for finfun_finType].
Canonical finfun_of_subFinType := Eval hnf in [subFinType of fT].

Lemma card_pfamily y0 D F :
  #|pfamily y0 D F| = foldr muln 1 [seq #|F x| | x in D].
Proof.
rewrite /image_mem; transitivity #|pfamily y0 (enum D) F|.
  by apply/eq_cardf; apply/eq_forallbx /=; rewrite mem_enum.
elim: {D}(enum D) (enum_uniq D) ⇒ /= [_|x0 s IHs /andP[s'x0 /IHs<-{IHs}]].
  apply: eq_card1 [ffun y0] _ _f.
  apply/familyP/eqP⇒ [y0_f|-> x]; last by rewrite ffunE inE.
  by apply/ffunPx; rewrite ffunE (eqP (y0_f x)).
pose g (xf : rT × fT) := finfun [eta xf.2 with x0 |-> xf.1].
have gK: cancel (fun f : fT(f x0, g (y0, f))) g.
  by movef; apply/ffunPx; do !rewrite ffunE /=; case: eqP ⇒ // →.
rewrite -cardX -(card_image (can_inj gK)); apply: eq_card ⇒ [] [y f] /=.
apply/imageP/andP⇒ [[f0 /familyP/=Ff0] [{f}-> ->]| [Fy /familyP/=Ff]].
  split; first by have:= Ff0 x0; rewrite /= mem_head.
  apply/familyPx; have:= Ff0 x; rewrite ffunE inE /=.
  by case: eqP ⇒ //= → _; rewrite ifN ?inE.
(g (y, f)).
  by apply/familyPx; have:= Ff x; rewrite ffunE /= inE; case: eqP ⇒ // →.
congr (_, _); last apply/ffunPx; do !rewrite ffunE /= ?eqxx //.
by case: eqP ⇒ // ->{x}; apply/eqP; have:= Ff x0; rewrite ifN.
Qed.

Lemma card_family F : #|family F| = foldr muln 1 [seq #|F x| | x : aT].
Proof.
have [y0 _ | rT0] := pickP rT; first exact: (card_pfamily y0 aT).
rewrite /image_mem; case DaT: (enum aT) ⇒ [{rT0}|x0 e] /=; last first.
  by rewrite !eq_card0 // ⇒ [f | y]; [have:= rT0 (f x0) | have:= rT0 y].
have{DaT} no_aT P (x : aT) : P by have:= mem_enum aT x; rewrite DaT.
apply: eq_card1 [ffun x no_aT rT x] _ _f.
by apply/familyP/eqP_; [apply/ffunP | ] ⇒ x; apply: no_aT.
Qed.

Lemma card_pffun_on y0 D R : #|pffun_on y0 D R| = #|R| ^ #|D|.
Proof.
rewrite (cardE D) card_pfamily /image_mem.
by elim: (enum D) ⇒ //= _ e ->; rewrite expnS.
Qed.

Lemma card_ffun_on R : #|ffun_on R| = #|R| ^ #|aT|.
Proof.
rewrite card_family /image_mem cardT.
by elim: (enum aT) ⇒ //= _ e ->; rewrite expnS.
Qed.

Lemma card_ffun : #|fT| = #|rT| ^ #|aT|.
Proof. by rewrite -card_ffun_on; apply/esym/eq_cardf; apply/forallP. Qed.

End FinTheory.
Canonical exp_finType (T : finType) n := [finType of T ^ n].