From HB Require Import structures.
From mathcomp Require Import boot order algebra boolp classical_sets.
From mathcomp Require Import fsbigop functions reals topology separation_axioms.
From mathcomp Require Import ereal sequences numfun measure measurable_realfun.
From mathcomp Require Import lebesgue_measure lebesgue_integral.

(**md**************************************************************************)
(* # The Giry monad                                                           *)
(*                                                                            *)
(* ```                                                                        *)
(*         m1 ≡μ m2 == m1 S = m2 S for any measurable set S                   *)
(*         giry T R == the type of subprobability measures over T             *)
(*      giry_ev P A == the evaluation map `giryM T R -> [0, 1]`               *)
(* preimg_giry_ev A == preimage of the sigma-algebra of extended real numbers *)
(*                     via the function giry_ev ^~ A                          *)
(*  giry_measurable == the sigma-algebra generated by the countable unions of *)
(*                     preimg_giry_ev                                         *)
(*     giry_display == display of giry_measurable                             *)
(*    giry_int mu f := \int[mu]_x f x                                         *)
(*      giry_map mf == the map of type giry T1 R -> giry T2 R                 *)
(*                     where mf is a proof that f : T1 -> T2 is measurable    *)
(*       giry_ret x == the unit of the Giry monad, i.e., \d_x                 *)
(* @giry_join _ T R == the multiplication of the Giry monad                   *)
(*                     type : giry (giry T R) R -> giry T R                   *)
(*  giry_bind mu mf == the bind with mu : giry T1 R and f : T1 -> giry T2 R   *)
(*        giry_prod == product of type                                        *)
(*                     giry T1 R * giry T2 R -> giry (T1 * T2) R              *)
(* ```                                                                        *)
(*                                                                            *)
(******************************************************************************)

Reserved Notation "m >>= f" (at level 49).

Unset SsrOldRewriteGoalsOrder.  (* remove the line when requiring MathComp >= 2.6 *)
Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Import Order.TTheory GRing.Theory Num.Def Num.Theory.

Definition measure_eq {d} {T : measurableType d} {R : realType} :
  measure T R -> measure T R -> Prop :=
  fun m1 m2 => forall S : set T, measurable S -> m1 S = m2 S.
Notation "x ≡μ y" := (measure_eq x y) (at level 70).

Global Hint Extern 0 (_ ≡μ _) => reflexivity : core.

Local Open Scope classical_set_scope.
Local Open Scope ereal_scope.

Section giry_def.
Context d (T : measurableType d) (R : realType).

Definition giry : Type := @subprobability d T R.

HB.instance Definition _ := gen_eqMixin giry.
HB.instance Definition _ := gen_choiceMixin giry.
HB.instance Definition _ := isPointed.Build giry mzero.

Definition giry_ev (mu : giry) (A : set T) := mu A.

Definition preimg_giry_ev (A : set T) : set_system giry :=
  preimage_set_system [set: giry] (giry_ev ^~ A) measurable.

Definition giry_measurable := <<s \bigcup_(A in measurable) preimg_giry_ev A >>.

Let giry_measurable0 : giry_measurable set0.
Proof. exact: sigma_algebra0. Qed.

Let giry_measurableC (U : set giry) :
  giry_measurable U -> giry_measurable (~` U).
Proof. exact: sigma_algebraC. Qed.

Let giry_measurableU (F : (set giry)^nat) :
  (forall i, giry_measurable (F i)) -> giry_measurable (\bigcup_i F i).
Proof. exact: sigma_algebra_bigcup. Qed.

Definition giry_display : measure_display.
Proof. by constructor. Qed.

HB.instance Definition _ :=
  @isMeasurable.Build giry_display giry giry_measurable
    giry_measurable0 giry_measurableC giry_measurableU.

Lemma measurable_giry_ev (A : set T) : measurable A ->
  measurable_fun [set: giry] (giry_ev ^~ A).
Proof.
move=> mS.
apply: (@measurability giry_display _ giry _ setT (giry_ev ^~ A) measurable).
  by rewrite smallest_id//; exact: sigma_algebra_measurable.
apply: subset_trans; last exact: sub_gen_smallest.
exact: (bigcup_sup mS).
Qed.

End giry_def.
Arguments giry_ev {d T R} mu A.

Section giry_integral.
Context {d} {T : measurableType d} {R : realType}.

Definition giry_int (mu : giry T R) (f : T -> \bar R) := \int[mu]_x f x.

Import HBNNSimple.
Import MeasurableR.

(**md The idea is to reconstruct f from simple functions, then use measurability
  of giry_ev. Reference: Tom Avery. Codensity and the Giry monad.
  https://arxiv.org/pdf/1410.4432.
 *)
Lemma measurable_giry_int (f : T -> \bar R) :
  measurable_fun [set: T] f -> (forall x, 0 <= f x) ->
  measurable_fun [set: giry T R] (giry_int ^~ f).
Proof.
move=> mf h0.
pose g := nnsfun_approx measurableT mf.
pose Eg := fun n => EFin \o g n.
pose intEg := fun n mu => giry_int mu (Eg n).
have mintgE n : measurable_fun [set: giry T R] (intEg n).
  rewrite /intEg /giry_int/=.
  under eq_fun do rewrite integralT_nnsfun sintegralE.
  apply: emeasurable_fsum => //= r.
  by apply: measurable_funeM => //=; exact: measurable_giry_ev.
apply: (emeasurable_fun_cvg _ (giry_int ^~ f) mintgE) => mu _.
rewrite (_ : giry_int mu f = \int[mu]_x limn (Eg ^~ x)).
  apply: eq_integral => t _.
  by apply/esym/cvg_lim  => //; exact/cvg_nnsfun_approx.
apply: cvg_monotone_convergence => //.
- by move=> n; exact/measurable_EFinP.
- by move=> n x _; rewrite lee_fin.
- by move=> t _ n m nm; apply/lefP/nd_nnsfun_approx.
Qed.

End giry_integral.
Arguments giry_int {d T R} mu f.

Section measurable_giry_codensity.
Context d1 {T1 : measurableType d1}.

Lemma measurable_giry_codensity d2 {T2 : measurableType d2} {R : realType}
    (D : set T1) (f : T1 -> giry T2 R) :
  measurable D ->
  (forall B, measurable B -> measurable_fun D (f ^~ B)) ->
  measurable_fun D f.
Proof.
move=> mD mf.
pose G : set_system (giry T2 R) := \bigcup_(B in measurable) preimg_giry_ev B.
apply: (measurability G) => //= _ [_ [C mC [Z mZ] <-] <-].
by rewrite setTI; exact: mf.
Qed.

End measurable_giry_codensity.

Section giry_map.
Context {d1} {d2} {T1 : measurableType d1} {T2 : measurableType d2}
  {R : realType}.
Variables (f : T1 -> T2) (mf : measurable_fun [set: T1] f) (mu1 : giry T1 R).

Let map := pushforward mu1 f.

Let map0 : map set0 = 0. Proof. exact: measure0. Qed.

Let map_ge0 A : 0 <= map A. Proof. exact: measure_ge0. Qed.

Let map_sigma_additive : semi_sigma_additive map.
Proof. exact: measure_semi_sigma_additive. Qed.

HB.instance Definition _ := isMeasure.Build _ _ _ map
  map0 map_ge0 map_sigma_additive.

(* used to work before the change of definition of pushforward
HB.instance Definition _ := Measure.on gMap_ev.
the change of definition was maybe not a good idea...
https://github.com/math-comp/analysis/pull/1661
*)

Let map_setT : map [set: T2] <= 1.
Proof. by rewrite /map /pushforward /= sprobability_setT. Qed.

HB.instance Definition _ := Measure_isSubProbability.Build _ _ _ map map_setT.

Definition giry_map : giry T2 R := map.

End giry_map.

Section giry_map_lemmas.
Context d1 d2 (T1 : measurableType d1) (T2 : measurableType d2) (R : realType).
Variable f : T1 -> T2.
Hypothesis mf : measurable_fun [set: T1] f.

Lemma measurable_giry_map : measurable_fun [set: giry T1 R] (giry_map mf).
Proof.
rewrite /giry_map.
apply: measurable_giry_codensity => // B mB.
apply: measurable_giry_ev.
by rewrite -(setTI (f @^-1` B)); exact: mf.
Qed.

