\<^marker>\<open>creator "Peter Lammich"\<close>
\<^marker>\<open>contributor "Maximilian P. L. Haslbeck"\<close>
section \<open>Introsort and its first Phase\<close>
theory Sorting_Quicksort_Scheme
imports Sorting_Setup Sorting_Partially_Sorted
begin

paragraph \<open>Summary\<close>
text \<open>This theory introduces the abstract quicksort-like first phase of the introsort algorithm.
  Also the abstract algorithm for introsort is introduced and proved correct.
    \<close>

paragraph \<open>Main Theorems/Definitions\<close>
text \<open>
\<^item> introsort_aux1: gives the abstract algorithm for the first phase
\<^item> introsort_aux_cost: specifies the cost of that algorithm
\<^item> introsort_aux1_correct': proves correctness
\<^item> introsort3: abstract introsort algorithm
\<^item> introsort3_correct: proves its correctness
\<close>

subsection \<open>introsort_aux1\<close>

  abbreviation "is_threshold \<equiv> 16::nat"

  context weak_ordering begin

  definition "partition1_spec xs \<equiv> doN {
    ASSERT (length xs \<ge> 4);
    SPEC (\<lambda>(xs1,xs2). mset xs = mset xs1 + mset xs2 \<and> xs1\<noteq>[] \<and> xs2\<noteq>[] \<and> slice_LT (\<^bold>\<le>) xs1 xs2)
         (\<lambda>_. cost ''partition_c'' (enat (length xs)))
  }"

  definition introsort_aux1 :: "(nat \<Rightarrow> nat) \<Rightarrow> 'a list \<Rightarrow> nat \<Rightarrow> ('a list,_) nrest" where
    "introsort_aux1 \<mu> xs d \<equiv> RECT' (\<lambda>introsort_aux1 (xs,d). doN {
      ASSERT (xs \<noteq> []);
      lxs \<leftarrow> SPECT [length xs \<mapsto> cost ''list_length'' 1];
      if\<^sub>N SPECc2 ''lt'' (<) is_threshold (length xs) then doN {
        if\<^sub>N SPECc2 ''eq'' (=) d 0 then
          mop_call (SPEC (sort_spec (\<^bold><) xs) (\<lambda>_. cost ''sort_c'' (enat (\<mu> (length xs)))))
        else doN {
          (xs1,xs2)\<leftarrow>partition1_spec xs;
          d' \<leftarrow> SPECc2 ''sub'' (-) d 1;
          xs1 \<leftarrow> introsort_aux1 (xs1,d');
          xs2 \<leftarrow> introsort_aux1 (xs2,d');
          SPECc2 ''list_append'' (@) xs1 xs2
        }
      }
      else
        RETURN xs
    }) (xs,d)"

    lemma slice_strict_LT_imp_LE: "slice_LT (\<^bold><) xs ys \<Longrightarrow> slice_LT (le_by_lt (\<^bold><)) xs ys"
      apply (erule slice_LT_mono)
      by (meson le_by_lt_def wo_less_asym)

definition introsort_aux_cost :: "_ \<Rightarrow> nat * nat \<Rightarrow> (char list, enat) acost"
  where "introsort_aux_cost \<mu> = (\<lambda>(lxs, d). lift_acost (
        ((d+1)*lxs) *m (cost ''if'' 2 + cost ''eq'' 1 + cost ''lt'' 1 + cost ''call'' 1
                    + cost ''list_length'' 1 + cost ''sub'' 1 + cost ''list_append'' 1)
            + cost ''sort_c'' (\<mu> (lxs)) + cost ''partition_c'' (d*(lxs))
         )
        )"

lemma introsort_aux1_aux5: "b>0 \<Longrightarrow> (b - Suc 0) * A + (b - Suc 0) * B + (B + A) = b*(A+B)"
proof -
  assume "b>0"
  then obtain b' where b: "b=b'+1"
    by (metis Groups.add_ac(2) Suc_eq_plus1_left Suc_pred)
  then have "(b - Suc 0) * A + (b - Suc 0) * B + (B + A)
      = (b') * A + (b') * B + (B + A)" by simp
  also have "\<dots> = (b'+1) * (B+A)" by (auto simp: ring_distribs)
  also have "\<dots> = b * (B+A)" using b by simp
  finally show ?thesis by simp
qed

lemma lengths_sums_if_msets_do: "mset a = mset b + mset c \<Longrightarrow> length a = length b + length c"
    by (metis add.commute length_append less_or_eq_imp_le mset_append mset_eq_length)

lemma
  introsort_aux1_aux:
  "b>0 \<Longrightarrow> lxs+lys>(16::nat) \<Longrightarrow> Suc ((b - Suc 0) * lxs + (b - Suc 0) * lys) \<le> b * (lxs + lys)"
  by (smt One_nat_def Suc_leI add.commute add_le_cancel_left less_imp_le_nat
          introsort_aux1_aux5 neq0_conv not_numeral_le_zero plus_1_eq_Suc)

lemma
  introsort_aux1_aux2:
  "b>0 \<Longrightarrow> lxs+lys>(16::nat) \<Longrightarrow> Suc ((b - Suc 0) * lys + (b - Suc 0) * lxs) \<le> b * (lxs + lys)"
  by (smt One_nat_def Suc_leI add.commute add_le_cancel_left less_imp_le_nat
          introsort_aux1_aux5 neq0_conv not_numeral_le_zero plus_1_eq_Suc)

lemma
  introsort_aux1_aux4:
  assumes "b>0" "lxs+lys>(16::nat)"
  shows "Suc(Suc ((b - Suc 0) * lxs * 2 + (b - Suc 0) * lys * 2)) \<le> b * (lxs + lys) * 2"
proof -
  have "Suc(Suc ((b - Suc 0) * lxs * 2 + (b - Suc 0) * lys * 2))
    = (Suc ((b - Suc 0) * lxs + (b - Suc 0) * lys)) + (Suc ((b - Suc 0) * lxs + (b - Suc 0) * lys))"
    by auto
  also have "\<dots> \<le> (b * (lxs + lys)) + (b * (lxs + lys))"
    apply(rule add_mono)
    subgoal apply(rule introsort_aux1_aux) using assms by auto
    subgoal apply(rule introsort_aux1_aux) using assms by auto
    done
  also have "\<dots> = b * (lxs + lys) * 2"
    by auto
  finally show ?thesis .
qed

subsubsection \<open>Correctness Theorem\<close>

