Theory Geometric_PMF

(*
  File:    Pi_pmf.thy
  Authors: Manuel Eberl, Max W. Haslbeck
*)
section ‹Theorems about the Geometric Distribution›
theory Geometric_PMF
  imports
    "HOL-Probability.Probability"
    Pi_pmf
    "Monad_Normalisation.Monad_Normalisation"
begin

lemma nn_integral_geometric_pmf:
  assumes "p  {0<..1}"
  shows   "nn_integral (geometric_pmf p) real = (1 - p) / p"
 using assms expectation_geometric_pmf integrable_real_geometric_pmf
  by (subst nn_integral_eq_integral) auto

lemma geometric_pmf_prob_atMost:
  assumes "p  {0<..1}"
  shows  "measure_pmf.prob (geometric_pmf p) {..n} = (1 - (1 - p)^(n + 1))"
proof -
  have "(xn. (1 - p) ^ x * p) = 1 - (1 - p) * (1 - p) ^ n"
    by (induction n) (auto simp add: algebra_simps)
  then show ?thesis
  using assms by (auto simp add: measure_pmf_conv_infsetsum)
qed

lemma geometric_pmf_prob_lessThan:
  assumes "p  {0<..1}"
  shows  "measure_pmf.prob (geometric_pmf p) {..<n} = 1 - (1 - p) ^ n"
proof -
  have "(x<n. (1 - p) ^ x * p) = 1 - (1 - p) ^ n"
    by (induction n) (auto simp add: algebra_simps)
  then show ?thesis
  using assms by (auto simp add: measure_pmf_conv_infsetsum)
qed

lemma geometric_pmf_prob_greaterThan:
  assumes "p  {0<..1}"
  shows  "measure_pmf.prob (geometric_pmf p) {n<..} = (1 - p)^(n + 1)"
proof -
  have "(UNIV - {..n}) = {n<..}"
    by auto
  moreover have "measure_pmf.prob (geometric_pmf p) (UNIV - {..n}) = (1 - p) ^ (n + 1)"
    using assms by (subst measure_pmf.finite_measure_Diff)
                   (auto simp add: geometric_pmf_prob_atMost)
  ultimately show ?thesis
    by auto
qed

lemma geometric_pmf_prob_atLeast:
  assumes "p  {0<..1}"
  shows  "measure_pmf.prob (geometric_pmf p) {n..} = (1 - p)^n"
proof -
  have "(UNIV - {..<n}) = {n..}"
    by auto
  moreover have "measure_pmf.prob (geometric_pmf p) (UNIV - {..<n}) = (1 - p) ^ n"
    using assms by (subst measure_pmf.finite_measure_Diff)
                   (auto simp add: geometric_pmf_prob_lessThan)
  ultimately show ?thesis
    by auto
qed

lemma bernoulli_pmf_of_set':
  assumes "finite A"
  shows "map_pmf (λb. {x  A. ¬ b x}) (Pi_pmf A P (λ_. bernoulli_pmf (1/2))) = pmf_of_set (Pow A)"
proof -
  have *: "Pi_pmf A P (λ_. pmf_of_set (UNIV :: bool set)) = pmf_of_set (PiE_dflt A P (λ_. UNIV :: bool set))"
    using assms by (intro Pi_pmf_of_set) auto
  have "map_pmf (λb. {x  A. ¬ b x}) (Pi_pmf A P (λ_. bernoulli_pmf (1 / 2))) =
        map_pmf (λb. {x  A. ¬ b x}) (Pi_pmf A P (λ_. pmf_of_set UNIV))"
    using bernoulli_pmf_half_conv_pmf_of_set by auto
  also have " = map_pmf (λb. {x  A. ¬ b x}) (pmf_of_set (PiE_dflt A P (λ_. UNIV)))"
    using assms by (subst Pi_pmf_of_set) (auto)
  also have " = pmf_of_set (Pow A)"
  proof -
    have "bij_betw (λb. {x  A. ¬ b x}) (PiE_dflt A P (λ_. UNIV)) (Pow A)"
      by (rule bij_betwI[of _ _ _ "λB b. if b  A then ¬ (b  B) else P"]) (auto simp add: PiE_dflt_def)
    then show ?thesis
      using assms by (intro map_pmf_of_set_bij_betw) auto
  qed
  finally show ?thesis
    by simp
qed