Lemma giry_int_map (mu : giry T1 R) (h : T2 -> \bar R) :
  measurable_fun [set: T2] h -> (forall x, 0 <= h x) ->
  giry_int (giry_map mf mu) h = giry_int mu (h \o f).
Proof. by move=> mh h0; exact: ge0_integral_pushforward. Qed.

Lemma giry_map_dirac (mu1 : giry T1 R) (B : set T2) : measurable B ->
  giry_map mf mu1 B = \int[mu1]_x (\d_(f x))%R B.
Proof.
move=> mB.
rewrite -[in LHS](setIT B) -[LHS]integral_indic// [LHS]giry_int_map//.
exact/measurable_EFinP/measurable_indic.
Qed.

End giry_map_lemmas.

Section giry_ret.
Context {d} {T : measurableType d} {R : realType}.

Definition giry_ret (x : T) : giry T R := (\d_x)%R.

Lemma measurable_giry_ret : measurable_fun [set: T] giry_ret.
Proof.
by apply: measurable_giry_codensity => //; exact: measurable_fun_dirac.
Qed.

Lemma giry_int_ret (x : T) (f : T -> \bar R) : measurable_fun [set: T] f ->
  giry_int (giry_ret x) f = f x.
Proof.
by move=> mf; rewrite /giry_int /giry_ret integral_dirac// diracT mul1e.
Qed.

