section \<open> Probabilistic Designs \<close>

text \<open> This is the mechanisation of \emph{probabilistic designs}~\cite{Jifeng2004,Woodcock2019} in
Isabelle/UTP.\<close>

theory utp_prob_des
  imports "UTP-Calculi.utp_wprespec" "UTP-Designs.utp_designs" "HOL-Probability.Probability_Mass_Function"
  "HOL-Probability.SPMF"
begin recall_syntax

purge_notation inner (infix "\<bullet>" 70)

declare [[coercion pmf]]

alphabet 's prss =
  prob :: "'s pmf"

text \<open> If the probabilities of two disjoint sample sets sums up to 1, then the probability of the
  first set is equal to 1 minus the probability of the second set. \<close>

lemma pmf_disj_set:
  assumes "X \<inter> Y = {}"
  shows "((\<Sum>\<^sub>a i\<in>(X \<union> Y). pmf M i) = 1) = ((\<Sum>\<^sub>a i\<in>X. pmf M i) = 1 - (\<Sum>\<^sub>a i\<in>Y. pmf M i))"
  by (metis assms diff_eq_eq infsetsum_Un_disjoint pmf_abs_summable)

no_utp_lift ndesign wprespec uwp

text \<open>
  Probabilistic designs (@{text "('s, 's) rel_pdes"}), that map the standard state space to the
probabilistic state space, are heterogeneous.
\<close>

type_synonym ('a, 'b) rel_pdes = "('a, 'b prss) rel_des"
type_synonym 's hrel_pdes = "('s, 's) rel_pdes"
type_synonym 's hrel_hpdes = "('s prss, 's prss) rel_des"

translations
  (type) "('a, 'b) rel_pdes" <= (type) "('a, 'b prss) rel_des"

text \<open>
  The @{term "forget_prob"} is a non-homogeneous design as a forgetful function that maps a discrete
probability distribution @{term "$prob"} at initial observation to a final state.
\<close>
definition forget_prob :: "('s prss, 's) rel_des" ("\<^bold>f\<^bold>p") where
[upred_defs]: "forget_prob = U(true \<turnstile>\<^sub>n ($prob($\<^bold>v\<acute>) > 0))"

text \<open>
  Here, we reverse the order of parameters in order to match with its definition in
original publications such the UTP-book.
\<close>
abbreviation wprespec' :: " ('a, 'c) urel \<Rightarrow> ('b, 'c) urel \<Rightarrow> ('a, 'b) urel"  (infixr "'/\<^sub>w" 70) where
"(Y /\<^sub>w K) \<equiv> (K \\ Y)"

text \<open>
  The weakest prespecification of a standard design @{text "D"} wrt @{text "\<^bold>f\<^bold>p"} is the weakest
probabilistic design, as an embedding of @{text "D"} in the probabilistic world through @{text "\<K>"}.
\<close>

definition pemb :: "('a, 'b) rel_des \<Rightarrow> ('a, 'b) rel_pdes" ("\<K>")
  where [upred_defs]: "pemb D = D /\<^sub>w \<^bold>f\<^bold>p"

lemma pemb_mono: "P \<sqsubseteq> Q \<Longrightarrow> \<K>(P) \<sqsubseteq> \<K>(Q)"
  by (metis (mono_tags, lifting) dual_order.trans order_refl pemb_def wprespec)

lemma wdprespec: "(p \<turnstile>\<^sub>n Q) /\<^sub>w (true \<turnstile>\<^sub>n R)  = (p \<turnstile>\<^sub>n (Q /\<^sub>w R))"
  by (rel_auto)

