Library QBS.qbs_prob_quot

From HB Require Import structures.
From mathcomp Require Import all_boot all_algebra reals classical_sets boolp
  ereal measurable_structure measurable_function probability_measure
  lebesgue_stieltjes_measure lebesgue_integral.
From QBS Require Import quasi_borel probability_qbs.


Import GRing.Theory Num.Def Num.Theory boolp.

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

Local Open Scope classical_set_scope.
Local Open Scope ereal_scope.

Section qbs_prob_quot.
Variable R : realType.
Local Notation mR := (measurableTypeR R).

Implicit Types (X Y Z : qbsType R).

The quotient type (setoid-style). A qbs_prob_space X wraps a qbs_prob X representative. Equality is qbs_prob_equiv: two wrapped triples are equal iff they induce the same pushforward measure.
Quotient probability space wrapping a representative.
Record qbs_prob_space (X : qbsType R) := QPS {
  qps_repr : qbs_prob X ;
}.

Arguments QPS {X}.
Arguments qps_repr {X}.

Equivalence: same pushforward measure.
Definition qps_eq (X : qbsType R) (p1 p2 : qbs_prob_space X) : Prop :=
  qbs_prob_equiv X (qps_repr p1) (qps_repr p2).

Arguments qps_eq {X}.

Lemma qps_eqxx (X : qbsType R) (p : qbs_prob_space X) :
  qps_eq p p.
Proof. exact: qbs_prob_equivxx. Qed.

Lemma qps_eqC (X : qbsType R) (p1 p2 : qbs_prob_space X) :
  qps_eq p1 p2 qps_eq p2 p1.
Proof. exact: qbs_prob_equivC. Qed.

Lemma qps_eq_trans (X : qbsType R) (p1 p2 p3 : qbs_prob_space X) :
  qps_eq p1 p2 qps_eq p2 p3 qps_eq p1 p3.
Proof. exact: qbs_prob_equiv_trans. Qed.

Definition qps_of (X : qbsType R) (p : qbs_prob X) : qbs_prob_space X :=
  QPS p.

Arguments qps_of {X}.

Lifted operations. Return, bind, and integral lifted to qbs_prob_space.

Definition qps_return (X : qbsType R) (x : X) (mu : probability mR R) :
  qbs_prob_space X :=
  qps_of (qbs_return X x mu).

Arguments qps_return {X}.

Definition qps_bind (X Y : qbsType R) (p : qbs_prob_space X)
  (f : X qbs_prob Y)
  (hdiag : qbs_Mx
    (fun rqbs_prob_alpha (f (qbs_prob_alpha (qps_repr p) r)) r)) :
  qbs_prob_space Y :=
  qps_of (qbs_bind X Y (qps_repr p) f hdiag).

Arguments qps_bind {X Y}.

Definition qps_bind_strong (X Y : qbsType R) (p : qbs_prob_space X)
  (f : X qbs_prob Y)
  (hf : qbs_morphism_strong X Y f) : qbs_prob_space Y :=
  qps_of (qbs_bind_strong X Y (qps_repr p) f hf).

Arguments qps_bind_strong {X Y}.

Definition qps_integral (X : qbsType R) (p : qbs_prob_space X)
  (h : X \bar R) : \bar R :=
  qbs_integral X (qps_repr p) h.

Arguments qps_integral {X}.

Functorial map lifted to the quotient.
Definition qps_map (X Y : qbsType R) (f : X Y)
  (hf : qbs_morphism f) (p : qbs_prob_space X) :
  qbs_prob_space Y :=
  qps_of (monadP_map X Y f hf (qps_repr p)).

Arguments qps_map {X Y}.

Well-definedness: operations respect qps_eq.

Lemma qps_return_equiv (X : qbsType R) (x : X)
  (mu1 mu2 : probability mR R) :
  qps_eq (qps_return x mu1) (qps_return x mu2).
Proof. exact: qbs_return_equiv. Qed.