End giry_ret.

Section giry_join.
Context {d} {T : measurableType d} {R : realType}.
Variable M : giry (giry T R) R.

Let join A := giry_int M (giry_ev ^~ A).

Let join0 : join set0 = 0.
Proof. by rewrite /join /giry_ev /giry_int/= integral0_eq. Qed.

Let join_ge0 A : 0 <= join A. Proof. by rewrite /join integral_ge0. Qed.

Let join_semi_sigma_additive : semi_sigma_additive join.
Proof.
move=> F mF tF _; rewrite [X in _ --> X](_ : _ =
    giry_int M (fun x => \sum_(0 <= k <oo) x (F k))).
  apply: eq_integral => mu _.
  by apply/esym/cvg_lim => //; exact: measure_sigma_additive.
rewrite [X in X @ _](_ : _ =
    (fun n => giry_int M (fun mu => \sum_(0 <= i < n) mu (F i)))).
  apply/funext => n; rewrite -ge0_integral_sum//.
  by move=> ?; exact: measurable_giry_ev.
apply: cvg_monotone_convergence => //.
- by move=> n; apply: emeasurable_sum => m; exact: measurable_giry_ev.
- by move=> n x _; rewrite sume_ge0.
- by move=> x _ m n mn; exact: ereal_nondecreasing_series.
Qed.

HB.instance Definition _ := isMeasure.Build d _ R join
  join0 join_ge0 join_semi_sigma_additive.