(*
lemma pemb_form:
  fixes R :: "('a, 'b) urel"
  shows "(R /\<^sub>w U($prob($\<^bold>v\<acute>) > 0)) = \<^U>((\<Sum>\<^sub>a i\<in>{s'.(R wp (&\<^bold>v = s'))\<^sup><}. $prob\<acute> i) = 1)" (is "?lhs = ?rhs")
proof -
  have "?lhs = \<^U>((\<not> (\<not> R) ;; (0 < $prob\<acute>$\<^bold>v)))"
    by (rel_auto)
  also have "... = \<^U>((\<Sum>\<^sub>a i\<in>{s'.(R wp (&\<^bold>v = s'))\<^sup><}. $prob\<acute> i) = 1)"
    apply (rel_auto)
    apply (metis (no_types, lifting) infsetsum_pmf_eq_1 mem_Collect_eq pmf_positive subset_eq)
    apply (metis AE_measure_pmf_iff UNIV_I measure_pmf.prob_eq_1 measure_pmf_conv_infsetsum mem_Collect_eq set_pmf_eq' sets_measure_pmf)
    done
  finally show ?thesis .
qed
*)

lemma pemb_form:
  fixes R :: "('a, 'b) urel"
  shows "(R /\<^sub>w U($prob($\<^bold>v\<acute>) > 0)) = \<^U>((\<Sum>\<^sub>a i | (R wp (&\<^bold>v = i))\<^sup><. $prob\<acute> i) = 1)" (is "?lhs = ?rhs")
proof -
  have "?lhs = \<^U>((\<not> (\<not> R) ;; (0 < $prob\<acute>$\<^bold>v)))"
    by (rel_auto)
  also have "... = \<^U>((\<Sum>\<^sub>a i\<in>{s'.(R wp (&\<^bold>v = s'))\<^sup><}. $prob\<acute> i) = 1)"
    apply (rel_auto)
    apply (metis (no_types, lifting) infsetsum_pmf_eq_1 mem_Collect_eq pmf_positive subset_eq)
    apply (metis AE_measure_pmf_iff UNIV_I measure_pmf.prob_eq_1 measure_pmf_conv_infsetsum
           mem_Collect_eq set_pmf_eq' sets_measure_pmf)
    done
  finally show ?thesis .
qed

text \<open>
  Embedded standard designs are probabilistic designs~\cite[Theorem 1]{Woodcock2019} and
~\cite[Theorem 3.6]{Jifeng2004}.
\<close>
lemma prob_lift [ndes_simp]:
  fixes R :: "('a, 'b) urel" and p :: "'a upred"
  shows "\<K>(p \<turnstile>\<^sub>n R) = \<^U>(p \<turnstile>\<^sub>n ((\<Sum>\<^sub>a i | (R wp (&\<^bold>v = i))\<^sup><. $prob\<acute> i) = 1))"
proof -
  have 1:"\<K>(p \<turnstile>\<^sub>n R) = (p \<turnstile>\<^sub>n (R  /\<^sub>w \<^U>($prob($\<^bold>v\<acute>) > 0)))"
    by (rel_auto)
  have 2:"R /\<^sub>w \<^U>(($prob($\<^bold>v\<acute>) > 0)) = \<^U>((\<Sum>\<^sub>a i | (R wp (&\<^bold>v = i))\<^sup><. $prob\<acute> i) = 1)"
    by (simp add: pemb_form)
  show ?thesis
    by (simp add: "1" "2")
qed

no_utp_lift usubst (0) subst (1)

subsection \<open> wplus \<close>
text \<open>
  Two pmfs can be joined into one by their corresponding weights via @{text "P +\<^bsub>w\<^esub> Q"}
where @{text "w"} is the weight of P.
\<close>
definition wplus :: "'a pmf \<Rightarrow> real \<Rightarrow> 'a pmf \<Rightarrow> 'a pmf" ("(_ +\<^bsub>_\<^esub> _)" [64, 0, 65] 64) where
"wplus P w Q = join_pmf (pmf_of_list [(P, w), (Q, 1 - w)])"

text \<open>
  Query of the probability value of a state @{text "i"} in a joined probability distribution is just
the summation of the query of @{text "i"} in P by its weight @{text "w"} and
the query of @{text "i"} in Q by its weight @{text "(1 - w)"}.
\<close>
lemma pmf_wplus:
  assumes "w \<in> {0..1}"
  shows "pmf (P +\<^bsub>w\<^esub> Q) i = pmf P i * w + pmf Q i * (1 - w)"
proof -
  from assms have pmf_wf_list: "pmf_of_list_wf [(P, w), (Q, 1 - w)]"
    by (auto intro!: pmf_of_list_wfI)
  show ?thesis
  proof (cases "w \<in> {0<..<1}")
    case True
    hence set_pmf:"set_pmf (pmf_of_list [(P, w), (Q, 1 - w)]) = {P, Q}"
      by (subst set_pmf_of_list_eq, auto simp add: pmf_wf_list)
    thus ?thesis
    proof (cases "P = Q")
      case True
      from assms show ?thesis
        apply (auto simp add: wplus_def join_pmf_def pmf_bind)
        apply (subst integral_measure_pmf[of "{P, Q}"])
          apply (auto simp add: set_pmf_of_list pmf_wf_list set_pmf pmf_pmf_of_list)
        apply (simp add: True)
        apply (metis distrib_right eq_iff_diff_eq_0 le_add_diff_inverse mult.commute mult_cancel_left1)
        done
    next
      case False
      then show ?thesis
        apply (auto simp add: wplus_def join_pmf_def pmf_bind)
        apply (subst integral_measure_pmf[of "{P, Q}"])
          apply (auto simp add: set_pmf_of_list pmf_wf_list set_pmf pmf_pmf_of_list)
        done
    qed
  next
    case False
    thm disjE
    with assms have "w = 0 \<or> w = 1"
      by (auto)
    with assms show ?thesis
    proof (erule_tac disjE, simp_all)
      assume w: "w = 0"
      with pmf_wf_list have "set_pmf (pmf_of_list [(P, w), (Q, 1 - w)]) = {Q}"
        apply (simp add: pmf_of_list_remove_zeros(2)[THEN sym])
        apply (subst set_pmf_of_list_eq, auto simp add: pmf_of_list_wf_def)
        done
      with w show "pmf (P +\<^bsub>0\<^esub> Q) i = pmf Q i"
        apply (auto simp add: wplus_def join_pmf_def pmf_bind pmf_wf_list pmf_of_list_remove_zeros(2)[THEN sym])
        apply (subst integral_measure_pmf[of "{Q}"])
          apply (simp_all add: set_pmf_of_list_eq pmf_pmf_of_list pmf_of_list_wf_def)
        done
    next
      assume w: "w = 1"
      with pmf_wf_list have "set_pmf (pmf_of_list [(P, w), (Q, 1 - w)]) = {P}"
        apply (simp add: pmf_of_list_remove_zeros(2)[THEN sym])
        apply (subst set_pmf_of_list_eq, auto simp add: pmf_of_list_wf_def)
        done
      with w show "pmf (P +\<^bsub>1\<^esub> Q) i = pmf P i"
        apply (auto simp add: wplus_def join_pmf_def pmf_bind pmf_wf_list pmf_of_list_remove_zeros(2)[THEN sym])
        apply (subst integral_measure_pmf[of "{P}"])
          apply (simp_all add: set_pmf_of_list_eq pmf_pmf_of_list pmf_of_list_wf_def)
        done
    qed
  qed
qed

lemma wplus_commute:
  assumes "w \<in>{0..1}"
  shows "P +\<^bsub>w\<^esub> Q = Q +\<^bsub>(1 - w)\<^esub> P"
  using assms by (auto intro: pmf_eqI simp add: pmf_wplus)

lemma wplus_idem:
  assumes "w \<in>{0..1}"
  shows "P +\<^bsub>w\<^esub> P = P"
  using assms
  apply (rule_tac pmf_eqI)
  apply (simp add: pmf_wplus)
  by (metis le_add_diff_inverse mult.commute mult_cancel_left2 ring_class.ring_distribs(2))

lemma wplus_zero: "P +\<^bsub>0\<^esub> Q = Q"
  by (auto intro: pmf_eqI simp add: pmf_wplus)

lemma wplus_one: "P +\<^bsub>1\<^esub> Q = P"
  by (auto intro: pmf_eqI simp add: pmf_wplus)

text \<open>
  This is used to prove the associativity of probabilistic choice: @{text "prob_choice_assoc"}.
\<close>
lemma wplus_assoc:
  assumes "w\<^sub>1 \<in> {0..1}" "w\<^sub>2 \<in> {0..1}"
  assumes "(1-w\<^sub>1)*(1-w\<^sub>2)=(1-r\<^sub>2)" "w\<^sub>1=r\<^sub>1*r\<^sub>2"
  shows "P +\<^bsub>w\<^sub>1\<^esub> (Q +\<^bsub>w\<^sub>2\<^esub> R) = (P +\<^bsub>r\<^sub>1\<^esub> Q) +\<^bsub>r\<^sub>2\<^esub> R"
proof (cases "w\<^sub>1 = 0 \<and> w\<^sub>2 = 0")
  case True
  then show ?thesis
    proof -
      from assms(3-4) have t1: "r\<^sub>2=0"
        by (simp add: True)
      then show ?thesis
        by (simp add: wplus_zero True t1)
    qed
next
  case False
  from assms(3) have f1: "r\<^sub>2 = w\<^sub>1+w\<^sub>2-w\<^sub>1*w\<^sub>2"
    proof -
      have f1: "\<forall>r ra. (ra::real) + - r = 0 \<or> \<not> ra = r"
        by simp
      have f2: "\<forall>r ra rb rc. (rc::real) \<cdot> rb + - (ra \<cdot> r) = rc \<cdot> (rb + - r) + (rc + - ra) \<cdot> r"
        by (simp add: mult_diff_mult)
      have f3: "\<forall>r ra. (ra::real) + (r + - ra) = r + 0"
        by fastforce
      have f4: "\<forall>r ra. (ra::real) + ra \<cdot> r = ra \<cdot> (1 + r)"
        by (simp add: distrib_left)
      have f5: "\<forall>r ra. (ra::real) + - r + 0 = ra + - r"
        by linarith
      have f6: "\<forall>r ra. (0::real) + (ra + - r) = ra + - r"
        by simp
      have "1 + - w\<^sub>2 + - (w\<^sub>1 \<cdot> (1 + - w\<^sub>2)) = 1 + (0 + - r\<^sub>2)"
        using f2 f1 by (metis (no_types) add.left_commute add_uminus_conv_diff assms(3) mult.left_neutral)
      then have "1 + (w\<^sub>1 + w\<^sub>1 \<cdot> - w\<^sub>2 + - r\<^sub>2) = 1 + - w\<^sub>2"
        using f6 f5 f4 f3 by (metis (no_types) add.left_commute)
    then show ?thesis
    by linarith
    qed
  then have f2: "r\<^sub>2 \<in> {0..1}"
    using assms(1-2) by (smt assms(3) atLeastAtMost_iff mult_le_one sum_le_prod1)
  from f1 have f2': "(w\<^sub>1+w\<^sub>2-w\<^sub>1*w\<^sub>2) \<ge> w\<^sub>1"
    using assms(1) assms(2) mult_left_le_one_le by auto
  from f1 have f3: "r\<^sub>1 = w\<^sub>1/(w\<^sub>1+w\<^sub>2-w\<^sub>1*w\<^sub>2)"
    by (metis False add.commute add_diff_eq assms(4) diff_add_cancel
        mult_zero_left mult_zero_right nonzero_eq_divide_eq)
  show ?thesis
  proof (cases "w\<^sub>1 = 0")
    case True
    from f3 have ft1: "r\<^sub>1 = 0"
      by (simp add: True)
    from f1 have ft2: "r\<^sub>2 = w\<^sub>2"
      by (simp add: True)
    then show ?thesis
      using ft1 ft2 assms(1-2)
      by (simp add: True wplus_zero)
  next
    case False
    from f3 f2' have ff1: "r\<^sub>1 \<le> 1"
      using False
      by (metis assms(4) atLeastAtMost_iff eq_iff f1 f2 le_cases le_numeral_extra(4) mult_cancel_right2 mult_right_mono)
    have ff2: "r\<^sub>1 \<ge> 0"
      by (smt False assms(1) assms(4) atLeastAtMost_iff f2 mult_not_zero zero_le_mult_iff)
    from ff1 and ff2 have ff3: "r\<^sub>1 \<in> {0..1}"
      by simp
    have ff4: "w\<^sub>2 * (1 - w\<^sub>1) = (1 - r\<^sub>1) * r\<^sub>2"
      using f1 f3 False assms
      by (metis (no_types, hide_lams) add_diff_eq diff_add_eq_diff_diff_swap diff_diff_add
          diff_diff_eq2 eq_iff_diff_eq_0 mult.commute mult.right_neutral right_diff_distrib' right_minus_eq)
    then show ?thesis
      using assms(1-2) f2 ff3 apply (rule_tac pmf_eqI)
      apply (simp add: assms(1-2) f2 ff3 pmf_wplus)
      using assms(3-4) ff4
      by (metis (no_types, hide_lams) add.commute add.left_commute mult.assoc mult.commute)
  qed
qed

subsection \<open> Probabilistic Choice \<close>
text \<open> We use parallel-by-merge in UTP to define the probabilistic choice operator.
The merge predicate is the join of two distributions by their weights.
\<close>

definition prob_merge :: "real \<Rightarrow> (('s, 's prss, 's prss) mrg, 's prss) urel" ("\<^bold>P\<^bold>M\<^bsub>_\<^esub>") where
[upred_defs]: "prob_merge r = \<^U>($prob\<acute> = $0:prob +\<^bsub>\<guillemotleft>r\<guillemotright>\<^esub> $1:prob)"

lemma swap_prob_merge:
  assumes "r \<in> {0..1}"
  shows "swap\<^sub>m ;; \<^bold>P\<^bold>M\<^bsub>r\<^esub> = \<^bold>P\<^bold>M\<^bsub>1 - r\<^esub>"
  by (rel_auto, (metis assms wplus_commute)+)

abbreviation prob_des_merge :: "real \<Rightarrow> (('s des, 's prss des, 's prss des) mrg, 's prss des) urel" ("\<^bold>P\<^bold>D\<^bold>M\<^bsub>_\<^esub>") where
"\<^bold>P\<^bold>D\<^bold>M\<^bsub>r\<^esub> \<equiv> \<^bold>D\<^bold>M(\<^bold>P\<^bold>M\<^bsub>r\<^esub>)"

lemma swap_prob_des_merge:
  assumes "r \<in> {0..1}"
  shows "swap\<^sub>m ;; \<^bold>P\<^bold>D\<^bold>M\<^bsub>r\<^esub> = \<^bold>P\<^bold>D\<^bold>M\<^bsub>1 - r\<^esub>"
  by (metis assms swap_des_merge swap_prob_merge)

text \<open> The probabilistic choice operator is defined conditionally in order to satisfy unit and zero
laws (@{text "prob_choice_one"} and @{term "prob_choice_zero"}) below. The definition of the operator
follows~\cite[Definition 3.14]{Jifeng2004}. Actually use of @{text "P \<parallel>\<^sup>D\<^bsub>\<^bold>P\<^bold>M\<^bsub>r\<^esub>\<^esub> Q"} directly for (r = 0)
or (r = 1) cannot get the desired result (P or Q) as the precondition of merged designs cannot be
discharged to the precondition of P or Q simply.
\<close>
definition prob_choice :: "'s hrel_pdes \<Rightarrow> real \<Rightarrow> 's hrel_pdes \<Rightarrow> 's hrel_pdes" ("(_ \<oplus>\<^bsub>_\<^esub> _)" [164, 0, 165] 164)
  where [upred_defs]:
"prob_choice P r Q \<equiv>
  if r \<in> {0<..<1}
  then P \<parallel>\<^sup>D\<^bsub>\<^bold>P\<^bold>M\<^bsub>r\<^esub>\<^esub> Q
  else (if r = 0
        then Q
        else (if r = 1
              then P
              else \<top>\<^sub>D))"

(* Use Morgan's logical constant to construct a probabilistic choice of which the weight is a UTP
expression *)
text \<open> The r in @{text "P \<oplus>\<^bsub>r\<^esub> Q"} is a real number (HOL terms). Sometimes, however, we want a
similar operator of which the weight is a UTP expression (therefore it depends on the values of
state variables). For example, @{text "P \<oplus>\<^bsub>U(1/real (\<guillemotleft>N\<guillemotright>-i))\<^esub> Q"} in a uniform selection algorithms
where @{text "&i"} is a state variable. Hence, @{text "(P \<oplus>\<^sub>e\<^bsub>E\<^esub> Q)"} is defined below, which is
inspired by Morgan's logical constant~\cite{Morgan1990a}.
\<close>
definition prob_choice_r :: "('a, 'a) rel_pdes \<Rightarrow> (real, 'a) uexpr \<Rightarrow> ('a, 'a) rel_pdes \<Rightarrow> ('a, 'a) rel_pdes"
  ("(_ \<oplus>\<^sub>e\<^bsub>_\<^esub> _)" [164, 0, 165] 164)
where [upred_defs]:
"prob_choice_r P E Q \<equiv> (con\<^sub>D R \<bullet> (II\<^sub>D \<triangleleft> U(\<guillemotleft>R\<guillemotright> = E) \<triangleright>\<^sub>D \<bottom>\<^sub>D) ;; (P \<oplus>\<^bsub>R\<^esub> Q))"

lemma prob_choice_commute: "r \<in> {0..1} \<Longrightarrow> P \<oplus>\<^bsub>r\<^esub> Q = Q \<oplus>\<^bsub>1 - r\<^esub> P"
  by (simp add: prob_choice_def swap_prob_des_merge[THEN sym], metis par_by_merge_commute_swap)

lemma prob_choice_one:
  "P \<oplus>\<^bsub>1\<^esub> Q = P"
  by (simp add: prob_choice_def)

lemma prob_choice_zero:
  "P \<oplus>\<^bsub>0\<^esub> Q = Q"
  by (simp add: prob_choice_def)

lemma prob_choice_r:
  "r \<in> {0<..<1} \<Longrightarrow> P \<oplus>\<^bsub>r\<^esub> Q = P \<parallel>\<^sup>D\<^bsub>\<^bold>P\<^bold>M\<^bsub>r\<^esub>\<^esub> Q"
  by (simp add: prob_choice_def)

lemma prob_choice_inf_simp:
  "(\<Sqinter> r \<in> {0<..<1} \<bullet> (P \<oplus>\<^bsub>r\<^esub> Q)) = (\<Sqinter> r \<in> {0<..<1} \<bullet>  P \<parallel>\<^sup>D\<^bsub>\<^bold>P\<^bold>M\<^bsub>r\<^esub>\<^esub> Q)"
  using prob_choice_r
  apply (simp add: prob_choice_def)
  by (simp add: UINF_as_Sup_collect image_def)

text \<open> @{text "inf_is_exists"} helps to establish the fact that our theorem regarding nondeterminism
~\cite[Sect. 8]{Woodcock2019} is the same as He's~\cite[Theorem 3.10]{Jifeng2004}.
\<close>
lemma inf_is_exists:
  "(\<Sqinter> r \<in> {0<..<1} \<bullet>  (p \<turnstile>\<^sub>n P) \<parallel>\<^sup>D\<^bsub>\<^bold>P\<^bold>M\<^bsub>r\<^esub>\<^esub> (q \<turnstile>\<^sub>n Q))
 = (\<^bold>\<exists> r \<in> \<^U>({0<..<1}) \<bullet>  (p \<turnstile>\<^sub>n P) \<parallel>\<^sup>D\<^bsub>\<^bold>P\<^bold>M\<^bsub>r\<^esub>\<^esub> (q \<turnstile>\<^sub>n Q))"
  by (pred_auto)

subsection \<open> Kleisli Lifting and Sequential Composition \<close>
utp_lit_vars
text \<open> The Kleisli lifting operator maps a probabilistic design (@{text "p \<turnstile>\<^sub>n R"}) into a ``lifted''
design that maps from
@{text "prob"} to @{text "prob"}. Therefore, one probabilistic design can be composed sequentially with
another lifted design. The precondition of the definition specifies that all states of the initial
distribution satisfy the predicate @{text "p"}. The postcondition specifies that there exists a
function @{text "Q"}, that maps states to distributions, such that
\begin{itemize}
  \item for any state $s$, if its probability in the initial distribution is larger than 0, then
R(s, Q(s)) must be held;
  \item any state $ss$ in final distribution @{text "$prob\<acute>"} is equal to summation of all paths
from any state $t$ in its initial distribution to $ss$ via @{text "Q t"}.
\end{itemize}

\begin{figure*}\centering
    \includegraphics[width=.45\textwidth]{images/klift.pdf}
    \caption{\label{fig:klift}Illustration of Kleisli lifting}
\end{figure*}

Figure~\ref{fig:klift} illustrates the lifting operation, provided that there are four states in
the state space. The blue states in @{text "$prob"} denotes their initial probabilities are larger
than 0, and the red states in @{text "$prob\<acute>"} denotes their final probabilities are larger than 0.
Q is defined as
  \[\{(s_1, Q(s_1)), (s_2, Q(s_2)), (s_4, Q(s_4))\}\]
and the relation between
$s_i$ and $Q(s_i)$ is established by R. In addition, the probability of $s_1$ in $Q(s_1)$ is larger
than 0, that of $s_1$ and $s_3$ in $Q(s_2)$, and that of $s_3$ and $s_4$ in $Q(s_4)$. Finally, the
finally distribution is given below.
\begin{align*}
    & prob'(s_1) = prob(s_1) * Q(s_1)(s_1) + prob(s_2) * Q(s_2)(s_1) & \\
    & prob'(s_3) = prob(s_2) * Q(s_2)(s_3) + prob(s_4) * Q(s_4)(s_3) & \\
    & prob'(s_4) = prob(s_2) * Q(s_2)(s_4) + prob(s_4) * Q(s_4)(s_4) &
\end{align*}
\<close>
(*
\<up>(p\<turnstile>R) = ((prob(p) = 1) \<turnstile>
  (\<exists>Q\<in>('s\<Rightarrow>'s pmf) .
  (\<forall>s. prob(s) > 0 \<Rightarrow> R(s, Q s))  (* "prob(s) > 0" has different types as "R(s, Q s)", so \<Rightarrow> doesn't work *)
  \<and>
  (prob' =\<lambda>s. (\<Sum>t. prob(t)* (Q(t) s))))))
*)
definition kleisli_lift2:: "'a upred \<Rightarrow> ('a, 'a prss) urel \<Rightarrow> ('a prss, 'a prss) rel_des"
  where "kleisli_lift2 p R =
  ( \<^U>((\<Sum>\<^sub>a i\<in>\<lbrakk>p\<rbrakk>\<^sub>p. $prob i) = 1)
    \<turnstile>\<^sub>r
  (\<^bold>\<exists> Q \<bullet> (
      (\<^bold>\<forall>ss \<bullet> U(($prob\<acute> ss) = (\<Sum>\<^sub>a t. (($prob t) * (pmf (Q t) ss))))) \<and>
      (\<^bold>\<forall>s \<bullet> (\<not>(U($prob $\<^bold>v\<acute> > 0 \<and> $\<^bold>v\<acute> = s) ;;
            ((((\<not>R) ;; (\<^bold>\<forall>t \<bullet> U(($prob t) = (pmf (Q s) t)))))))
    ))
  )))"

named_theorems kleisli_lift

text \<open> Alternatively, we can define the lifting operator as a normal design, instead of a design in
previous definition. \<close>
definition kleisli_lift2':: "'a upred \<Rightarrow> ('a, 'a prss) urel \<Rightarrow> ('a prss, 'a prss) rel_des" where
[kleisli_lift]: "kleisli_lift2' p R =
  ( \<^U>((\<Sum>\<^sub>a i\<in>\<lbrakk>p\<rbrakk>\<^sub>p. &prob i) = 1)
    \<turnstile>\<^sub>n
  (\<^bold>\<exists> Q \<bullet> (
      (\<^bold>\<forall>ss \<bullet> U(($prob\<acute> ss) = (\<Sum>\<^sub>a t. (($prob t) * (pmf (Q t) ss))))) \<and>
      (\<^bold>\<forall>s \<bullet> (\<not>(U($prob $\<^bold>v\<acute> > 0 \<and> $\<^bold>v\<acute> = s) ;;
            ((\<not>R) ;; (\<^bold>\<forall>t \<bullet> U(($prob t) = (pmf (Q s) t)))))
    ))
  )))"

(*
print_theorems
find_theorems "kleisli_lift2'"
*)

text \<open> Two definitions actually are equal. \<close>
lemma kleisli_lift2_eq: "kleisli_lift2' p R = kleisli_lift2 p R"
  apply (simp add: kleisli_lift2_def)
  apply (simp add: utp_prob_des.kleisli_lift2'_def)
  by (rel_auto)

utp_expr_vars

text \<open> Then the lifting operator @{text "\<up>"} is defined upon @{text "kleisli_lift2"}.
\<close>
definition kleisli_lift ("\<up>") where
  "kleisli_lift P = kleisli_lift2 (\<lfloor>pre\<^sub>D(P)\<rfloor>\<^sub><) (pre\<^sub>D(P) \<and> post\<^sub>D(P))"

text \<open> The alternative definition of the lifting operator @{text "\<up>"} is based on
@{text "kleisli_lift2'"}.
\<close>
lemma kleisli_lift_alt_def:
  "kleisli_lift P = kleisli_lift2' (\<lfloor>pre\<^sub>D(P)\<rfloor>\<^sub><) (pre\<^sub>D(P) \<and> post\<^sub>D(P))"
  by (simp add: kleisli_lift_def kleisli_lift2_eq)

text \<open>
Sequential composition of two probabilistic designs (P and Q) is composition of P with the lifted Q
through the Kleisli lifting operator.
\<close>
abbreviation pseqr :: "('b, 'b) rel_pdes \<Rightarrow> ('b, 'b) rel_pdes \<Rightarrow> ('b, 'b) rel_pdes" (infix ";;\<^sub>p" 60)
  where  "pseqr P Q \<equiv> (P ;; (\<up> Q))"

text \<open>
@{text "II\<^sub>p"} is the identity of sequence of probabilistic designs.
\<close>
abbreviation skip_p ("II\<^sub>p") where
  "skip_p \<equiv> \<K>(II\<^sub>D)"

text \<open>
The top of probabilistic designs is still the top of designs.
\<close>
abbreviation falsep :: "('b, 'b) rel_pdes" ("false\<^sub>p") where
"falsep \<equiv> false"

end


section \<open>Introduction\<close>

theory Equivalence_Relation_Enumeration
  imports "HOL-Library.Sublist" "HOL-Library.Disjoint_Sets"
    "Card_Equiv_Relations.Card_Equiv_Relations"
begin

text \<open>As mentioned in the abstract the enumeration algorithm relies on the bijection between
restricted growth functions (RGFs) of length @{term "n"} and the equivalence relations on
@{term "{..<n}"}, where the bijection is the operation that forms the equivalence kernels of an RGF.
The method is being dicussed, for example, by~\<^cite>\<open>"hutchinson1963" and "milne1977"\<close> or
\<^cite>\<open>\<open>\textsection 1.5\<close> in "white1986"\<close>.

An enumeration algorithm for RGFs is less convoluted than one for equivalence relations or
partitions and the representation has the advantage that checking whether a pair of elements are
equivalent can be done by performing two list lookup operations.

After a few preliminary results in the following section, Section~\ref{sec:enum_rgf} introduces the
enumeration algorithm for RGFs and shows that the function enumerates all of them (for the given
length) without repetition. Section~\ref{sec:bij_kernel} shows that the operation of forming the
equivalence kernel is a bijection and concludes with the correctness of the entire algorithm. In
Section~\ref{sec:app} an interesting application is being discussed, where the enumeration of
partitions is applied within a proof. Section~\ref{sec:add} contains a few additional results,
such as the fact that the length of the enumerated list is a Bell number. The latter result relies
on the formalization of the cardinality of equivalence relations by
Bulwahn~\<^cite>\<open>"Card_Equiv_Relations-AFP"\<close>.\<close>

section \<open>Preliminary Results\<close>

text \<open>This section contains a few preliminary results used in the proofs below.\<close>

lemma length_filter:"length (filter p xs) = sum_list (map (\<lambda>x. of_bool ( p x)) xs)"
  by (induct xs, simp_all)

lemma count_list_expand:"count_list xs x = length (filter ((=) x) xs)"
  by (induct xs, simp_all)

text \<open>An induction schema (similar to @{thm [source] list_induct2} and @{thm [source] rev_induct})
for two lists of equal length, where induction step is shown appending elements at the end.\<close>

lemma list_induct_2_rev[consumes 1, case_names Nil Cons]:
  assumes "length x = length y"
  assumes "P [] []"
  assumes "\<And>x xs y ys. length xs = length ys \<Longrightarrow> P xs ys \<Longrightarrow> P (xs@[x]) (ys@[y])"
  shows "P x y"
  using assms(1)
proof (induct "length x" arbitrary: x y)
  case 0
  then show ?case using assms(2) by simp
next
  case (Suc n)
  obtain x1 x2 where a:"x = x1@[x2]" and c:"length x1 = n"
    by (metis Suc(2) append_butlast_last_id length_append_singleton
        length_greater_0_conv nat.inject zero_less_Suc)

  obtain y1 y2 where b:"y = y1@[y2]" and d:"length y1 = n"
    by (metis Suc(2,3) append_butlast_last_id length_append_singleton
        length_greater_0_conv nat.inject zero_less_Suc)

  have "P x1 y1" using c d Suc by simp
  hence "P (x1@[x2]) (y1@[y2])" using assms(3) c d by simp
  thus ?case using a b by simp
qed

text \<open>If all but one value of a sum is zero then it can be evaluated on the remaining point:\<close>

lemma sum_collapse:
  fixes f :: "'a \<Rightarrow> 'b::{comm_monoid_add}"
  assumes "finite A"
  assumes "z \<in> A"
  assumes "\<And>y. y \<in> A \<Longrightarrow> y \<noteq> z \<Longrightarrow> f y = 0"
  shows "sum f A = f z"
  using sum.union_disjoint[where A="A-{z}" and B="{z}" and g="f"]
  by (simp add: assms sum.insert_if)

text \<open>Number of occurrences of elements in lists is preserved under injective maps.\<close>

lemma count_list_inj_map:
  assumes "inj_on f (set x)"
  assumes "y \<in> set x"
  shows "count_list (map f x) (f y) = count_list x y"
  using assms by (induction x, simp_all, fastforce)

text \<open>A relation cannot be an equivalence relation on two distinct sets.\<close>

lemma equiv_on_unique:
  assumes "equiv A p"
  assumes "equiv B p"
  shows "A = B"
  by (meson assms equalityI equiv_class_eq_iff subsetI)

text \<open>The restriction of an equivalence relation is itself an equivalence relation.\<close>

lemma equiv_subset:
  assumes "B \<subseteq> A"
  assumes "equiv A p"
  shows "equiv B (Restr p B)"
proof -
  have "refl_on B (Restr p B)" using assms by (simp add:refl_on_def equiv_def, blast)
  moreover have "sym (Restr p B)" using assms by (simp add:sym_def equiv_def)
  moreover have "trans (Restr p B)"
    using assms by (simp add:trans_def equiv_def, blast)
  ultimately show ?thesis by (simp add:equiv_def)
qed

section \<open>Enumerating Restricted Growth Functions\label{sec:enum_rgf}\<close>

fun rgf_limit :: "nat list \<Rightarrow> nat"
  where
    "rgf_limit [] = 0" |
    "rgf_limit (x#xs) = max (x+1) (rgf_limit xs)"

lemma rgf_limit_snoc: "rgf_limit (x@[y]) = max (y+1) (rgf_limit x)"
  by (induction x, simp_all)

lemma rgf_limit_ge: "y \<in> set xs \<Longrightarrow> y < rgf_limit xs"
  by (induction xs, simp_all, metis lessI max_less_iff_conj not_less_eq)

definition rgf :: "nat list \<Rightarrow> bool"
  where "rgf x = (\<forall>ys y. prefix (ys@[y]) x \<longrightarrow> y \<le> rgf_limit ys)"

text \<open>The function @{term "rgf_limit"} returns the smallest natural number larger than all list
elements, it is the largest allowed value following @{term "xs"} for restricted growth functions.
The definition @{term "rgf"} is the predicate capturing the notion.\<close>

fun enum_rgfs :: "nat \<Rightarrow> (nat list) list"
  where
    "enum_rgfs 0 = [[]]" |
    "enum_rgfs (Suc n) = [(x@[y]). x \<leftarrow> enum_rgfs n,  y \<leftarrow> [0..<rgf_limit x+1]]"

text \<open>The function @{term "enum_rgfs n"} returns all RGFs of length @{term "n"} without repetition.
The fact is verified in the three lemmas at the end of this section.\<close>

lemma rgf_snoc:
  "rgf (xs@[x]) \<longleftrightarrow> rgf xs \<and> x < rgf_limit xs + 1"
  unfolding rgf_def by (rule order_antisym, (simp add:less_Suc_eq_le)+)

lemma rgf_imp_initial_segment:
  "rgf xs \<Longrightarrow> set xs = {..<rgf_limit xs}"
proof (induction xs rule:rev_induct)
  case Nil
  then show ?case by simp
next
  case (snoc x xs)
  have c:"rgf xs" using snoc(2) rgf_snoc by simp
  hence a:"set xs = {..<rgf_limit xs}" using snoc(1) by simp
  have b: "x \<le> rgf_limit xs" using snoc(2) rgf_snoc c by simp
  have "set (xs@[x]) = insert x {..<rgf_limit xs}"
    using a by simp
  also have "... = {..<max (x+1) (rgf_limit xs)}" using b
    by (cases "x < rgf_limit xs", simp add:insert_absorb, simp add:lessThan_Suc)
  also have "... = {..<rgf_limit (xs@[x])}"
    using rgf_limit_snoc by simp
  finally show ?case by simp
qed

lemma enum_rgfs_returns_rgfs:
  assumes "x \<in> set (enum_rgfs n)"
  shows "rgf x"
  using assms
proof (induction n arbitrary: x)
  case 0
  then show ?case by (simp add:rgf_def)
next
  case (Suc n)
  obtain x1 x2 where
    x_def:"x = x1@[x2]" "x2 < rgf_limit x1 + 1" "x1 \<in> set (enum_rgfs n)"
    using Suc by (simp add:image_iff, force)
  have a:"rgf x1" using Suc x_def by blast
  thus ?case using x_def by (simp add:rgf_snoc)
qed

lemma enum_rgfs_len:
  assumes "x \<in> set (enum_rgfs n)"
  shows "length x = n"
  using assms by (induction n arbitrary: x, simp_all, fastforce)

lemma equiv_rels_enum:
  assumes "rgf x"
  shows "count_list (enum_rgfs (length x)) x = 1"
  using assms
proof (induction x rule:rev_induct)
  case Nil
  then show ?case by simp
next
  case (snoc x xs)
  have b:"rgf xs" using snoc(2) rgf_def by simp
  hence "x < rgf_limit xs + 1" using rgf_snoc snoc by blast
  hence a:"card ({0..<rgf_limit xs + 1} \<inter> {x}) = 1" by force
  have "1 = count_list (enum_rgfs (length xs)) xs" using snoc b by simp
  also have "... = (\<Sum>r1\<leftarrow>enum_rgfs (length xs). of_bool (xs = r1) *
      card ({0..<rgf_limit xs + 1} \<inter> {x}))"
    using a by (simp add:length_concat filter_concat count_list_expand length_filter)
  also have "... = (\<Sum>r1\<leftarrow>enum_rgfs (length xs). of_bool (xs = r1) *
      card ({0..<rgf_limit r1 + 1} \<inter> {x}))"
    by (metis (mono_tags, opaque_lifting) mult_eq_0_iff of_bool_eq_0_iff)
  also have "... = (\<Sum>r1\<leftarrow>enum_rgfs (length xs). of_bool (xs = r1) *
      (\<Sum>r2\<leftarrow>[0..<rgf_limit r1 + 1]. of_bool (x = r2)))"
    by (simp add:interv_sum_list_conv_sum_set_nat del:One_nat_def)
  also have "... = length (filter ((=) (xs@[x])) (enum_rgfs (length (xs@[x]))))"
    by (simp add:length_concat filter_concat length_filter comp_def
        of_bool_conj sum_list_const_mult del:upt_Suc)
  also have "... = count_list (enum_rgfs (length (xs@[x]))) (xs@[x])"
    by (simp add:count_list_expand length_filter del:enum_rgfs.simps)
  finally show ?case by presburger
qed

section \<open>Enumerating Equivalence Relations\label{sec:bij_kernel}\<close>

text \<open>The following definition returns the equivalence relation induced by a list, for example, by
a restricted growth function.\<close>

definition kernel_of :: "'a list \<Rightarrow> nat rel"
  where "kernel_of xs = {(i,j). i < length xs \<and> j < length xs \<and> xs ! i = xs ! j}"

text \<open>Using that the enumeration function for equivalence relations on @{term "{..<n}"} is
straight-forward to define:\<close>

definition equiv_rels where "equiv_rels n = map kernel_of (enum_rgfs n)"

text \<open>The following lemma shows that the image of
@{term "kernel_of"} is indeed an equivalence relation:\<close>

lemma kernel_of_equiv: "equiv {..<length xs} (kernel_of xs)"
proof -
  have "kernel_of xs \<subseteq> {..<length xs} \<times> {..<length xs}"
    by (rule subsetI, simp add:kernel_of_def mem_Times_iff case_prod_beta)
  thus ?thesis by (simp add:equiv_def refl_on_def sym_def trans_def kernel_of_def)
qed

lemma kernel_of_eq_len:
  assumes "kernel_of x = kernel_of y"
  shows "length x = length y"
proof -
  have "{..<length x} = {..<length y}"
    by (metis kernel_of_equiv equiv_on_unique assms)
  thus ?thesis by simp
qed

lemma kernel_of_eq:
  "(kernel_of x = kernel_of y) \<longleftrightarrow>
  (length x = length y \<and> (\<forall>j < length x. \<forall>i < j. (x ! i = x ! j) = (y ! i = y ! j)))"
proof (cases "length x = length y")
  case True
  have "(kernel_of x = kernel_of y)  \<longleftrightarrow>
    (\<forall>j < length x. \<forall>i < length x. (x ! i = x ! j) = (y ! i = y ! j))"
    unfolding set_eq_iff kernel_of_def using True by (simp, blast)
  also have "... \<longleftrightarrow> (\<forall>j < length x. \<forall>i < j. (x ! i = x ! j) = (y ! i = y ! j))"
    by (metis (no_types, lifting) linorder_cases order.strict_trans)
  finally show ?thesis using True by simp
next
  case False
  then show ?thesis using kernel_of_eq_len by blast
qed

lemma kernel_of_snoc:
  "kernel_of (xs) = Restr (kernel_of (xs@[x]))  {..<length xs}"
  by (simp add:kernel_of_def nth_append set_eq_iff)

lemma kernel_of_inj_on_rgfs_aux:
  assumes "length x = length y"
  assumes "rgf x"
  assumes "rgf y"
  assumes "kernel_of x = kernel_of y"
  shows "x = y"
  using assms
proof (induct x y rule: list_induct_2_rev)
  case Nil
  then show ?case by simp
next
  case (Cons x xs y ys)
  have a:"kernel_of xs = kernel_of ys"
    using Cons(1,5) kernel_of_snoc by metis
  have d:"rgf xs" "rgf ys" using Cons rgf_def by auto
  hence b:"xs = ys" using Cons(2) a by auto
  have "\<And>i. i < length xs \<Longrightarrow> (xs ! i = x) = (ys ! i = y)"
  proof -
    fix i
    assume i_l:"i < length xs"
    have "(xs ! i = x) \<longleftrightarrow> (i,length xs) \<in> kernel_of (xs@[x])" using i_l
      by (simp add:kernel_of_def less_Suc_eq nth_append)
    also have "... \<longleftrightarrow> (i,length xs) \<in> kernel_of (ys@[y])"
      using Cons(5) by simp
    also have "... \<longleftrightarrow> (ys ! i= y)" using  i_l Cons(1)
      by (simp add:kernel_of_def less_Suc_eq nth_append)
    finally show "(xs ! i = x) = (ys ! i = y)" by simp
  qed
  hence c:"(x \<in> set xs \<longrightarrow> x = y) \<and> (x \<notin> set xs \<longrightarrow> y \<notin> set ys)"
    by (metis b in_set_conv_nth)
  have x_bound:"x < rgf_limit xs + 1"
    using Cons(3) rgf_snoc d by simp
  have y_bound:"y < rgf_limit ys + 1"
    using Cons(4) rgf_snoc d by simp
  have "x = y" using b c d rgf_imp_initial_segment Cons x_bound y_bound
    apply (cases "x < rgf_limit xs", simp)
    by (cases "y < rgf_limit ys", simp+)
  then show ?case using b by simp
qed

lemma kernel_of_inj_on_rgfs:
  "inj_on kernel_of {x. rgf x}"
  by (rule inj_onI, simp, metis kernel_of_eq_len kernel_of_inj_on_rgfs_aux)

text \<open>Applying an injective map to a list preserves the induced relation:\<close>

lemma kernel_of_under_inj_map:
  assumes "inj_on f (set x)"
  shows "kernel_of x = kernel_of (map f x)"
proof -
  have "\<And>i j. i < length x \<Longrightarrow> j < length x
    \<Longrightarrow> (map f x) ! i = (map f x) ! j \<Longrightarrow> x ! i = x ! j"
    using assms by (simp add: inj_on_eq_iff)
  thus ?thesis unfolding kernel_of_def by fastforce
qed

lemma all_rels_are_kernels:
  assumes "equiv {..<n} p"
  shows "\<exists>(x :: nat set list). kernel_of x = p \<and> length x = n"
proof -
  define r where "r = map (\<lambda>k. p``{k}) [0..<n]"

  have "\<And> u v. (u,v) \<in> kernel_of r \<longleftrightarrow> (u,v) \<in> p"
  proof -
    fix u v :: nat
    have "(u,v) \<in> kernel_of r \<longleftrightarrow> ((u,v) \<in> {..<n}\<times>{..<n} \<and> p``{u} = p``{v})"
      unfolding kernel_of_def r_def by auto
    also have "... \<longleftrightarrow> (u,v) \<in> p" by (metis assms equiv_class_eq_iff mem_Sigma_iff)
    finally show "(u,v) \<in> kernel_of r \<longleftrightarrow> (u,v) \<in> p" by simp
  qed
  hence "kernel_of r = p" by auto
  moreover have "length r = n" using r_def by simp
  ultimately show ?thesis by auto
qed

text \<open>For any list there is always an injective map on its set, such that its image is an RGF.\<close>

lemma map_list_to_rgf:
  "\<exists>f. inj_on f (set x) \<and> rgf (map f x)"
proof (induction "length x" arbitrary: x)
  case 0
  then show ?case by (simp add:rgf_def)
next
  case (Suc n)
  obtain x1 x2 where x_def: "x = x1@[x2]" and l_x1: "length x1 = n"
    by (metis append_butlast_last_id length_append_singleton Suc(2)
        length_greater_0_conv nat.inject zero_less_Suc)
  obtain f where inj_f: "inj_on f (set x1)" and pc_f: "rgf (map f x1)"
    using Suc(1) l_x1 by blast
  show ?case
  proof (cases "x2 \<in> set x1")
    case True
    have a:"set x = set x1" using x_def True by auto
    hence b:"inj_on f (set x)" using inj_f by auto

    have "f x2 < rgf_limit (map f x1)" using rgf_limit_ge True by auto
    hence "rgf (map f x)"
      by (simp add:x_def rgf_snoc pc_f)
    then show ?thesis using b by blast
  next
    case False
    define f' where "f' = (\<lambda>y. if y \<in> set x1 then f y else rgf_limit (map f x1))"
    have "inj_on f' (set x1)" using f'_def inj_f by (simp add: inj_on_def)
    moreover have "rgf_limit (map f x1) \<notin> set (map f x1)"
      using rgf_limit_ge by blast
    hence "f' x2 \<notin> f' ` set x1" using False by (simp add:f'_def)
    ultimately have "inj_on f' (insert x2 (set x1))" using False by simp
    hence a:"inj_on f' (set x)" using False x_def by simp

    have b:"map f x1 = map f' x1" using f'_def by simp

    have c:"f' x2 < Suc (rgf_limit (map f x1))" by (simp add:f'_def False)
    have "rgf (map f' x)" by (simp add:x_def b[symmetric] rgf_snoc pc_f c)
    then show ?thesis using a by blast
  qed
qed

text \<open>For any relation there is a corresponding RGF:\<close>

lemma rgf_exists:
  assumes "equiv {..<n} r"
  shows "\<exists>x. rgf x \<and> length x = n \<and> kernel_of x = r"
proof -
  obtain y :: "nat set list" where a:"kernel_of y = r" "length y = n"
    using all_rels_are_kernels assms by blast
  then obtain f where b:"inj_on f (set y)" "rgf (map f y)"
    using map_list_to_rgf by blast
  have "kernel_of (map f y) = r"
    using kernel_of_under_inj_map a b by blast
  moreover have "length (map f y) = n" using a by simp
  ultimately show ?thesis
    using b by blast
qed

text \<open>These are the main result of this entry: The function @{term "equiv_rels n"} enumerates the
equivalence relations on @{term "{..<n}"} without repetition.\<close>

theorem equiv_rels_set:
  assumes "x \<in> set (equiv_rels n)"
  shows "equiv {..<n} x"
  using assms equiv_rels_def kernel_of_equiv enum_rgfs_len by auto

theorem equiv_rels:
  assumes "equiv {..<n} r"
  shows "count_list (equiv_rels n) r = 1"
proof -
  obtain y where y_def: "rgf y" "length y = n" "kernel_of y = r"
    using rgf_exists assms by blast

  have a: "\<And>x. x \<in> set (enum_rgfs n) \<Longrightarrow> (kernel_of y = kernel_of x) = (y=x)"
    using enum_rgfs_returns_rgfs y_def(1,2) enum_rgfs_len inj_onD[OF kernel_of_inj_on_rgfs]
    by auto

  have "count_list (equiv_rels n) r =
    length (filter (\<lambda>x. r = kernel_of x) (enum_rgfs n))"
    by (simp add:equiv_rels_def count_list_expand length_filter comp_def)
  also have "... = length (filter (\<lambda>x. kernel_of y = kernel_of x) (enum_rgfs n))"
    using y_def(3) by simp
  also have "... = length (filter (\<lambda>x. y = x) (enum_rgfs n))"
    using a by (simp cong:filter_cong)
  also have "... = count_list (enum_rgfs n) y"
    by (simp add:count_list_expand length_filter)
  also have "... = 1"
    using equiv_rels_enum y_def(1,2) by auto
  finally show ?thesis by simp
qed

text \<open>A corollary of the previous theorem is that the sum of the indicator function for a relation
over @{term "equiv_rels n"} is always one.\<close>

corollary equiv_rels_2:
  assumes "n = length xs"
  shows "(\<Sum>x\<leftarrow>equiv_rels n. of_bool (kernel_of xs = x)) = (1 :: 'a :: {semiring_1})"
proof -
  have "length (filter (\<lambda>x. kernel_of xs = x) (equiv_rels (length xs))) = 1"
    using equiv_rels[OF kernel_of_equiv[where xs="xs"]] assms by (simp add:count_list_expand)
  thus ?thesis
    using assms by (simp add:of_bool_def sum_list_map_filter'[symmetric] sum_list_triv)
qed

section \<open>Example Application\label{sec:app}\<close>

text \<open>In this section, I wanted to discuss an interesting application within the context of
a proof in Isabelle. This is motivated by a real-world example \<^cite>\<open>\<open>\textsection 2.2\<close> in "alon1999"\<close>,
where a function in a 4-times iterated sum could only be reduced by splitting it according to the
equivalence relation formed by the indices. The notepad below illustrates how this can be done
(in the case of 3 index variables).\<close>

notepad
begin \<^marker>\<open>tag visible\<close>
  fix f :: "nat \<times> nat \<times> nat  \<Rightarrow> nat"
  fix I :: "nat set"
  assume a:"finite I"

  text \<open>To be able to break down such a sum by partitions let us introduce the function $P$
    which is defined to be sum of an indicator function over all possible equivalence relations
    its argument can form:\<close>

  define P :: "nat list \<Rightarrow> nat"
    where "P = (\<lambda>xs. (\<Sum>x \<leftarrow> equiv_rels (length xs). of_bool (kernel_of xs = x) ))"

  text \<open>Note that its value is always one, hence we can introduce it in an algebraic equation easily:\<close>

  have P_one: "\<And>xs. P xs = 1"
    by (simp add: P_def equiv_rels_2)

  note unfold_equiv_rels = P_def equiv_rels_def numeral_eq_Suc kernel_of_eq
    neq_commute All_less_Suc comp_def

  define r where "r = (\<Sum>i \<in> I. (\<Sum>j \<in> I. (\<Sum>k \<in> I.  f (i,j,k))))"

  text \<open>As a first step, we just introduce the factor @{term "P [i,j,k]"}.\<close>

  have "r = (\<Sum>i \<in> I. (\<Sum>j \<in> I. (\<Sum>k \<in> I.  f (i,j,k) * P [i,j,k])))"
    by (simp add:P_one r_def cong:sum.cong)

  text \<open>By expanding the definition of P and distributing, the sum can be expanded into 5 sums
    each representing a distinct equivalence relation formed by the indices.\<close>

  also have "... =
    (\<Sum>i\<in>I. f (i, i, i)) +
    (\<Sum>i\<in>I. \<Sum>j\<in>I. f (i, i, j) * of_bool (i \<noteq> j)) +
    (\<Sum>i\<in>I. \<Sum>j\<in>I. f (i, j, i) * of_bool (i \<noteq> j)) +
    (\<Sum>i\<in>I. \<Sum>j\<in>I. f (i, j, j) * of_bool (i \<noteq> j)) +
    (\<Sum>i\<in>I. \<Sum>j\<in>I. \<Sum>k\<in>I. f (i, j, k) * of_bool (j \<noteq> k \<and> i \<noteq> k \<and> i \<noteq> j))"
    (is "_ = ?rhs")
    by (simp add:unfold_equiv_rels sum.distrib distrib_left sum_collapse[OF a])
  finally have "r = ?rhs" by simp
end

section \<open>Additional Results\label{sec:add}\<close>

text \<open>If two lists induce the same equivalence relation, then there is a bijection between the sets
that preserves the multiplicities of its elements.\<close>

lemma kernel_of_eq_imp_bij:
  assumes "kernel_of x = kernel_of y"
  shows "\<exists>f. bij_betw f (set x) (set y) \<and>
    (\<forall>z \<in> set x. count_list x z = count_list y (f z))"
proof -
  obtain x' where x'_def: "inj_on x' (set x)" "rgf (map x' x)"
    using map_list_to_rgf by blast
  obtain y' where y'_def: "inj_on y' (set y)" "rgf (map y' y)"
    using map_list_to_rgf by blast

  have "kernel_of (map x' x) = kernel_of (map y' y)"
    using assms x'_def(1) y'_def(1)
    by (simp add: kernel_of_under_inj_map[symmetric])
  hence b:"map x' x = map y' y"
    using inj_onD[OF kernel_of_inj_on_rgfs] x'_def(2) y'_def(2) length_map by simp
  hence f: "x' ` set x = y' ` set y"
    by (metis list.set_map)
  define f where "f = the_inv_into (set y) y' \<circ> x'"
  have g:"\<And>z. z \<in> set x \<Longrightarrow> count_list x z = count_list y (f z)"
  proof -
    fix z
    assume a:"z \<in> set x"
    have e: "x' z \<in> y' ` set y"
      by (metis a b imageI image_set)
    have c: "the_inv_into (set y) y' (x' z) \<in> set y"
      using e the_inv_into_into[OF y'_def(1)] by simp
    have d: "(y' (the_inv_into (set y) y' (x' z))) = x' z"
      using e f_the_inv_into_f y'_def(1) by force

    have "count_list x z = count_list (map x' x) (x' z)"
      using a x'_def by (simp add: count_list_inj_map)
    also have "... = count_list (map y' y) (x' z)"
      by (simp add:b)
    also have "... = count_list (map y' y) (y' (the_inv_into (set y) y' (x' z)))"
      by (simp add:d)
    also have "... = count_list y (the_inv_into (set y) y' (x' z))"
      using c count_list_inj_map[OF y'_def(1)] by simp
    also have "... = count_list y (f z)" by (simp add:f_def)
    finally show "count_list x z = count_list y (f z)" by simp
  qed

  have "bij_betw x' (set x) (x' ` set x)"
    using x'_def(1) bij_betw_imageI by auto
  moreover have "bij_betw (the_inv_into (set y) y') (y' ` set y) (set y)"
    using bij_betw_the_inv_into[OF bij_betw_imageI] y'_def(1) by auto
  hence "bij_betw (the_inv_into (set y) y') (x' ` set x) (set y)"
    using f by simp
  ultimately have "bij_betw f (set x) (set y)"
    using bij_betw_trans f_def by blast
  thus ?thesis using g by blast
qed

text \<open>As expected the length of @{term "equiv_rels n"} is the $n$-th Bell number.\<close>

lemma len_equiv_rels: "length (equiv_rels n) = Bell n"
proof -
  have a:"finite {p. equiv {..<n} p}"
    by (simp add: finite_equiv)
  have b: "set (equiv_rels n)  \<subseteq>  {p. equiv {..<n} p}"
    using equiv_rels_set by blast
  have "length (equiv_rels n) =
    (\<Sum>x \<in> {p. equiv {..<n} p}. count_list (equiv_rels n) x)"
    using a b by (simp add:sum_count_set)
  also have "... = card {p. equiv {..<n} p}"
    by (simp add: equiv_rels)
  also have "... = Bell (card {..<n})"
    using card_equiv_rel_eq_Bell by blast
  also have "... = Bell n" by simp
  finally show ?thesis by simp
qed

text \<open>Instead of forming an equivalence relation from a list, it is also possible to induce a
partition from it:\<close>

definition induced_par :: "'a list \<Rightarrow> nat set set"  where
  "induced_par xs = (\<lambda>k. {i. i < length xs \<and> xs ! i = k}) ` (set xs)"

text \<open>The following lemma verifies the commutative diagram, i.e.,
@{term "induced_par xs"} is the same partition as the quotient of @{term "{..<length xs}"} over
the corresponding equivalence relation.\<close>

lemma quotient_of_kernel_is_induced_par:
  "{..<length xs} // (kernel_of xs) = (induced_par xs)"
proof (rule set_eqI)
  fix x
  have "x \<in> {..<length xs} // (kernel_of xs) \<longleftrightarrow>
    (\<exists>i < length xs. x = {j. j < length xs \<and> xs ! i = xs ! j})"
    unfolding quotient_def kernel_of_def by blast
  also have "... \<longleftrightarrow> (\<exists>y \<in> set xs. x = {j. j < length xs \<and> y = xs ! j})"
    unfolding in_set_conv_nth Bex_def by (rule order_antisym, force+)
  also have "... \<longleftrightarrow> (x \<in> induced_par xs)"
    unfolding induced_par_def by auto
  finally show "x \<in> {..<length xs} // (kernel_of xs) \<longleftrightarrow> (x \<in> induced_par xs)"
    by simp
qed

end


theory TreeSimp
imports "$HIPSTER_HOME/IsaHipster"

begin

datatype 'a Tree =
  Leaf 'a
  | Node "'a Tree""'a Tree"

fun mirror :: "'a Tree => 'a Tree"
where
  "mirror (Leaf x) = Leaf x"
| "mirror (Node l r) = Node (mirror r) (mirror l)"

fun tmap :: "('a => 'b) => 'a Tree => 'b Tree"
where
  "tmap f (Leaf x) = Leaf (f x)"
| "tmap f (Node l r) = Node (tmap f l) (tmap f r)"

ML\<open>Hipster_Explore_Simp.explore  @{context} ["TreeSimp.tmap", "TreeSimp.mirror"];\<close>
lemma lemma_a [thy_expl]: "mirror (tmap x2 y2) = tmap x2 (mirror y2)"
by (tactic \<open>Hipster_Tacs.induct_simp_metis @{context} @{thms TreeSimp.tmap.simps TreeSimp.mirror.simps thy_expl}\<close>)

lemma lemma_aa [thy_expl]: "mirror (mirror x2) = x2"
by (tactic \<open>Hipster_Tacs.induct_simp_metis @{context} @{thms TreeSimp.tmap.simps TreeSimp.mirror.simps thy_expl}\<close>)

fun rigthmost :: "'a Tree \<Rightarrow> 'a"
where
  "rigthmost (Leaf x) = x"
|  "rigthmost (Node l r) = rigthmost r"

fun leftmost :: "'a Tree \<Rightarrow> 'a"
where
  "leftmost (Leaf x) = x"
|  "leftmost (Node l r) = leftmost l"

ML\<open>Hipster_Explore_Simp.explore  @{context} ["TreeSimp.mirror","TreeSimp.tmap", "TreeSimp.rigthmost", "TreeSimp.leftmost"];\<close>
lemma lemma_ab [thy_expl]: "leftmost (mirror x2) = rigthmost x2"
by (tactic \<open>Hipster_Tacs.induct_simp_metis @{context} @{thms TreeSimp.mirror.simps TreeSimp.tmap.simps TreeSimp.rigthmost.simps TreeSimp.leftmost.simps thy_expl}\<close>)

lemma lemma_ac [thy_expl]: "rigthmost (mirror x2) = leftmost x2"
by (tactic \<open>Hipster_Tacs.induct_simp_metis @{context} @{thms TreeSimp.mirror.simps TreeSimp.tmap.simps TreeSimp.rigthmost.simps TreeSimp.leftmost.simps thy_expl}\<close>)

fun flat_tree :: "'a Tree => 'a list"
where
  "flat_tree (Leaf x) = Cons x []"
| "flat_tree (Node l r) = (flat_tree l) @ (flat_tree r)"

ML\<open>Hipster_Explore_Simp.explore  @{context} ["TreeSimp.flat_tree", "TreeSimp.mirror", "TreeSimp.tmap", "TreeSimp.leftmost", "TreeSimp.rigthmost","List.rev", "List.map", "List.hd", "List.append"];\<close>

lemma lemma_ad [thy_expl]: "flat_tree (tmap x2 y2) = map x2 (flat_tree y2)"
by (tactic \<open>Hipster_Tacs.induct_simp_metis @{context} @{thms TreeSimp.flat_tree.simps TreeSimp.mirror.simps TreeSimp.tmap.simps TreeSimp.leftmost.simps TreeSimp.rigthmost.simps List.rev.simps List.append.simps thy_expl}\<close>)

lemma lemma_ae [thy_expl]: "map x2 (rev xs2) = rev (map x2 xs2)"
by (tactic \<open>Hipster_Tacs.induct_simp_metis @{context} @{thms TreeSimp.flat_tree.simps TreeSimp.mirror.simps TreeSimp.tmap.simps TreeSimp.leftmost.simps TreeSimp.rigthmost.simps List.rev.simps List.append.simps thy_expl}\<close>)

lemma lemma_af [thy_expl]: "flat_tree (mirror x2) = rev (flat_tree x2)"
by (tactic \<open>Hipster_Tacs.induct_simp_metis @{context} @{thms TreeSimp.flat_tree.simps TreeSimp.mirror.simps TreeSimp.tmap.simps TreeSimp.leftmost.simps TreeSimp.rigthmost.simps List.rev.simps List.append.simps thy_expl}\<close>)

lemma lemma_ag [thy_expl]: "hd (xs2 @ xs2) = hd xs2"
by (tactic \<open>Hipster_Tacs.induct_simp_metis @{context} @{thms TreeSimp.flat_tree.simps TreeSimp.mirror.simps TreeSimp.tmap.simps TreeSimp.leftmost.simps TreeSimp.rigthmost.simps List.rev.simps List.append.simps thy_expl}\<close>)

lemma unknown [thy_expl]: "hd (flat_tree x) = leftmost x"
apply (tactic \<open>Hipster_Explore.explore_goal @{context} ["TreeSimp.flat_tree", "TreeSimp.leftmost"]\<close>)
by (hipster_induct_simp_metis flat_tree.simps)
oops

end