definition "introsort_aux_cost' \<mu>  = (\<lambda>(xs,d). introsort_aux_cost \<mu> (length xs,d) )"

  lemma introsort_aux1_correct':
    assumes tf_suplinear: "\<And>a b. \<mu> a + \<mu> b \<le> \<mu> (a+b)"
      and "xs \<noteq> []"
    shows
   "introsort_aux1 \<mu> xs d
     \<le> SPEC (\<lambda>xs'. mset xs' = mset xs \<and> part_sorted_wrt (le_by_lt (\<^bold><)) is_threshold xs')
               (\<lambda>_. introsort_aux_cost' \<mu> (xs, d))"

    unfolding introsort_aux1_def partition1_spec_def sort_spec_def

    apply (rule RECT'_rule_arb[where V="measure (\<lambda>(xs,d). d+1)"
                                  and pre="\<lambda>xss (xs',d). length xs' > 0 \<and> xss=xs'"])
       apply refine_mono
      apply (all \<open>(auto intro: sorted_wrt_imp_part_sorted part_sorted_wrt_init; fail)?\<close>)
    subgoal using  assms by simp

    unfolding SPEC_REST_emb'_conv SPECc2_alt
    subgoal premises prems for f xss x
      apply(rule gwp_specifies_I)
      apply(refine_vcg \<open>-\<close> rules: prems(1)[THEN gwp_specifies_rev_I, THEN gwp_conseq_0])
      thm prems(1)[THEN gwp_specifies_rev_I, THEN gwp_conseq_0]
      using  prems(2)
      subgoal unfolding emb'_def apply auto
        subgoal  by (simp add: sorted_wrt_imp_part_sorted)
        subgoal unfolding introsort_aux_cost_def introsort_aux_cost'_def
          apply (simp add: norm_cost lift_acost_propagate lift_acost_cost)
          apply sc_solve by (auto simp: one_enat_def)
        done
             apply simp apply simp
           apply simp apply simp
      subgoal
        apply (auto simp add: emb'_def handy_if_lemma)

        subgoal  using prems(2) by simp
        subgoal
          apply (rule part_sorted_concatI; assumption?)
          apply (subst slice_LT_mset_eq1, assumption)
          apply (subst slice_LT_mset_eq2, assumption)
          using le_by_lt by blast
        subgoal premises p
          unfolding introsort_aux_cost_def introsort_aux_cost'_def
          using p(3,8)
          apply (simp add: norm_cost)
          apply sc_solve_debug apply safe apply(all \<open>((auto simp: sc_solve_debug_def),fail)?\<close>)
          subgoal (* append *)
            unfolding sc_solve_debug_def
            apply(drule lengths_sums_if_msets_do)
            using p(2)
            apply (simp add: one_enat_def  )
            apply(rule introsort_aux1_aux) by auto
          subgoal (* partition *)
            unfolding sc_solve_debug_def
            apply(drule lengths_sums_if_msets_do)
            apply auto apply(simp only: introsort_aux1_aux5) apply simp done
          subgoal (* sub *)
            unfolding sc_solve_debug_def
            apply(drule lengths_sums_if_msets_do)
            using p(2)
            apply (simp add: one_enat_def  )
            apply(rule introsort_aux1_aux) by auto
          subgoal (* ''sort_c'' *) unfolding sc_solve_debug_def
            apply (simp add: one_enat_def)
            apply(drule lengths_sums_if_msets_do) apply simp apply(subst add.commute)
            by(rule tf_suplinear)
          subgoal  (* if *)
            unfolding sc_solve_debug_def
            apply(drule lengths_sums_if_msets_do)
            using p(2)
            apply (simp add: one_enat_def)
            apply(rule introsort_aux1_aux4) by auto
          subgoal  (* list_length *)
            unfolding sc_solve_debug_def
            apply(drule lengths_sums_if_msets_do)
            using p(2)
            apply (simp add: one_enat_def  )
            apply(rule introsort_aux1_aux2) by auto
          subgoal (* eq *)
            unfolding sc_solve_debug_def
            apply(drule lengths_sums_if_msets_do)
            using p(2)
            apply (simp add: one_enat_def  )
            apply(rule introsort_aux1_aux) by auto
          subgoal (* call *)
            unfolding sc_solve_debug_def
            apply(drule lengths_sums_if_msets_do)
            using p(2)
            apply (simp add: one_enat_def  )
            apply(rule introsort_aux1_aux2) by auto
          subgoal  (* lt *)
            unfolding sc_solve_debug_def
            apply(drule lengths_sums_if_msets_do)
            using p(2)
            apply (simp add: one_enat_def  )
            apply(rule introsort_aux1_aux) by auto
          done
        done
      subgoal by auto
      subgoal
        using prems(2) apply (auto simp add: emb'_def handy_if_lemma)
        subgoal by (simp add: part_sorted_wrt_init)
        subgoal
          unfolding introsort_aux_cost_def introsort_aux_cost'_def
          apply (simp add: norm_cost)
          apply sc_solve apply (auto simp: one_enat_def)
          subgoal
            by (simp add: Suc_leI)
          subgoal
            by (simp add: Suc_leI)
          done
        done
      subgoal using prems by simp
      done
    done

  lemma introsort_aux1_correct:
    assumes tf_suplinear: "\<And>a b. \<mu> a + \<mu> b \<le> \<mu> (a+b)"
      and "xs \<noteq> []" "lxs = length xs"
    shows
      "introsort_aux1 \<mu> xs d
         \<le> SPEC (\<lambda>xs'. mset xs' = mset xs \<and> part_sorted_wrt (le_by_lt (\<^bold><)) is_threshold xs')
                 (\<lambda>_. introsort_aux_cost \<mu> (lxs, d))"
    unfolding assms(3)
    by (rule  introsort_aux1_correct'[where \<mu>=\<mu>, OF assms(1,2), unfolded introsort_aux_cost'_def, simplified])

subsection \<open>introsort_aux2\<close>

    definition "partition2_spec xs \<equiv> doN {
      ASSERT (length xs \<ge> 4);
      SPEC (\<lambda>(xs',i). mset xs' = mset xs \<and> 0<i \<and> i<length xs \<and> slice_LT (\<^bold>\<le>) (take i xs') (drop i xs'))
             (\<lambda>_. cost ''partition_c'' (enat (length xs)))
    }"

    lemma partition2_spec_refine:
      assumes "(xs,xs')\<in>Id"
      shows "partition2_spec xs
             \<le>  \<Down>(br (\<lambda>(xs,i). (take i xs, drop i xs)) (\<lambda>(xs,i). 0<i \<and> i<length xs))
                     (timerefine TId (partition1_spec xs'))"
      using assms
      unfolding partition1_spec_def partition2_spec_def
      apply (intro refine0 SPEC_both_refine)
       apply (auto dest: mset_eq_length simp: in_br_conv simp flip: mset_append)
      subgoal
        using mset_eq_length by fastforce
      subgoal
        using mset_eq_length by force
      done

  definition introsort_aux2 :: "(nat \<Rightarrow> nat) \<Rightarrow> 'a list \<Rightarrow> nat \<Rightarrow> ('a list,_) nrest" where
    "introsort_aux2 \<mu> xs d \<equiv> RECT' (\<lambda>introsort_aux (xs,d). doN {
      lxs \<leftarrow> SPECT [length xs \<mapsto> cost ''list_length'' 1];
      if\<^sub>N SPECc2 ''lt'' (<) is_threshold lxs then
        if\<^sub>N SPECc2 ''eq'' (=) d 0 then
          mop_call (SPEC (sort_spec (\<^bold><) xs) (\<lambda>_. cost ''sort_c'' (enat (\<mu> (length xs)))))
        else doN {
          (xs,m)\<leftarrow>partition2_spec xs;
          ASSERT (m\<le>length xs);
          d' \<leftarrow> SPECc2 ''sub'' (-) d 1;
          xs1 \<leftarrow> introsort_aux (take m xs,d');
          xs2 \<leftarrow> introsort_aux (drop m xs,d');
          SPECc2 ''list_append'' (@) xs1 xs2
        }
      else
        RETURN xs
    }) (xs,d)"

subsubsection \<open>Refinement Theorem\<close>

    lemma introsort_aux2_refine: "introsort_aux2 \<mu> xs d \<le>\<Down>Id (timerefine TId (introsort_aux1 \<mu> xs d))"
      unfolding introsort_aux2_def introsort_aux1_def
      apply (refine_rcg mop_call_refine partition2_spec_refine SPEC_both_refine MIf_refine SPECc2_refine RECT'_refine_t bindT_refine_conc_time_my_inres)
      apply refine_dref_type
      apply refine_mono
      apply (auto simp: in_br_conv cost_n_leq_TId_n)
      done

subsection \<open>introsort_aux3\<close>

    definition "partition3_spec xs l h \<equiv> doN {
      ASSERT (h-l\<ge>4 \<and> h\<le>length xs);
      SPEC (\<lambda>(xs',i). slice_eq_mset l h xs' xs \<and> l<i \<and> i<h \<and> slice_LT (\<^bold>\<le>) (slice l i xs') (slice i h xs')) (\<lambda>_. cost ''partition_c'' (enat (h-l)))
    }"

abbreviation "TR_i_aux \<equiv> (TId(''list_append'':=0,''list_length'':=cost ''sub'' 1))"

    lemma partition3_spec_refine: "(xsi,xs) \<in> slice_rel xs\<^sub>0 l h \<Longrightarrow> partition3_spec xsi l h  \<le>\<Down>(slice_rel xs\<^sub>0 l h \<times>\<^sub>r idx_shift_rel l) (timerefine TR_i_aux (partition2_spec xs))"
      unfolding partition3_spec_def partition2_spec_def
      apply (intro refine0 SPEC_both_refine)
       apply (auto simp: slice_rel_def in_br_conv) [2]
      subgoal for xs'i ii
        apply (rule exI[where x="slice l h xs'i"])
        apply (rule conjI)
        subgoal by (auto simp: slice_eq_mset_def)
        apply (simp add: idx_shift_rel_alt)
        by (auto simp: slice_eq_mset_def take_slice drop_slice intro!: cost_mono)
      done

    lemma partition3_spec_refine': "\<lbrakk>(xsi,xs) \<in> slicep_rel l h; xsi'=xsi; l'=l; h'=h\<rbrakk>
      \<Longrightarrow> partition3_spec xsi l h  \<le>\<Down>(slice_rel xsi' l' h' \<times>\<^sub>r idx_shift_rel l') (timerefine TR_i_aux (partition2_spec xs))"
      unfolding partition3_spec_def partition2_spec_def
      apply (intro refine0 SPEC_both_refine )
      apply (auto simp: slicep_rel_def in_br_conv) [2]
      subgoal for xs'i ii
        apply (rule exI[where x="slice l h xs'i"])
        apply (rule conjI)
        subgoal by (auto simp: slice_eq_mset_def)
        apply (simp add: idx_shift_rel_alt)
        apply (auto simp: slice_eq_mset_def take_slice drop_slice intro!: cost_mono)
        by (smt less_or_eq_imp_le less_trans slice_eq_mset_alt slice_eq_mset_def slice_in_slice_rel slice_rel_alt)
      done

    definition introsort_aux3 :: "(nat \<Rightarrow> nat) \<Rightarrow> 'a list \<Rightarrow> nat \<Rightarrow> nat \<Rightarrow> nat \<Rightarrow> ('a list,_) nrest" where "introsort_aux3 \<mu> xs l h d
    \<equiv>  RECT' (\<lambda>introsort_aux (xs,l,h,d). doN {
        ASSERT (l\<le>h);
        lxs \<leftarrow> SPECc2 ''sub'' (-) h l;
        if\<^sub>N SPECc2 ''lt'' (<) is_threshold lxs then doN {
          if\<^sub>N SPECc2 ''eq'' (=) d 0 then
            mop_call (slice_sort_specT (cost ''sort_c'' (enat (\<mu> lxs))) (\<^bold><) xs l h)
          else doN {
            (xs,m)\<leftarrow>partition3_spec xs l h;
            d' \<leftarrow> SPECc2 ''sub'' (-) d 1;
            xs \<leftarrow> introsort_aux (xs,l,m,d');
            xs \<leftarrow> introsort_aux (xs,m,h,d');
            RETURN xs
          }
        }
        else
          RETURN xs
      }) (xs,l,h,d)"

definition "introsort_aux3d th xs l h \<equiv>
  doN {
    d \<leftarrow> SPECT [Discrete.log (h-l) * 2 \<mapsto> cost ''depth'' 1];
    introsort_aux3 th xs l h d
  }"

corollary my_slice_sort_spec_refine_sort':
  fixes T :: "nat \<Rightarrow> ecost"
  assumes "\<And>x. T x \<le> timerefineA (TId(''list_append'' := 0, ''list_length'' := cost ''sub'' 1)) (T x)"
  shows
 "\<lbrakk>(xsi,xs) \<in> slicep_rel l h; xsi'=xsi; l'=l; h'=h; xa=length xs\<rbrakk>
\<Longrightarrow> slice_sort_specT (T xa) lt xsi l h \<le>\<Down>(slice_rel xsi' l' h') (timerefine TR_i_aux (SPEC (sort_spec lt xs) (\<lambda>_. T (length xs))))"
  unfolding slice_sort_specT_def sort_spec_def
  apply(intro refine0 SPEC_both_refine)
  subgoal by (auto simp: slicep_rel_def)
  subgoal for x apply (auto simp: slice_rel_alt)
    apply(rule exI[where x="slice l h x"]) apply (auto simp: eq_outside_range_def slicep_rel_def)
    using assms by blast
  done

subsubsection \<open>Refinement Theorem\<close>

lemma introsort_aux3_refine: "(xsi,xs)\<in>slicep_rel l h
       \<Longrightarrow> introsort_aux3 \<mu> xsi l h d \<le> \<Down>(slice_rel xsi l h) (timerefine TR_i_aux (introsort_aux2 \<mu> xs d))"
      unfolding introsort_aux3_def introsort_aux2_def

      supply recref = RECT'_dep_refine[where
          R="\<lambda>_. {((xsi::'a list, l, h, di::nat), (xs, d)). (xsi, xs) \<in> slicep_rel l h \<and> di=d}" and
          S="\<lambda>_ (xsi::'a list, l, h, di::nat). slice_rel xsi l h" and
          arb\<^sub>0 = "()"
          ]
      thm recref
      unfolding SPECc2_def
      apply (refine_rcg bindT_refine_conc_time_my_inres MIf_refine mop_call_refine
        recref
        partition3_spec_refine'
        my_slice_sort_spec_refine_sort'
        ; (intro refl)?; (intro wfR''_TId wfR''_upd sp_TId struct_preserving_upd_I )?
        )
            apply refine_mono

      apply(all \<open>(simp add: timerefineA_update_cost timerefineA_update_apply_same_cost;fail)?\<close>)

     focus  apply (auto simp: slicep_rel_def)    solved
     focus apply refine_dref_type apply (auto simp: slicep_rel_def)    solved
     focus apply refine_dref_type apply (auto simp: slicep_rel_def)   solved
     focus apply simp   solved
     focus apply refine_dref_type apply simp   solved
     focus   apply simp   solved
     focus apply simp solved
     focus apply simp apply (rule IdI) solved
     focus (* first recursive call *)
      apply(rprems)
      apply(auto simp: slice_rel_alt idx_shift_rel_def slicep_rel_take)
      solved
     focus (* second recursive call *)
      apply(rprems)
      apply(auto simp: slice_rel_alt idx_shift_rel_def slicep_rel_eq_outside_range slicep_rel_drop)
      solved
     focus
      unfolding RETURNT_alt
      apply(rule SPECT_refine_t)
      subgoal
        apply (clarsimp simp: slice_rel_alt idx_shift_rel_def)
        apply (rule conjI)
        subgoal
          apply (rule slicep_rel_append)
          apply (subst slicep_rel_eq_outside_range; assumption?)
          by auto
        subgoal
          apply (drule (1) eq_outside_range_gen_trans[OF _ _ refl refl])
          apply (erule (1) eq_outside_range_gen_trans)
          apply (auto simp: max_def algebra_simps slicep_rel_def split: if_splits)
          done
        done
      apply (simp add: timerefineA_update_apply_same_cost)
      solved
     focus
      subgoal by (auto simp: slicep_rel_def)
      solved
      done

    definition "slice_part_sorted_spec xsi l h \<equiv> doN { ASSERT (l<h \<and> h\<le>length xsi); SPEC (\<lambda>xsi'.
        eq_outside_range xsi' xsi l h
      \<and> mset (slice l h xsi') = mset (slice l h xsi)
      \<and> part_sorted_wrt (le_by_lt (\<^bold><)) is_threshold (slice l h xsi')) (\<lambda>_. cost ''slice_part_sorted'' 1)}"

  definition "slice_part_sorted_specT T xsi l h \<equiv> doN { ASSERT (l<h \<and> h\<le>length xsi); SPEC (\<lambda>xsi'.
      eq_outside_range xsi' xsi l h
    \<and> mset (slice l h xsi') = mset (slice l h xsi)
    \<and> part_sorted_wrt (le_by_lt (\<^bold><)) is_threshold (slice l h xsi')) (\<lambda>_. T)}"

lemma spss_eq1: "slice_part_sorted_spec xs l h = slice_part_sorted_specT (cost ''slice_part_sorted'' 1) xs l h"
  unfolding slice_part_sorted_specT_def slice_part_sorted_spec_def by auto

lemma spss_eq2: "timerefine R (slice_part_sorted_specT T xs l h) = slice_part_sorted_specT (timerefineA R T) xs l h"
  unfolding slice_part_sorted_specT_def
  by(cases "l < h \<and> h \<le> length xs", auto simp add: SPEC_timerefine_conv)

definition introsort_aux_cost2 :: "_ \<Rightarrow> 'a list * nat \<Rightarrow> (char list, enat) acost"
  where "introsort_aux_cost2 \<mu> = (\<lambda>(xs, d). lift_acost (
        cost ''if'' (2^(d+1)-1) + cost ''eq'' (2^(d+1)-1) + cost ''if'' (2^(d+1)-1)
         + cost ''lt'' (2^(d+1)-1) + cost ''call'' ((2^(d+1)-1))
        + cost ''sub'' (2^(d+1)-1)
        +   cost ''sort_c'' (\<mu> (length xs)) + cost ''partition_c'' (d*(length xs))
         )
        )"

lemma wfR''TR_i_aux[simp]: "wfR'' TR_i_aux"
  by (auto intro!: wfR''_upd)

subsubsection \<open>Correctness Theorem\<close>

lemma introsort_aux3_correct:
  assumes tf_sums: "\<And>a b. \<mu> a + \<mu> b \<le> \<mu> (a+b)"
  shows
    "introsort_aux3 \<mu> xsi l h d \<le> \<Down>Id (timerefine TR_i_aux (timerefine (TId(''slice_part_sorted'':=introsort_aux_cost \<mu> (h-l,d))) (slice_part_sorted_spec xsi l h)))"
  apply(subst timerefine_iter2)
  subgoal by simp
  subgoal by (auto intro!: wfR''_upd)
  unfolding slice_part_sorted_spec_def
  apply(rule refine0)
  apply(rule order.trans)
   apply(rule introsort_aux3_refine[of xsi "slice l h xsi" l h \<mu> d])
  subgoal by (auto simp: slicep_rel_def)
  apply(rule order.trans)
  apply(rule nrest_Rel_mono)
  apply(rule timerefine_mono2)
  subgoal by (auto intro!: wfR''_upd)
  apply(rule introsort_aux2_refine)
  apply(rule order.trans)
  apply(rule nrest_Rel_mono)
  apply(rule timerefine_mono2)
  subgoal by (auto intro!: wfR''_upd)
  apply(rule order.trans)
  apply(rule nrest_Rel_mono)
  apply(rule timerefine_mono2)
  subgoal by (auto intro!: wfR''_upd)
    apply(rule introsort_aux1_correct)
  subgoal using tf_sums by blast
  subgoal by (metis Sorting_Misc.slice_len length_greater_0_conv zero_less_diff)
   apply (simp add: timerefine_id)
   apply(rule order.refl)
  apply(simp add: SPEC_timerefine_conv introsort_aux_cost_def timerefineA_propagate lift_acost_propagate lift_acost_cost
          timerefineA_update_apply_same_cost' costmult_cost pp_fun_upd )
  unfolding slice_rel_def
  apply(subst build_rel_SPEC_conv) apply simp
  apply(rule SPEC_leq_SPEC_I_strong)
  subgoal apply auto
    by (simp add: eq_outside_range_def)
  subgoal by(sc_solve)
  done

lemma TR_sps_important:
  assumes "TR ''slice_part_sorted'' = timerefineA TR_i_aux (introsort_aux_cost \<mu> (h-l,d))"
  shows "(timerefine TR_i_aux (timerefine (TId(''slice_part_sorted'':=introsort_aux_cost \<mu> (h-l,d))) (slice_part_sorted_spec xsi l h)))
    = timerefine TR (slice_part_sorted_spec xsi l h)"
  unfolding slice_part_sorted_spec_def
  apply(cases "l < h \<and> h \<le> length xsi") apply auto
  apply(simp only: SPEC_timerefine_conv)
  apply(rule SPEC_cong, simp)
  apply(rule ext)
  apply(simp add: norm_cost)
  apply(subst assms(1))
  apply simp
  done

lemma introsort_aux3_correct_flexible:
  assumes  tf_sums: "\<And>a b. \<mu> a + \<mu> b \<le> \<mu> (a+b)"
   and  "TR ''slice_part_sorted'' = timerefineA TR_i_aux (introsort_aux_cost \<mu> (h-l,d))"
  shows
    "introsort_aux3 \<mu> xsi l h d \<le> \<Down>Id (timerefine TR (slice_part_sorted_spec xsi l h))"
  apply(subst TR_sps_important[symmetric, where TR=TR, OF assms(2)])
  apply(rule introsort_aux3_correct) by fact+

subsection \<open>introsort3\<close>

    definition "final_sort_spec xs l h \<equiv> doN {
      ASSERT (h-l>1 \<and> part_sorted_wrt (le_by_lt (\<^bold><)) is_threshold (slice l h xs));
      slice_sort_spec (\<^bold><) xs l h
      }"

    definition "introsort3 xs l h \<equiv> doN {
      ASSERT(l\<le>h);
      hml \<leftarrow> SPECc2 ''sub'' (-) h l;
      if\<^sub>N SPECc2 ''lt'' (<) 1 hml then doN {
        xs \<leftarrow> slice_part_sorted_spec xs l h;
        xs \<leftarrow> final_sort_spec xs l h;
        RETURN xs
      } else RETURN xs
    }"

definition "introsort3_cost =  cost ''sub'' 1 + cost ''lt'' 1 + cost ''if'' 1
                               + cost ''slice_sort'' 1 +  cost ''slice_part_sorted'' (1::enat)"

subsubsection \<open>Correctness Theorem\<close>

lemma introsort3_correct: "introsort3 xs l h \<le> \<Down>Id (timerefine (TId(''slice_sort'':=introsort3_cost)) (slice_sort_spec (\<^bold><) xs l h))"
  apply (cases "l\<le>h \<and> h\<le>length xs")
  subgoal
    apply (cases "1<h-l")
    subgoal
      unfolding introsort3_def slice_part_sorted_spec_def final_sort_spec_def slice_sort_spec_alt
      apply(auto intro: refine0 simp: SPEC_timerefine_conv)
      unfolding SPEC_REST_emb'_conv SPECc2_def
      apply(rule gwp_specifies_I)
      apply(refine_vcg \<open>-\<close> rules: )
      subgoal unfolding emb'_def apply(rule If_le_Some_rule2)
         apply (auto simp: introsort3_cost_def timerefineA_update_apply_same_cost' add.commute)
        subgoal apply sc_solve  apply safe by auto
        by (auto elim: eq_outside_range_gen_trans[of _ _ l h _ l h l h, simplified])
      subgoal using eq_outside_range_list_rel_on_conv by blast
      subgoal by blast
      subgoal by linarith
      subgoal by(auto simp: emb'_def)
      done
    subgoal
      unfolding introsort3_def slice_sort_spec_alt slice_part_sorted_spec_def final_sort_spec_def
      apply(auto intro: refine0 simp: SPEC_timerefine_conv)
      unfolding SPEC_REST_emb'_conv SPECc2_def
      apply(rule gwp_specifies_I)
      apply(refine_vcg \<open>-\<close> rules: )
      apply (auto simp add: emb'_def eq_outside_range_triv sorted_wrt01 le_fun_def ecost_nneg)
      apply (auto simp: introsort3_cost_def timerefineA_update_apply_same_cost' add.commute)
      apply sc_solve apply auto
      done
    done
  subgoal
    unfolding slice_sort_spec_alt
    by(auto intro: refine0)
  done

  end

end


theory Random_Stream_Monad
  imports
    "HOL-Library.Stream"
    "HOL-Library.While_Combinator"
    "HOL-Library.Log_Nat"
    "HOL-Probability.Stream_Space"
    "HOL-Probability.Probability_Mass_Function"
    "HOL-Probability.Giry_Monad"
begin

text \<open>Abbreviation for the discrete $\sigma$-algebra. (Do not use, if the measure is important.)\<close>

abbreviation discrete where "discrete \<equiv> count_space UNIV"

text \<open>Preliminary measure theoretic results:\<close>

lemma (in sigma_finite_measure) restrict_space_pair_lift:
  assumes "A' \<in> sets A"
  shows "restrict_space A A' \<Otimes>\<^sub>M M = restrict_space (A \<Otimes>\<^sub>M M) (A' \<times> space M)" (is "?L = ?R")
proof -
  let ?X = "((\<inter>) (A' \<times> space M) ` {a \<times> b |a b. a \<in> sets A \<and> b \<in> sets M})"
  have 0: "A' \<subseteq> space A"
    using assms sets.sets_into_space by blast

  have "?X \<subseteq> {a \<times> b |a b. a \<in> sets (restrict_space A A') \<and> b \<in> sets M}"
  proof (rule image_subsetI)
    fix x assume "x \<in> {a \<times> b |a b. a \<in> sets A \<and> b \<in> sets M}"
    then obtain u v where uv_def: "x = u \<times> v" "u \<in> sets A" "v \<in> sets M"
      by auto
    have 8:"u \<inter> A' \<in> sets (restrict_space A A')"
      using uv_def(2) unfolding sets_restrict_space by auto
    have "v \<subseteq> space M"
      using uv_def(3) sets.sets_into_space by auto
    hence "A' \<times> space M \<inter> x = (u \<inter> A') \<times> v"
      unfolding uv_def(1) by auto
    also have "... \<in> {a \<times> b |a b. a \<in> sets (restrict_space A A') \<and> b \<in> sets M}"
      using 8 uv_def(3) by auto

    finally show "A' \<times> space M \<inter> x \<in> {a \<times> b |a b. a \<in> sets (restrict_space A A') \<and> b \<in> sets M}"
      by simp
  qed
  moreover have "{a \<times> b |a b. a \<in> sets (restrict_space A A') \<and> b \<in> sets M} \<subseteq> ?X"
  proof (rule subsetI)
    fix x assume "x \<in> {a \<times> b |a b. a \<in> sets (restrict_space A A') \<and> b \<in> sets M}"
    then obtain u v where uv_def: "x = u \<times> v" "u \<in> sets (restrict_space A A')" "v \<in> sets M"
      by auto

    have "x = (A' \<times> space M) \<inter> x"
      unfolding uv_def(1) using uv_def(2,3) sets.sets_into_space
      by (intro Int_absorb1[symmetric]) (auto simp add:sets_restrict_space)
    moreover have "u \<in> sets A" using uv_def(2) assms unfolding sets_restrict_space by blast
    hence "x \<in> {a \<times> b |a b. a \<in> sets A \<and> b \<in> sets M}"
      unfolding uv_def(1) using uv_def(3) by auto
    ultimately show "x \<in> ?X"
      by simp
  qed
  ultimately have 1: "?X = {a \<times> b |a b. a \<in> sets (restrict_space A A') \<and> b \<in> sets M}" by simp

  have "sets ?R = sigma_sets (A'\<times>space M) ((\<inter>) (A'\<times>space M) ` {a\<times>b |a b. a \<in> sets A\<and>b \<in> sets M})"
    unfolding sets_restrict_space sets_pair_measure using assms  sets.sets_into_space
    by (intro sigma_sets_Int sigma_sets.Basic) auto
  also have "... = sets (restrict_space A A' \<Otimes>\<^sub>M M)"
    unfolding sets_pair_measure space_restrict_space Int_absorb2[OF 0] sets_restrict_space 1
    by auto
  finally have 2:"sets (restrict_space (A \<Otimes>\<^sub>M M) (A' \<times> space M)) = sets (restrict_space A A' \<Otimes>\<^sub>M M)"
    by simp

  have 3: "emeasure (restrict_space A A'\<Otimes>\<^sub>MM) S = emeasure (restrict_space (A\<Otimes>\<^sub>MM) (A'\<times>space M)) S"
    (is "?L1 = ?R1") if 4:"S \<in> sets (restrict_space A A' \<Otimes>\<^sub>M M)" for S
  proof -
    have "Pair x -` S = {}" if "x \<notin> A'" "x \<in> space A" for x
      using that 4 by (auto simp add:2[symmetric] sets_restrict_space)
    hence 5: "emeasure M (Pair x -` S) = 0" if "x \<notin> A'" "x \<in> space A" for x
      using that by auto
    have "?L1 = (\<integral>\<^sup>+ x. emeasure M (Pair x -` S) \<partial>restrict_space A A')"
      by (intro emeasure_pair_measure_alt[OF that])
    also have "... = (\<integral>\<^sup>+x\<in>A'. emeasure M (Pair x -` S) \<partial>A)"
      using assms by (intro nn_integral_restrict_space) auto
    also have "... = (\<integral>\<^sup>+x. emeasure M (Pair x -` S) \<partial>A)"
      using 5 by (intro nn_integral_cong) force
    also have "... = emeasure (A \<Otimes>\<^sub>M M) S"
      using that assms by (intro emeasure_pair_measure_alt[symmetric])
        (auto simp add:2[symmetric] sets_restrict_space)
    also have "... = ?R1"
      using assms that by (intro emeasure_restrict_space[symmetric])
        (auto simp add:2[symmetric] sets_restrict_space)
    finally show ?thesis by simp
  qed

  show ?thesis using 2 3
    by (intro measure_eqI) auto
qed

lemma pair_prob_spaceI:
  assumes "prob_space M" "prob_space N"
  shows "pair_prob_space M N"
  using assms
  by (simp add: pair_prob_space.intro pair_sigma_finite_def prob_space_imp_sigma_finite)

text \<open>
  The following is highly experimental.

  It probably makes sense to abstract over the stream type, or use a more general interface
  for randomness sources.
\<close>

type_synonym 'a random_monad = "bool stream \<Rightarrow> 'a \<times> bool stream"

definition rm_return :: "'a \<Rightarrow> 'a random_monad"
  where "rm_return x f = (x, f)"

definition rm_bind :: "'a random_monad \<Rightarrow> ('a \<Rightarrow> 'b random_monad) \<Rightarrow> 'b random_monad"
  where "rm_bind x y f = case_prod y (x f)"

adhoc_overloading Monad_Syntax.bind rm_bind

definition rm_coin :: "bool random_monad"
  where "rm_coin f = (shd f, stl f)"

definition coin :: "bool pmf"
  where "coin = pmf_of_set UNIV"

definition coin_space :: "bool stream measure"
  where "coin_space = stream_space (measure_pmf (pmf_of_set UNIV))"

lemma coin_space: "prob_space coin_space"
  unfolding coin_space_def
  by (intro prob_space.prob_space_stream_space prob_space_measure_pmf)

lemma space_coin_space: "space coin_space = UNIV"
  unfolding coin_space_def space_stream_space by simp

text \<open>This is the relation expressing that a randomized algorithm is a sampler for a given pmf.\<close>

definition sampler_for where
  "sampler_for r m = (
    distr coin_space (discrete \<Otimes>\<^sub>M coin_space) r = measure_pmf m \<Otimes>\<^sub>M coin_space \<and>
    r \<in> coin_space \<rightarrow>\<^sub>M discrete \<Otimes>\<^sub>M coin_space \<and> countable (range (fst \<circ> r)))"

lemma rm_range_intro:
  assumes "sampler_for r m"
  assumes "x \<in> set_pmf m"
  shows "x \<in> range (fst \<circ> r)"
proof (rule ccontr)
  assume "x \<notin> range (fst \<circ> r)"
  hence "False " if  "y \<in> r -` ({x} \<times> UNIV)" for y
    by (metis rangeI set.compositionality that vimage_fst vimage_singleton_eq)
  hence a:"r -` ({x} \<times> UNIV) = {}"
    by blast

  have "0 < emeasure (measure_pmf m) {x} * 1"
    using assms(2) by (simp add: emeasure_pmf_single_eq_zero_iff order_less_le)
  also have "... = emeasure (measure_pmf m) {x} * emeasure (coin_space) (space coin_space)"
    using coin_space
    by (intro arg_cong2[where f="(*)"] refl prob_space.emeasure_space_1[symmetric])
  also have "... = emeasure (measure_pmf m \<Otimes>\<^sub>M coin_space) ({x} \<times> space coin_space)"
    using prob_space_imp_sigma_finite[OF coin_space]
    by (intro sigma_finite_measure.emeasure_pair_measure_Times[symmetric]) simp_all
  also have "... = emeasure (distr coin_space (discrete \<Otimes>\<^sub>M coin_space) r) ({x} \<times> space coin_space)"
    using assms(1) unfolding sampler_for_def by simp
  also have "... = emeasure coin_space (r -` ({x} \<times> space coin_space) \<inter> space coin_space)"
    using assms(1) unfolding sampler_for_def by (intro emeasure_distr pair_measureI) auto
  also have "... = emeasure coin_space {}"
    using a by (intro arg_cong2[where f="emeasure"]) auto
  also have "... = 0"
    by simp
  finally show "False" by simp
qed

lemma sampler_for_bind:
  assumes "sampler_for f' f"  (* TODO: Rename args r for randomized algorith, m for measure *)
  assumes "\<And>x. x \<in> range (fst \<circ> f') \<Longrightarrow> sampler_for (g' x) (g x)"
  shows "sampler_for (rm_bind f' g') (f \<bind> g) "
proof -
  let ?M = "measure_pmf f"
  let ?N = "\<lambda>x. measure_pmf (g x)"
  let ?C = "coin_space"
  let ?D = "discrete"
  let ?R = "restrict_space ?M (range (fst \<circ> f'))"

  have m_g: "case_prod g' \<in> ?R \<Otimes>\<^sub>M ?C \<rightarrow>\<^sub>M ?D \<Otimes>\<^sub>M ?C"
    using assms unfolding sampler_for_def
    by (intro measurable_pair_restrict_pmf1) simp_all

  have "g \<in> ?R \<rightarrow>\<^sub>M ?D"
    by (intro measurable_restrict_space1) simp
  hence m_g_2: "(\<lambda>x. measure_pmf (g x)) \<in> ?R \<rightarrow>\<^sub>M subprob_algebra discrete"
    by measurable

  have 3:"x \<in> range (fst \<circ> f')" if "x \<in> set_pmf f" for x
    using rm_range_intro[OF assms(1) that] by simp

  have 2:"prob_space ?R" using 3
    by (intro Probability_Measure.prob_space_restrict_space measure_pmf.emeasure_eq_1_AE AE_pmfI)
     simp_all
  hence 0:"prob_space (?R \<bind> ?N)"
    by (intro prob_space.prob_space_bind[where S="discrete"] AE_I2 m_g_2 prob_space_measure_pmf)
      auto

  have m_f_3: "f' x \<in> range (\<lambda>x. fst (f' x)) \<times> space coin_space" for x
    unfolding space_coin_space using vimage_fst by fastforce
  hence m_f: "f' \<in> ?C \<rightarrow>\<^sub>M ?R \<Otimes>\<^sub>M ?C"
    using coin_space prob_space_imp_sigma_finite assms(1) unfolding sampler_for_def
    by (subst sigma_finite_measure.restrict_space_pair_lift)
       (auto intro: measurable_restrict_space2)

  have "f' \<in> coin_space \<rightarrow>\<^sub>M discrete \<Otimes>\<^sub>M coin_space"
    using assms(1) unfolding sampler_for_def by auto
  hence m_f_2: "f' \<in> ?C \<rightarrow>\<^sub>M ?M \<Otimes>\<^sub>M ?C"
    by measurable

  have "distr ?C (?R \<Otimes>\<^sub>M ?C) f' = distr ?C (restrict_space (?M\<Otimes>\<^sub>M?C) (range (fst\<circ>f')\<times>space ?C)) f'"
    using prob_space_imp_sigma_finite[OF coin_space]
    by (subst sigma_finite_measure.restrict_space_pair_lift) auto
  also have "... = restrict_space (distr ?C (?M \<Otimes>\<^sub>M ?C) f') (range (fst \<circ> f') \<times> space coin_space)"
    using m_f_2 m_f_3
    by (intro restrict_distr[symmetric]) (simp_all add: Pi_def)
  also have "... = restrict_space (distr ?C (?D \<Otimes>\<^sub>M ?C) f') (range (fst \<circ> f') \<times> space coin_space)"
    by (intro arg_cong2[where f="restrict_space"] distr_cong refl) (simp add:sets_pair_measure)
  also have "... = restrict_space (?M \<Otimes>\<^sub>M ?C) (range (fst \<circ> f') \<times> space coin_space)"
    using assms(1) unfolding sampler_for_def by simp
  also have "... = ?R \<Otimes>\<^sub>M ?C"
    using prob_space_imp_sigma_finite[OF coin_space]
    by (intro sigma_finite_measure.restrict_space_pair_lift[symmetric]) simp_all
  finally have 1: "distr ?C (?R \<Otimes>\<^sub>M ?C) f' = ?R \<Otimes>\<^sub>M ?C"
    by simp

  have "distr ?C (?D \<Otimes>\<^sub>M ?C) (rm_bind f' g') = distr ?C (?D \<Otimes>\<^sub>M ?C) (\<lambda>\<omega>. case_prod g' (f' \<omega>))"
    unfolding rm_bind_def by simp
  also have "... =  distr (distr ?C (?R \<Otimes>\<^sub>M ?C) f') (?D \<Otimes>\<^sub>M ?C) (case_prod g')"
    using m_g m_f by (subst distr_distr) (simp_all add:comp_def)
  also have "... = distr (?R \<Otimes>\<^sub>M ?C) (?D \<Otimes>\<^sub>M ?C) (case_prod g')"
    unfolding 1 by simp
  also have "... = distr (?R \<bind> (\<lambda>x. ?C \<bind> (\<lambda>\<omega>. return(?R\<Otimes>\<^sub>M?C) (x,\<omega>)))) (?D\<Otimes>\<^sub>M?C) (case_prod g')"
    using pair_prob_spaceI[OF 2 coin_space]
    by (subst pair_prob_space.pair_measure_eq_bind) simp_all
  also have "... = ?R \<bind>(\<lambda>x. distr(?C \<bind>(\<lambda>\<omega>. return(?R\<Otimes>\<^sub>M?C)(x, \<omega>)))(?D \<Otimes>\<^sub>M ?C)(\<lambda>(x, y). g' x y))"
    using prob_space_imp_subprob_space[OF coin_space]
    by (intro distr_bind[where K="?R \<Otimes>\<^sub>M ?C"] m_g measurable_bind[where N="?C"] measurable_const)
       (simp_all add:space_subprob_algebra)
  also have "... =?R \<bind>(\<lambda>x.?C \<bind>(\<lambda>y. distr (return(?R\<Otimes>\<^sub>M?C)(x, y))(?D \<Otimes>\<^sub>M ?C)(\<lambda>(x, y). g' x y)))"
    using m_g
    by (intro bind_cong ext distr_bind[where K="?R \<Otimes>\<^sub>M ?C"] refl) (simp_all add:space_coin_space)
  also have "... =?R \<bind> (\<lambda>x. ?C \<bind> (\<lambda>y. return (?D \<Otimes>\<^sub>M ?C) (case (x, y) of (x, y) \<Rightarrow> g' x y))) "
    using m_g
    by (intro bind_cong distr_return) (simp_all add:space_restrict_space space_pair_measure)
  also have "... = ?R \<bind> (\<lambda>x. ?C \<bind> (\<lambda>y. return (?D \<Otimes>\<^sub>M ?C) (g' x y)))"
    by simp
  also have "... = ?R \<bind> (\<lambda>x. distr ?C (?D \<Otimes>\<^sub>M ?C) (g' x))"
    using assms(2) unfolding sampler_for_def
    by (intro bind_cong bind_return_distr' refl) (simp_all add:space_coin_space)
  also have "... = ?R \<bind> (\<lambda>x. ?N x \<Otimes>\<^sub>M ?C)"
    using assms(2) unfolding sampler_for_def
    by (intro bind_cong refl) simp
  also have "... = ?R \<bind> (\<lambda>x. ?N x \<bind> (\<lambda>y. ?C \<bind> (\<lambda>\<omega>. return (?N x \<Otimes>\<^sub>M ?C) (y,\<omega>))))"
   using coin_space prob_space_measure_pmf by (intro arg_cong2[where f="bind"]
     pair_prob_space.pair_measure_eq_bind ext pair_prob_space.intro)
     (auto simp add: pair_sigma_finite.intro prob_space_imp_sigma_finite)
  also have "... = ?R \<bind> (\<lambda>x. ?N x \<bind> (\<lambda>y. ?C \<bind> (\<lambda>\<omega>. return (?D \<Otimes>\<^sub>M ?C) (y,\<omega>))))"
    by (intro arg_cong2[where f="bind"] ext refl return_cong sets_pair_measure_cong) simp
  also have "... = (?R \<bind> ?N) \<bind> (\<lambda>y. ?C \<bind> (\<lambda>\<omega>. return (?D \<Otimes>\<^sub>M ?C) (y, \<omega>)))"
    using prob_space_imp_subprob_space[OF coin_space]
    by (intro bind_assoc[symmetric, where N="discrete" and R="discrete \<Otimes>\<^sub>M coin_space"] m_g_2)
      (auto intro!:subprob_space.bind_in_space simp add:Pi_def)
  also have "... = (?R \<bind> ?N) \<bind> (\<lambda>y. ?C \<bind> (\<lambda>\<omega>. return ((?R \<bind> ?N) \<Otimes>\<^sub>M ?C) (y, \<omega>)))"
    by (intro arg_cong2[where f="bind"] ext refl return_cong sets_pair_measure_cong) auto
  also have "... = (?R \<bind> ?N) \<Otimes>\<^sub>M ?C"
    using pair_prob_spaceI[OF 0 coin_space]
    by (subst (2) pair_prob_space.pair_measure_eq_bind) simp_all
  also have "... = (?M \<bind> (\<lambda>x. if x \<in> range (fst \<circ> f') then ?N x else
    null_measure (measure_pmf (g (SOME x. x \<in> range (fst \<circ> f') \<and> x \<in> space ?M))))) \<Otimes>\<^sub>M ?C"
    using m_g_2 by (subst bind_restrict_space[where N="discrete"]) auto
  also have "... = (?M \<bind> ?N)  \<Otimes>\<^sub>M ?C"
    using 3 by (intro bind_cong_AE[where B="discrete"] arg_cong2[where f="(\<Otimes>\<^sub>M)"] refl AE_pmfI)
      (simp_all add:space_subprob_algebra Pi_def subprob_spaceI)+
  also have "... = measure_pmf (bind_pmf f g) \<Otimes>\<^sub>M coin_space"
    unfolding measure_pmf_bind by simp
  finally have 4:"distr coin_space (?D \<Otimes>\<^sub>M ?C) (rm_bind f' g') = measure_pmf (bind_pmf f g) \<Otimes>\<^sub>M ?C"
    by simp

  have 5:"(\<lambda>\<omega>. rm_bind f' g' \<omega>) \<in> ?C \<rightarrow>\<^sub>M (?D \<Otimes>\<^sub>M ?C)"
    using m_f m_g unfolding rm_bind_def by simp

  have "range (fst \<circ> rm_bind f' g') = range (\<lambda>x. fst (g' (fst (f' x)) (snd (f' x))))"
    unfolding rm_bind_def by (simp add:case_prod_beta)
  also have "... \<subseteq> (\<Union>x \<in> range (fst \<circ> f'). range (fst \<circ> g' x))"
    by (intro subsetI) auto
  finally have "range (fst \<circ> rm_bind f' g') \<subseteq> (\<Union>x \<in> range (fst \<circ> f'). range (fst \<circ> g' x))"
    by simp
  moreover have "countable (\<Union>x \<in> range (fst \<circ> f'). range (fst \<circ> g' x))"
    using assms unfolding sampler_for_def by (intro countable_UN) auto

  ultimately have 6:"countable (range (fst \<circ> rm_bind f' g'))"
    using countable_subset by auto

  show ?thesis using 4 5 6
    unfolding sampler_for_def by blast
qed

lemma sampler_for_coin: "sampler_for rm_coin coin"
proof -
  let ?C = "coin_space"
  let ?D = "discrete"
  let ?M = "measure_pmf (pmf_of_set (UNIV :: bool set))"

  have "emeasure (distr ?C (?D \<Otimes>\<^sub>M ?C) rm_coin) A = emeasure (measure_pmf coin \<Otimes>\<^sub>M ?C) A"
    (is "?L = ?R") if "A \<in> sets (distr ?C (?D \<Otimes>\<^sub>M ?C) rm_coin)" for A
  proof -
    have "A \<in> sets (?D \<Otimes>\<^sub>M ?C)"
      using that by simp
    moreover have "(\<lambda>x. (shd x, stl x)) \<in> ?C \<rightarrow>\<^sub>M ?D \<Otimes>\<^sub>M ?C"
      unfolding coin_space_def by simp
    ultimately have 0: "{x. (shd x, stl x) \<in> A} \<in> sets ?C"
      unfolding measurable_def space_coin_space vimage_def by auto

    have "?L = emeasure ?C (rm_coin -` A \<inter> space ?C)"
      using that unfolding rm_coin_def coin_space_def by (subst emeasure_distr) simp_all
    also have "... = emeasure ?C {x. (shd x, stl x) \<in> A}"
      unfolding rm_coin_def vimage_def space_coin_space by simp
    also have "... = \<integral>\<^sup>+ t. emeasure ?C {x \<in> space ?C. t ## x \<in> {x. (shd x, stl x) \<in> A}} \<partial>?M"
      using 0 unfolding coin_space_def
      by (intro prob_space.emeasure_stream_space prob_space_measure_pmf) simp
    also have "... = \<integral>\<^sup>+ t. emeasure ?C {x. (t,x) \<in> A} \<partial>?M"
      unfolding space_coin_space by simp
    also have "... = \<integral>\<^sup>+ x. emeasure coin_space (Pair x -` A) \<partial>measure_pmf coin"
      unfolding coin_def vimage_def by simp
    also have "... = ?R"
      using prob_space_imp_sigma_finite[OF coin_space] that
      by (intro sigma_finite_measure.emeasure_pair_measure_alt[symmetric]) simp_all
    finally show ?thesis
      by simp
  qed
  hence "distr ?C (?D \<Otimes>\<^sub>M ?C) rm_coin = measure_pmf coin \<Otimes>\<^sub>M ?C"
    by (intro measure_eqI) (auto simp add:sets_pair_measure)

  thus ?thesis
    unfolding sampler_for_def coin_space_def rm_coin_def by simp
qed

lemma sampler_for_return: "sampler_for (rm_return x) (return_pmf x) "
proof -
  let ?C = "coin_space"
  let ?D = "discrete"

  have "emeasure (distr ?C (?D \<Otimes>\<^sub>M ?C) (Pair x)) A = emeasure (measure_pmf (return_pmf x) \<Otimes>\<^sub>M ?C) A"
    (is "?L = ?R") if "A \<in> sets (distr ?C (?D \<Otimes>\<^sub>M ?C) (Pair x))" for A
  proof -
    have "?L = emeasure ?C (Pair x -` A \<inter> space ?C)"
      using that by (intro emeasure_distr) simp_all
    also have "... = \<integral>\<^sup>+ x. emeasure ?C (Pair x -` A) \<partial>measure_pmf (return_pmf x)"
      unfolding space_coin_space by simp
    also have "... = ?R"
      using prob_space_imp_sigma_finite[OF coin_space] that
      by (intro sigma_finite_measure.emeasure_pair_measure_alt[symmetric]) simp_all
    finally show ?thesis
      by simp
  qed
  hence "distr coin_space (?D \<Otimes>\<^sub>M ?C) (Pair x) = measure_pmf (return_pmf x) \<Otimes>\<^sub>M ?C"
    by (intro measure_eqI)  (simp_all add:sets_pair_measure)
  thus ?thesis
    unfolding sampler_for_def rm_return_def by simp
qed

lemmas sampler_for_simps = sampler_for_return sampler_for_bind sampler_for_coin

subsection \<open>Sampler for numbers in {0..<2^n}\<close>

fun rm_uniform_pow :: "nat \<Rightarrow> nat random_monad"
  where
    "rm_uniform_pow 0 = rm_return 0" |
    "rm_uniform_pow (Suc n) =
      do {
        x \<leftarrow> rm_uniform_pow n;
        y \<leftarrow> rm_coin;
        rm_return (of_bool y + 2*x)
      }"

fun pmf_uniform_pow :: "nat \<Rightarrow> nat pmf"
  where
    "pmf_uniform_pow 0 = return_pmf 0" |
    "pmf_uniform_pow (Suc n) =
      do {
        x \<leftarrow> pmf_uniform_pow n;
        y \<leftarrow> coin;
        return_pmf (of_bool y + 2*x)
      }"

lemma "sampler_for (rm_uniform_pow n) (pmf_uniform_pow n)"
  by (induction n) (simp_all add:sampler_for_simps)

text \<open>Sampler for pmf_of_set {0..n} *inclusive*\<close>

function rm_uniform :: "nat \<Rightarrow> nat random_monad"
  where "rm_uniform n =
    do {
      x \<leftarrow> rm_uniform_pow (floorlog 2 n);
      if x \<le> n then rm_return x else rm_uniform n
    }"
    by pat_completeness auto

(* TODO: Write this using the while combinator or a counter-part of it in the random_monad, the
above variant does not allow code-gen. *)

end