lemma Pi_pmf_pmf_of_set_Suc:
  assumes "finite A"
  shows "Pi_pmf A 0 (λ_. geometric_pmf (1/2)) =
       do {
         B  pmf_of_set (Pow A);
         Pi_pmf B 0 (λ_. map_pmf Suc (geometric_pmf (1/2))) }"
proof -
  have "Pi_pmf A 0 (λ_. geometric_pmf (1/2)) =
        Pi_pmf A 0 (λ_. bernoulli_pmf (1/2) 
        (λb. if b then return_pmf 0 else map_pmf Suc (geometric_pmf (1/2))))"
    using assms by (subst geometric_bind_pmf_unfold) auto
  also have " =
             Pi_pmf A False (λ_. bernoulli_pmf (1/2)) 
             (λb. Pi_pmf A 0 (λx. if b x then return_pmf 0 else map_pmf Suc (geometric_pmf (1/2))))"
    using assms by (subst Pi_pmf_bind[of _ _ _ _ False]) auto
  also have " =
             do {b  Pi_pmf A False (λ_. bernoulli_pmf (1/2));
                 Pi_pmf {x  A. ~b x} 0 (λx. map_pmf Suc (geometric_pmf (1/2)))}"
    using assms by (subst Pi_pmf_if_set') auto
  also have " =
             do {B  map_pmf (λb. {x  A. ¬ b x}) (Pi_pmf A False (λ_. bernoulli_pmf (1/2)));
                 Pi_pmf B 0 (λx. map_pmf Suc (geometric_pmf (1/2)))}"
    unfolding map_pmf_def apply(subst bind_assoc_pmf) apply(subst bind_return_pmf) by auto
  also have " = pmf_of_set (Pow A)  (λB. Pi_pmf B 0 (λx. map_pmf Suc (geometric_pmf (1 / 2))))"
    unfolding assms using assms by (subst bernoulli_pmf_of_set') auto
  finally show ?thesis
    by simp
qed

lemma Pi_pmf_pmf_of_set_Suc':
  assumes "finite A"
  shows "Pi_pmf A 0 (λ_. geometric_pmf (1/2)) =
       do {
         B  pmf_of_set (Pow A);
         Pi_pmf B 0 (λ_. map_pmf Suc (geometric_pmf (1/2))) }"
proof -
  have "Pi_pmf A 0 (λ_. geometric_pmf (1/2)) =
        Pi_pmf A 0 (λ_. bernoulli_pmf (1/2) 
        (λb. if b then return_pmf 0 else map_pmf Suc (geometric_pmf (1/2))))"
    using assms by (subst geometric_bind_pmf_unfold) auto
  also have " =
             Pi_pmf A False (λ_. bernoulli_pmf (1/2)) 
             (λb. Pi_pmf A 0 (λx. if b x then return_pmf 0 else map_pmf Suc (geometric_pmf (1/2))))"
    using assms by (subst Pi_pmf_bind[of _ _ _ _ False]) auto
  also have " =
             do {b  Pi_pmf A False (λ_. bernoulli_pmf (1/2));
                 Pi_pmf {x  A. ~b x} 0 (λx. map_pmf Suc (geometric_pmf (1/2)))}"
    using assms by (subst Pi_pmf_if_set') auto
  also have " =
             do {B  map_pmf (λb. {x  A. ¬ b x}) (Pi_pmf A False (λ_. bernoulli_pmf (1/2)));
                 Pi_pmf B 0 (λx. map_pmf Suc (geometric_pmf (1/2)))}"
    unfolding map_pmf_def by (auto simp add: bind_assoc_pmf bind_return_pmf)
  also have " = pmf_of_set (Pow A)  (λB. Pi_pmf B 0 (λx. map_pmf Suc (geometric_pmf (1 / 2))))"
    unfolding assms using assms by (subst bernoulli_pmf_of_set') auto
  finally show ?thesis
    by simp
qed

lemma binomial_pmf_altdef':
  fixes A :: "'a set"
  assumes "finite A" and "card A = n" and p: "p  {0..1}"
  shows   "binomial_pmf n p =
             map_pmf (λf. card {xA. f x}) (Pi_pmf A dflt (λ_. bernoulli_pmf p))" (is "?lhs = ?rhs")
proof -
  from assms have "?lhs = binomial_pmf (card A) p"
    by simp
  also have " = ?rhs"
  using assms(1)
  proof (induction rule: finite_induct)
    case empty
    with p show ?case by (simp add: binomial_pmf_0)
  next
    case (insert x A)
    from insert.hyps have "card (insert x A) = Suc (card A)"
      by simp
    also have "binomial_pmf  p = do {
                                     b  bernoulli_pmf p;
                                     f  Pi_pmf A dflt (λ_. bernoulli_pmf p);
                                     return_pmf ((if b then 1 else 0) + card {y  A. f y})
                                   }"
      using p by (simp add: binomial_pmf_Suc insert.IH bind_map_pmf)
    also have " = do {
                      b  bernoulli_pmf p;
                      f  Pi_pmf A dflt (λ_. bernoulli_pmf p);
                      return_pmf (card {y  insert x A. (f(x := b)) y})
                    }"
    proof (intro bind_pmf_cong refl, goal_cases)
      case (1 b f)
      have "(if b then 1 else 0) + card {yA. f y} = card ((if b then {x} else {})  {yA. f y})"
        using insert.hyps by auto
      also have "(if b then {x} else {})  {yA. f y} = {yinsert x A. (f(x := b)) y}"
        using insert.hyps by auto
      finally show ?case by simp
    qed
    also have " = map_pmf (λf. card {yinsert x A. f y})
                      (Pi_pmf (insert x A) dflt (λ_. bernoulli_pmf p))"
      using insert.hyps by (subst Pi_pmf_insert) (simp_all add: pair_pmf_def map_bind_pmf)
    finally show ?case .
  qed
  finally show ?thesis .
qed

lemma bernoulli_pmf_Not:
  assumes "p  {0..1}"
  shows "bernoulli_pmf p = map_pmf Not (bernoulli_pmf (1 - p))"
proof -
  have *: "(Not -` {True}) = {False}" "(Not -` {False}) = {True}"
    by blast+
  have "pmf (bernoulli_pmf p) b = pmf (map_pmf Not (bernoulli_pmf (1 - p))) b" for b
    using assms by (cases b) (auto simp add: pmf_map * measure_pmf_single)
  then show ?thesis
    by (rule pmf_eqI)
qed

lemma binomial_pmf_altdef'':
  assumes p: "p  {0..1}"
  shows   "binomial_pmf n p =
           map_pmf (λf. card {x. x < n  f x}) (Pi_pmf {..<n} dflt (λ_. bernoulli_pmf p))"
  using assms by (subst binomial_pmf_altdef'[of "{..<n}"]) (auto)

context includes monad_normalisation
begin

lemma Pi_pmf_geometric_filter:
  assumes "finite A" "p  {0<..1}"
  shows "Pi_pmf A 0 (λ_. geometric_pmf p) =
       do {
         fb  Pi_pmf A dflt (λ_. bernoulli_pmf p);
         Pi_pmf {x  A. ¬ fb x} 0 (λ_. map_pmf Suc (geometric_pmf p)) }"
proof -
  have "Pi_pmf A 0 (λ_. geometric_pmf p) =
        Pi_pmf A 0 (λ_. bernoulli_pmf p 
                   (λb. if b then return_pmf 0 else map_pmf Suc (geometric_pmf p)))"
    using assms by (subst geometric_bind_pmf_unfold) auto
  also have " =
             Pi_pmf A dflt (λ_. bernoulli_pmf p) 
             (λb. Pi_pmf A 0 (λx. if b x then return_pmf 0 else map_pmf Suc (geometric_pmf p)))"
    using assms by (subst Pi_pmf_bind[of _ _ _ _ dflt]) auto
  also have " =
             do {b  Pi_pmf A dflt (λ_. bernoulli_pmf p);
                 Pi_pmf {x  A. ¬ b x} 0 (λx. map_pmf Suc (geometric_pmf p))}"
    using assms by (subst Pi_pmf_if_set') (auto)
  finally show ?thesis
    by simp
qed

lemma Pi_pmf_geometric_filter':
  assumes "finite A" "p  {0<..1}"
  shows "Pi_pmf A 0 (λ_. geometric_pmf p) =
       do {
         fb  Pi_pmf A dflt (λ_. bernoulli_pmf (1 - p));
         Pi_pmf {x  A. fb x} 0 (λ_. map_pmf Suc (geometric_pmf p)) }"
  using assms by (auto simp add: Pi_pmf_geometric_filter[of _ _ "¬ dflt"] bernoulli_pmf_Not[of p]
      Pi_pmf_map[of _ _ dflt] map_pmf_def[of "((∘) Not)"])

end

end