Let join_setT : join [set: T] <= 1.
Proof.
rewrite (@le_trans _ _ (\int[M]_x `|giry_ev x [set: T]|))//; last first.
  rewrite (le_trans _ (@sprobability_setT _ _ _ M))//.
  rewrite -[leRHS]mul1e integral_le_bound//.
    exact: measurable_giry_ev.
  by apply/aeW => x _; rewrite gee0_abs// sprobability_setT.
rewrite ge0_le_integral//=.
- exact: measurable_giry_ev.
- by apply: measurableT_comp => //; exact: measurable_giry_ev.
- by move=> x _; rewrite gee0_abs.
Qed.

HB.instance Definition _ := Measure_isSubProbability.Build _ _ _ join join_setT.

Definition giry_join : giry T R := join.

End giry_join.
Arguments giry_join {d T R}.

Section measurable_giry_join.
Context {d} {T : measurableType d} {R : realType}.

Lemma measurable_giry_join : measurable_fun [set: giry (giry T R) R] giry_join.
Proof.
apply: measurable_giry_codensity => //= B mB.
by apply: measurable_giry_int => //; exact: measurable_giry_ev.
Qed.

Import HBNNSimple.
Import MeasurableR.

Lemma sintegral_giry_join (M : giry (giry T R) R) (h : {nnsfun T >-> R}) :
  sintegral (giry_join M) h = \int[M]_mu sintegral mu h.
Proof.
under eq_integral do rewrite sintegralE.
rewrite ge0_integral_fsum//.
- by move=> r; apply: measurable_funeM; exact: measurable_giry_ev.
- by move=> n x _; exact: nnsfun_mulemu_ge0.
rewrite sintegralE /=; apply: eq_fsbigr => // r rh.
rewrite integralZl//.
have := finite_measure_integrable_cst M 1 measurableT.
apply: le_integrable => //; first exact: measurable_giry_ev.
move=> mu _ /=.
rewrite normr1 (le_trans _ (@sprobability_setT _ _ _ mu))// gee0_abs//.
by rewrite le_measure// ?inE.
Qed.

Lemma giry_int_join (M : giry (giry T R) R) (h : T -> \bar R) :
    measurable_fun [set: T] h -> (forall x, 0 <= h x) ->
  giry_int (giry_join M) h = giry_int M (giry_int ^~ h).
Proof.
move=> mh h0.
pose g := nnsfun_approx measurableT mh.
pose gE := fun n => EFin \o g n.
have mgE n : measurable_fun [set: T] (gE n) by exact/measurable_EFinP.
have gE_ge0 n x : 0 <= gE n x by rewrite lee_fin.
have nd_gE x : {homo gE ^~ x : n m / (n <= m)%O >-> n <= m}.
  by move=> *; exact/lefP/nd_nnsfun_approx.
rewrite /giry_int.
transitivity (limn (fun n => \int[giry_join M]_x gE n x)).
  rewrite -monotone_convergence//; apply: eq_integral => t _.
  by apply/esym/cvg_lim => //; exact: cvg_nnsfun_approx.
transitivity (limn (fun n => \int[M]_mu \int[mu]_x gE n x)).
  apply: congr_lim; apply/funext => n.
  rewrite integralT_nnsfun sintegral_giry_join; apply: eq_integral => x _.
  by rewrite integralT_nnsfun.
rewrite -[LHS]monotone_convergence//.
- by move=> n; exact: measurable_giry_int.
- by move=> n x _; exact: integral_ge0.
- by move=> x _ m n mn; apply: ge0_le_integral => // t _; exact: nd_gE.
apply: eq_integral => mu _.
rewrite -monotone_convergence//.
apply: eq_integral => t _.
by apply/cvg_lim => //; exact: cvg_nnsfun_approx.
Qed.

End measurable_giry_join.

Section giry_bind.
Context d1 d2 (T1 : measurableType d1) (T2 : measurableType d2) (R : realType).
Implicit Types (mu : giry T1 R) (f : T1 -> giry T2 R).

Definition giry_bind mu f (mf : measurable_fun [set: T1] f) : giry T2 R :=
  (giry_join \o giry_map mf) mu.

Local Notation "m >>= mf" := (giry_bind m mf).

Lemma measurable_giry_bind f (mf : measurable_fun [set: T1] f) :
  measurable_fun [set: giry T1 R] (fun mu => mu >>= mf).
Proof.
apply: (@measurableT_comp _ _ _ _ _ _ _ _ (giry_map mf)) => //=.
  exact: measurable_giry_join.
exact: measurable_giry_map.
Qed.

Lemma giry_int_bind mu f (mf : measurable_fun [set: T1] f) (h : T2 -> \bar R) :
    measurable_fun [set: T2] h -> (forall x, 0 <= h x)%E ->
  giry_int (mu >>= mf) h = giry_int mu (fun x => giry_int (f x) h).
Proof.
move=> mh h0; rewrite /giry_bind /= giry_int_join// giry_int_map//.
  exact: measurable_giry_int.
by move=> ?; exact: integral_ge0.
Qed.

End giry_bind.

Section giry_monad.
Context d1 d2 d3 (T1 : measurableType d1) (T2 : measurableType d2)
  (T3 : measurableType d3) (R : realType).

Lemma giry_joinA (x : giry (giry (giry T1 R) R) R) :
  (giry_join \o giry_map measurable_giry_join) x ≡μ
  (giry_join \o giry_join) x.
Proof.
move=> A mA/=.
rewrite giry_int_map//; first exact: measurable_giry_ev.
by rewrite giry_int_join//; exact: measurable_giry_ev.
Qed.

Lemma giry_join_id1 (x : giry T1 R) :
  (giry_join \o giry_map measurable_giry_ret) x ≡μ
  (giry_join \o giry_ret) x.
Proof.
move=> A mA/=.
rewrite giry_int_map//; first exact: measurable_giry_ev.
rewrite giry_int_ret//; first exact: measurable_giry_ev.
by rewrite /giry_int /giry_ev /giry_ret/= /dirac integral_indic// setIT.
Qed.

Lemma giry_join_id2 (x : giry (giry T1 R) R) (f : T1 -> T2)
    (mf : measurable_fun [set: T1] f) :
  (giry_join \o giry_map (measurable_giry_map mf)) x ≡μ
  (giry_map mf \o giry_join) x.
Proof.
by move=> X mS /=; rewrite giry_int_map//; exact: measurable_giry_ev.
Qed.

End giry_monad.

(* https://en.wikipedia.org/wiki/Giry_monad#Product_distributions  *)
Definition giry_prod {d1} {d2} {T1 : measurableType d1} {T2 : measurableType d2}
    {R : realType} (m : giry T1 R * giry T2 R) : giry (T1 * T2)%type R :=
  @product_subprobability _ _ T1 T2 R m.

Section measurable_giry_prod.
Context {d1} {d2} {T1 : measurableType d1} {T2 : measurableType d2}
  {R : realType}.

(* See: Tobias Fritz. A synthetic approach to Markov kernels, conditional
   independence and theorems on sufficient statistics.
   https://arxiv.org/abs/1908.07021 *)
Lemma measurable_giry_prod :
  measurable_fun [set: giry T1 R * giry T2 R] giry_prod.
Proof.
apply: measurable_giry_codensity => //=.
rewrite prod_measurable_rectangle.
apply: dynkin_induction => /=.
- by rewrite prod_measurable_rectangle.
- move=> _ _ [A1 mA1 [A2 mA2 <-]] [B1 mB1 [B2 mB2 <-]].
  exists (A1 `&` B1); first exact: measurableI.
  exists (A2 `&` B2); first exact: measurableI.
  by rewrite setXI.
- apply: (eq_measurable_fun (fun x : giry T1 R * giry T2 R =>
      x.1 [set: T1] * x.2 [set: T2])).
    by move=> x _; rewrite -setXTT product_measure1E.
  by apply: emeasurable_funM => /=;
    apply: (@measurableT_comp _ _ _ _ _ _ (giry_ev ^~ _)) => //;
    exact: measurable_giry_ev.
- move=> _ [A mA [B mB <-]].
  apply: (eq_measurable_fun (fun x : giry T1 R * giry T2 R => x.1 A * x.2 B)).
    by move=> x _; rewrite product_measure1E.
  by apply: emeasurable_funM;
    apply: (@measurableT_comp _ _ _ _ _ _ (giry_ev ^~ _)) => //;
    exact: measurable_giry_ev.
- move=> S mS HS.
  apply: (eq_measurable_fun (fun x : giry T1 R * giry T2 R =>
      x.1 [set: T1] * x.2 [set: T2] - (x.1 \x x.2) S)).
    move=> /= x _; rewrite product_subprobability_setC//.
    by rewrite -setXTT product_measure1E.
  apply emeasurable_funB => //=.
  by apply: emeasurable_funM => //=;
    apply: (@measurableT_comp _ _ _ _ _ _ (giry_ev ^~ _)) => //;
    exact: measurable_giry_ev.
- move=> F mF tF Fn.
  apply: (eq_measurable_fun (fun x : giry T1 R * giry T2 R =>
      \sum_(0 <= k <oo) (x.1 \x x.2) (F k))).
    by move=> x _; rewrite measure_semi_bigcup//; exact: bigcup_measurable.
  exact: ge0_emeasurable_sum.
Qed.

End measurable_giry_prod.

Section giry_prod_int.
Context {d1} {d2} {T1 : measurableType d1} {T2 : measurableType d2}
  {R : realType} (m1 : giry T1 R) (m2 : giry T2 R) (h : T1 * T2 -> \bar R).
Hypotheses (mh : measurable_fun [set: T1 * T2] h) (h0 : forall x, 0 <= h x).

Lemma giry_int_prod1 : giry_int (giry_prod (m1, m2)) h =
  giry_int m1 (fun x => giry_int m2 (fun y => h (x, y))).
Proof. exact: fubini_tonelli1. Qed.

Lemma giry_int_prod2 : giry_int (giry_prod (m1, m2)) h =
  giry_int m2 (fun y => giry_int m1 (fun x => h (x, y))).
Proof. exact: fubini_tonelli2. Qed.

End giry_prod_int.
