(* mathcomp analysis (c) 2026 Inria and AIST. License: CeCILL-C.              *)
From HB Require Import structures.
From mathcomp Require Import boot order algebra.
From mathcomp Require Import boolp classical_sets functions cardinality reals.
From mathcomp Require Import interval_inference ereal topology normedtype.
From mathcomp Require Import measurable_structure measure_function dirac_measure.

(**md**************************************************************************)
(* # Probability Measures                                                     *)
(*                                                                            *)
(* ```                                                                        *)
(*          isSubProbability == interface for functions that satisfy the      *)
(*                              property of subprobability                    *)
(*                              The HB class is SubProbability.               *)
(*        subprobability T R == subprobability measure over the               *)
(*                              measurableType T with values in \bar R with   *)
(*                              R : realType                                  *)
(*                              The HB class is SubProbability.               *)
(*  Measure_isSubProbability == interface that extends measures to            *)
(*                              subprobability measures                       *)
(*             isProbability == interface for functions that satisfy the      *)
(*                              property of probability measures              *)
(*                              The HB class is Probability.                  *)
(*           probability T R == type of probability measure over the          *)
(*                              measurableType T with values in \bar R        *)
(*                              with R : realType                             *)
(*     Measure_isProbability == interface that extends measures to            *)
(*                              probability measures                          *)
(*             mnormalize mu == normalization of a measure to a probability   *)
(* ```                                                                        *)
(*                                                                            *)
(* ```                                                                        *)
(*                  mset U r == the set of probability measures mu such that  *)
(*                              mu U < r                                      *)
(*                      pset == the sets mset U r with U measurable and       *)
(*                              r \in [0,1]                                   *)
(*              pprobability == the measurable type generated by pset         *)
(* ```                                                                        *)
(******************************************************************************)

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.Theory.

Local Open Scope classical_set_scope.
Local Open Scope ring_scope.

HB.mixin Record isSubProbability d (T : sigmaRingType d) (R : realType)
  (P : set T -> \bar R) := { sprobability_setT : (P setT <= 1)%E }.

#[short(type=subprobability)]
HB.structure Definition SubProbability d (T : measurableType d) (R : realType)
  := {mu of @FiniteMeasure d T R mu & isSubProbability d T R mu }.

HB.factory Record Measure_isSubProbability d (T : measurableType d)
    (R : realType) (P : set T -> \bar R) & isMeasure _ _ _ P :=
  { sprobability_setT : (P setT <= 1)%E }.

HB.builders Context d (T : measurableType d) (R : realType)
  P & Measure_isSubProbability d T R P.

Let finite : @Measure_isFinite d T R P.
Proof.
split; apply: lty_fin_num_fun.
by rewrite (le_lt_trans (@sprobability_setT))// ltey.
Qed.

HB.instance Definition _ := finite.

HB.instance Definition _ := @isSubProbability.Build _ _ _ P sprobability_setT.

HB.end.

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

Let mzero_setT : (@mzero d T R setT <= 1)%E.
Proof. by rewrite /mzero/=. Qed.

HB.instance Definition _ :=
  Measure_isSubProbability.Build _ _ _ (@mzero d T R) mzero_setT.

End mzero_subprobability.

HB.mixin Record isProbability d (T : measurableType d) (R : realType)
  (P : set T -> \bar R) := { probability_setT : P setT = 1%E }.

#[short(type=probability)]
HB.structure Definition Probability d (T : measurableType d) (R : realType) :=
  {P of @SubProbability d T R P & isProbability d T R P }.

Arguments probability_setT {d T R} s.

HB.instance Definition _ d (T : measurableType d) (R : realType) :=
  gen_eqMixin (probability T R).
HB.instance Definition _ d (T : measurableType d) (R : realType) :=
  gen_choiceMixin (probability T R).

Section probability_lemmas.
Local Open Scope ereal_scope.
Context d (T : measurableType d) (R : realType) (P : probability T R).

Lemma probability_le1 (A : set T) : measurable A -> P A <= 1.
Proof. by move=> mA; rewrite -(probability_setT P) ?le_measure ?in_setE. Qed.

Lemma probability_setC (A : set T) : measurable A -> P (~` A) = 1 - P A.
Proof.
move=> mA; rewrite -(probability_setT P) -(setvU A) measureU ?addeK ?setICl//.
  exact: measurableC.
by rewrite fin_num_measure.
Qed.

End probability_lemmas.

HB.factory Record Measure_isProbability d (T : measurableType d)
    (R : realType) (P : set T -> \bar R) & isMeasure _ _ _ P :=
  { probability_setT : P setT = 1%E }.

HB.builders Context d (T : measurableType d) (R : realType)
  P & Measure_isProbability d T R P.

Let subprobability : @Measure_isSubProbability d T R P.
Proof. by split; rewrite probability_setT. Qed.

HB.instance Definition _ := subprobability.

HB.instance Definition _ := @isProbability.Build _ _ _ P probability_setT.

HB.end.

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

HB.instance Definition _ x :=
  Measure_isProbability.Build _ _ _ (@dirac _ T x R) (diracT R x).

End pdirac.

Section mnormalize.
Local Open Scope ereal_scope.
Context d (T : measurableType d) (R : realType).
Variables (mu : {measure set T -> \bar R}) (P : probability T R).

Definition mnormalize :=
  let evidence := mu [set: T] in
  if (evidence == 0) || (evidence == +oo) then fun U => P U
  else fun U => mu U * (fine evidence)^-1%:E.

Let mnormalize0 : mnormalize set0 = 0.
Proof.
by rewrite /mnormalize; case: ifPn => // _; rewrite measure0 mul0e.
Qed.

Let mnormalize_ge0 U : 0 <= mnormalize U.
Proof. by rewrite /mnormalize; case: ifPn. Qed.

Let mnormalize_sigma_additive : semi_sigma_additive mnormalize.
Proof.
move=> F mF tF mUF; rewrite /mnormalize/=.
case: ifPn => [_|_]; first exact: measure_semi_sigma_additive.
rewrite [X in X @ _ --> _](_ : _ = (fun n => \sum_(0 <= i < n) mu (F i)) \*
                               cst (fine (mu setT))^-1%:E).
  by apply/funext => n; rewrite -ge0_sume_distrl.
by apply: cvgeZr => //; exact: measure_semi_sigma_additive.
Qed.

HB.instance Definition _ := isMeasure.Build _ _ _ mnormalize
  mnormalize0 mnormalize_ge0 mnormalize_sigma_additive.

Let mnormalize1 : mnormalize [set: T] = 1.
Proof.
rewrite /mnormalize; case: ifPn; first by rewrite probability_setT.
rewrite negb_or => /andP[ft0 ftoo].
have ? : mu setT \is a fin_num by rewrite ge0_fin_numE// ltey.
by rewrite -{1}(@fineK _ (mu setT))// -EFinM divff// fine_eq0.
Qed.

HB.instance Definition _ :=
  Measure_isProbability.Build _ _ _ mnormalize mnormalize1.

End mnormalize.

Lemma mnormalize_id d (T : measurableType d) (R : realType)
  (P P' : probability T R) : mnormalize P P' = P.
Proof.
apply/funext => x; rewrite /mnormalize/= probability_setT.
by rewrite onee_eq0/= invr1 mule1.
Qed.

HB.instance Definition _ d (T : pmeasurableType d) (R : realType) :=
  isPointed.Build (probability T R) (dirac point).

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

Definition mset (U : set T) (r : R) :=
  [set mu : probability T R | mu U < r%:E]%E.

Lemma lt0_mset (U : set T) (r : R) : (r < 0)%R -> mset U r = set0.
Proof.
move=> r0; apply/seteqP; split => // x/=.
by apply/negP; rewrite -leNgt (@le_trans _ _ 0%E)// lee_fin ltW.
Qed.

Lemma gt1_mset (U : set T) (r : R) :
  measurable U -> (1 < r)%R -> mset U r = [set: probability T R].
Proof.
move=> mU r1; apply/seteqP; split => // x/= _.
by rewrite /mset/= (le_lt_trans (probability_le1 _ _)).
Qed.

Definition pset : set_system (probability T R) :=
  [set mset U r | r in `[0%R,1%R] & U in measurable].

Definition pprobability : measurableType pset.-sigma :=
  g_sigma_algebraType pset.

End dist_sigma_algebra_instance.