Lemma qps_integral_equiv (X : qbsType R) (p1 p2 : qbs_prob_space X)
  (h : X \bar R)
  (hm : qbs_measurable X h)
  (hint1 : (qbs_prob_mu (qps_repr p1)).-integrable setT
    (h \o qbs_prob_alpha (qps_repr p1)))
  (hint2 : (qbs_prob_mu (qps_repr p2)).-integrable setT
    (h \o qbs_prob_alpha (qps_repr p2))) :
  qps_eq p1 p2
  qps_integral p1 h = qps_integral p2 h.
Proof.
moveheq.
exact: (qbs_integral_equiv hm hint1 hint2 heq).
Qed.

Lemma qps_map_equiv (X Y : qbsType R) (f : X Y)
  (hf : qbs_morphism f)
  (p1 p2 : qbs_prob_space X) :
  qps_eq p1 p2 qps_eq (qps_map f hf p1) (qps_map f hf p2).
Proof.
moveheq U hU /=.
rewrite /qps_eq /= /qbs_prob_equiv /=.
have hpreimg : (p : qbs_prob X),
  (f \o qbs_prob_alpha p) @^-1` U = qbs_prob_alpha p @^-1` (f @^-1` U).
  by movep.
by rewrite !hpreimg; apply: heq; movealpha halpha; apply: hU; apply: hf.
Qed.

Monad laws on the quotient (as qps_eq equalities).

Lemma qps_bind_returnl (X Y : qbsType R) (x : X)
  (f : X qbs_prob Y)
  (hf : qbs_morphism f) :
  qps_eq
    (qps_bind (qps_return x (qbs_prob_mu (f x))) f
      (qbs_bind_alpha_random_const x f))
    (qps_of (f x)).
Proof. exact: qbs_bind_returnl. Qed.

Lemma qps_bind_returnr (X : qbsType R) (m : qbs_prob_space X)
  (mu : probability mR R) :
  qps_eq
    (qps_bind m (qbs_return X ^~ mu)
      (qbs_bind_alpha_random_return (qps_repr m) mu))
    m.
Proof. exact: qbs_bind_returnr. Qed.

Lemma qps_bindA (X Y Z : qbsType R) (m : qbs_prob_space X)
  (f : X qbs_prob Y) (g : Y qbs_prob Z)
  (hf_diag : qbs_Mx
    (fun rqbs_prob_alpha (f (qbs_prob_alpha (qps_repr m) r)) r))
  (hg_bind : (p : qbs_prob Y),
    qbs_Mx
      (fun rqbs_prob_alpha (g (qbs_prob_alpha p r)) r))
  (hfg_diag : qbs_Mx
    (fun rqbs_prob_alpha
      (g (qbs_prob_alpha (f (qbs_prob_alpha (qps_repr m) r)) r)) r)) :
  qps_eq
    (qps_bind (qps_bind m f hf_diag) g (hg_bind _))
    (qps_of (mkQBSProb
      (fun rqbs_prob_alpha
        (g (qbs_prob_alpha
          (f (qbs_prob_alpha (qps_repr m) r)) r)) r)
      (qbs_prob_mu (qps_repr m))
      hfg_diag)).
Proof. exact: qbs_bindA. Qed.

Expectation and probability of events on the quotient.

Definition qps_expect (X : qbsType R) (p : qbs_prob_space X)
  (h : X R) : \bar R :=
  qbs_expect X (qps_repr p) h.

Arguments qps_expect {X}.

Definition qps_prob_event (X : qbsType R) (p : qbs_prob_space X)
  (U : set X) : \bar R :=
  qbs_prob_event X (qps_repr p) U.

Arguments qps_prob_event {X}.

Lemma qps_prob_event_equiv (X : qbsType R) (p1 p2 : qbs_prob_space X)
  (U : set X) :
  sigma_Mx U
  qps_eq p1 p2
  qps_prob_event p1 U = qps_prob_event p2 U.
Proof.
movehU heq.
rewrite /qps_prob_event /qbs_prob_event.
exact: heq.
Qed.

The quotient forms a QBS via monadP. Since qbs_prob_space X is isomorphic to qbs_prob X as a type, and monadP X already equips qbs_prob X with a QBS structure, we transfer the QBS structure to qbs_prob_space X.

Section qbs_prob_space_instance.
Variable (X : qbsType R).

Let Mx : set (mR qbs_prob_space X) :=
  [set beta | @monadP_random_pw R X (fun rqps_repr (beta r))].

Let Mx_comp : beta (f : {mfun mR >-> mR}),
    Mx beta Mx (beta \o f).
Proof. movebeta f hbeta r /=; exact: hbeta. Qed.

Let Mx_const : x : qbs_prob_space X, Mx (fun _x).
Proof. movex r /=; exact: (qbs_prob_alpha_random (qps_repr x)). Qed.

Let Mx_glue : (P : {mfun mR >-> nat})
    (Fi : nat mR qbs_prob_space X),
    ( i, Mx (Fi i)) Mx (fun rFi (P r) r).
Proof. moveP Fi hFi r /=; exact: hFi. Qed.

HB.instance Definition _ :=
  @isQBS.Build R (qbs_prob_space X) Mx Mx_comp Mx_const Mx_glue.

Definition qbs_prob_space_qbs : qbsType R := qbs_prob_space X.

End qbs_prob_space_instance.

Integration on the quotient: linearity properties.

Lemma qps_integral_const (X : qbsType R) (p : qbs_prob_space X) (c : \bar R) :
  qps_integral p (fun _c) = \int[qbs_prob_mu (qps_repr p)]_x c.
Proof. by []. Qed.

Lemma qps_integral_return (X : qbsType R) (x : X)
  (mu : probability mR R) (h : X \bar R) :
  qps_integral (qps_return x mu) h = \int[mu]_r h x.
Proof. by []. Qed.

Lemma qps_integral_bind (X Y : qbsType R) (p : qbs_prob_space X)
  (f : X qbs_prob Y)
  (hdiag : qbs_Mx
    (fun rqbs_prob_alpha (f (qbs_prob_alpha (qps_repr p) r)) r))
  (h : Y \bar R) :
  qps_integral (qps_bind p f hdiag) h =
  \int[qbs_prob_mu (qps_repr p)]_r
    h (qbs_prob_alpha (f (qbs_prob_alpha (qps_repr p) r)) r).
Proof. by []. Qed.

Canonical representative via classical choice. Given a qbs_prob, we can pick a canonical representative from its equivalence class using constructive_indefinite_description. This does NOT give a well-defined function (different inputs in the same class may get different representatives), but it provides a representative that IS equivalent to the original.

Definition qps_pick_repr (X : qbsType R) (p : qbs_prob_space X) :
  qbs_prob_space X :=
  QPS (proj1_sig (constructive_indefinite_description
    (ex_intro (fun qqbs_prob_equiv X (qps_repr p) q)
      (qps_repr p) (qbs_prob_equivxx (qps_repr p))))).

Lemma qps_pick_repr_equiv (X : qbsType R) (p : qbs_prob_space X) :
  qps_eq p (qps_pick_repr p).
Proof.
rewrite /qps_pick_repr /qps_eq /=.
exact: (proj2_sig (constructive_indefinite_description _ )).
Qed.

End qbs_prob_quot.

Arguments qbs_prob_space {R}.
Arguments QPS {R X}.
Arguments qps_repr {R X}.
Arguments qps_eq {R X}.
Arguments qps_of {R X}.
Arguments qps_return {R X}.
Arguments qps_bind {R X Y}.
Arguments qps_bind_strong {R X Y}.
Arguments qps_integral {R X}.
Arguments qps_map {R X Y}.
Arguments qps_expect {R X}.
Arguments qps_prob_event {R X}.
Arguments qbs_prob_space_qbs {R}.
Arguments qps_pick_repr {R X}.