(*
  File:     Random_BSTs.thy
  Author:   Manuel Eberl <manuel@pruvisto.org>

  Expected shape of random Binary Search Trees
*)
section \<open>Expected shape of random Binary Search Trees\<close>
theory Random_BSTs
  imports
    Complex_Main
    "HOL-Probability.Random_Permutations"
    "HOL-Data_Structures.Tree_Set"
    Quick_Sort_Cost.Quick_Sort_Average_Case
begin

(* TODO: Hide this in the proper place *)
hide_const (open) Tree_Set.insert

subsection \<open>Auxiliary lemmas\<close>

(* TODO: Move? *)
lemma linorder_on_linorder_class [intro]:
  "linorder_on UNIV {(x, y). x \<le> (y :: 'a :: linorder)}"
  by (auto simp: linorder_on_def refl_on_def antisym_def trans_def total_on_def)

lemma Nil_in_permutations_of_set_iff [simp]: "[] \<in> permutations_of_set A \<longleftrightarrow> A = {}"
  by (auto simp: permutations_of_set_def)

lemma max_power_distrib_right:
  fixes a :: "'a :: linordered_semidom"
  shows "a > 1 \<Longrightarrow> max (a ^ b) (a ^ c) = a ^ max b c"
  by (auto simp: max_def)

lemma set_tree_empty_iff [simp]: "set_tree t = {} \<longleftrightarrow> t = Leaf"
  by (cases t) auto

lemma card_set_tree_bst: "bst t \<Longrightarrow> card (set_tree t) = size t"
proof (induction t)
  case (Node l x r)
  have "set_tree \<langle>l, x, r\<rangle> = insert x (set_tree l \<union> set_tree r)" by simp
  also from Node.prems have "card \<dots> = Suc (card (set_tree l \<union> set_tree r))"
    by (intro card_insert_disjoint) auto
  also from Node have "card (set_tree l \<union> set_tree r) = size l + size r"
    by (subst card_Un_disjoint) force+
  finally show ?case by simp
qed simp_all

lemma pair_pmf_cong:
  "p = p' \<Longrightarrow> q = q' \<Longrightarrow> pair_pmf p q = pair_pmf p' q'"
  by simp

lemma expectation_add_pair_pmf:
  fixes f :: "'a \<Rightarrow> 'c::{banach, second_countable_topology}"
  assumes "finite (set_pmf p)" and "finite (set_pmf q)"
  shows "measure_pmf.expectation (pair_pmf p q) (\<lambda>(x,y). f x + g y) =
           measure_pmf.expectation p f + measure_pmf.expectation q g"
proof -
  have "measure_pmf.expectation (pair_pmf p q) (\<lambda>(x,y). f x + g y) =
          measure_pmf.expectation (pair_pmf p q) (\<lambda>z. f (fst z) + g (snd z))"
    by (simp add: case_prod_unfold)
  also have "\<dots> = measure_pmf.expectation (pair_pmf p q) (\<lambda>z. f (fst z)) +
                  measure_pmf.expectation (pair_pmf p q) (\<lambda>z. g (snd z))"
    by (intro Bochner_Integration.integral_add integrable_measure_pmf_finite) (auto intro: assms)
  also have "measure_pmf.expectation (pair_pmf p q) (\<lambda>z. f (fst z)) =
               measure_pmf.expectation (map_pmf fst (pair_pmf p q)) f" by simp
  also have "map_pmf fst (pair_pmf p q) = p" by (rule map_fst_pair_pmf)
  also have "measure_pmf.expectation (pair_pmf p q) (\<lambda>z. g (snd z)) =
               measure_pmf.expectation (map_pmf snd (pair_pmf p q)) g" by simp
  also have "map_pmf snd (pair_pmf p q) = q" by (rule map_snd_pair_pmf)
  finally show ?thesis .
qed

subsection \<open>Creating a BST from a list\<close>

text \<open>
  The following recursive function creates a binary search tree from a given list of
  elements by inserting them into an initially empty BST from left to right. We will prove
  that this is the case later, but the recursive definition has the advantage of giving us
  a useful induction rule, so we chose that definition and prove the alternative definitions later.

  This recursion, which already almost looks like QuickSort, will be key in analysing the
  shape distributions of random BSTs.
\<close>
fun bst_of_list :: "'a :: linorder list \<Rightarrow> 'a tree" where
  "bst_of_list [] = Leaf"
| "bst_of_list (x # xs) =
     Node (bst_of_list [y \<leftarrow> xs. y < x]) x (bst_of_list [y \<leftarrow> xs. y > x])"

lemma bst_of_list_eq_Leaf_iff [simp]: "bst_of_list xs = Leaf \<longleftrightarrow> xs = []"
  by (induction xs) auto

lemma bst_of_list_snoc [simp]:
  "bst_of_list (xs @ [y]) = Tree_Set.insert y (bst_of_list xs)"
  by (induction xs rule: bst_of_list.induct) auto

lemma bst_of_list_append:
  "bst_of_list (xs @ ys) = fold Tree_Set.insert ys (bst_of_list xs)"
proof (induction ys arbitrary: xs)
  case (Cons y ys)
  have "bst_of_list (xs @ (y # ys)) = bst_of_list ((xs @ [y]) @ ys)" by simp
  also have "\<dots> = fold Tree_Set.insert ys (bst_of_list (xs @ [y]))"
    by (rule Cons.IH)
  finally show ?case by simp
qed simp_all

text \<open>
  The following now shows that the recursive function indeed corresponds to the
  notion of inserting the elements from the list from left to right.
\<close>
lemma bst_of_list_altdef: "bst_of_list xs = fold Tree_Set.insert xs Leaf"
  using bst_of_list_append[of "[]" xs] by simp

lemma size_bst_insert: "x \<notin> set_tree t \<Longrightarrow> size (Tree_Set.insert x t) = Suc (size t)"
  by (induction t) auto

lemma set_bst_insert [simp]: "set_tree (Tree_Set.insert x t) = insert x (set_tree t)"
  by (induction t) auto

lemma set_bst_of_list [simp]: "set_tree (bst_of_list xs) = set xs"
  by (induction xs rule: rev_induct) simp_all

lemma size_bst_of_list_distinct [simp]:
  assumes "distinct xs"
  shows   "size (bst_of_list xs) = length xs"
  using assms by (induction xs rule: rev_induct) (auto simp: size_bst_insert)

lemma strict_mono_on_imp_less_iff:
  assumes "strict_mono_on A f" "x \<in> A" "y \<in> A"
  shows   "f x < (f y :: 'b :: linorder) \<longleftrightarrow> x < (y :: 'a :: linorder)"
  using assms by (cases x y rule: linorder_cases; force simp: strict_mono_on_def)+

lemma bst_of_list_map:
  fixes f :: "'a :: linorder \<Rightarrow> 'b :: linorder"
  assumes "strict_mono_on A f" "set xs \<subseteq> A"
  shows   "bst_of_list (map f xs) = map_tree f (bst_of_list xs)"
  using assms
proof (induction xs rule: bst_of_list.induct)
  case (2 x xs)
  have "[xa\<leftarrow>xs . f xa < f x] = [xa\<leftarrow>xs . xa < x]" and "[xa\<leftarrow>xs . f xa > f x] = [xa\<leftarrow>xs . xa > x]"
    using "2.prems" by (auto simp: strict_mono_on_imp_less_iff intro!: filter_cong)
  with 2 show ?case by (auto simp: filter_map o_def)
qed auto

subsection \<open>Random BSTs\<close>

text \<open>
  Analogously to the previous section, we can now view the concept of a random BST
  (i.\,e.\ a BST obtained by inserting a given set of elements in random order) in two
  different ways.

  We again start with the recursive variant:
\<close>
function random_bst :: "'a :: linorder set \<Rightarrow> 'a tree pmf" where
  "random_bst A =
     (if \<not>finite A \<or> A = {} then
        return_pmf Leaf
      else do {
        x \<leftarrow> pmf_of_set A;
        l \<leftarrow> random_bst {y \<in> A. y < x};
        r \<leftarrow> random_bst {y \<in> A. y > x};
        return_pmf (Node l x r)
     })"
  by auto
termination by (relation finite_psubset) auto

declare random_bst.simps [simp del]

lemma random_bst_empty [simp]: "random_bst {} = return_pmf Leaf"
  by (simp add: random_bst.simps)

lemma set_pmf_random_permutation [simp]:
  "finite A \<Longrightarrow> set_pmf (pmf_of_set (permutations_of_set A)) = {xs. distinct xs \<and> set xs = A}"
  by (subst set_pmf_of_set) (auto dest: permutations_of_setD)

text \<open>
  The alternative characterisation is the more intuitive one where we simply pick a
  random permutation of the set elements uniformly at random and insert them into an empty
  tree from left to right:
\<close>
lemma random_bst_altdef:
  assumes "finite A"
  shows   "random_bst A = map_pmf bst_of_list (pmf_of_set (permutations_of_set A))"
using assms
proof (induction A rule: finite_psubset_induct)
  case (psubset A)
  define L R where "L = (\<lambda>x. {y\<in>A. y < x})" and "R = (\<lambda>x. {y\<in>A. y > x})"
  {
    fix x assume x: "x \<in> A"
    hence *: "L x \<subset> A" "R x \<subset> A" by (auto simp: L_def R_def)
    note this [THEN psubset.IH]
  } note IH = this

  show ?case
  proof (cases "A = {}")
    case False
    note A = \<open>finite A\<close> \<open>A \<noteq> {}\<close>
    have "random_bst A =
            do {
              x \<leftarrow> pmf_of_set A;
              (l, r) \<leftarrow> pair_pmf (random_bst (L x)) (random_bst (R x));
              return_pmf (Node l x r)
            }" using A unfolding pair_pmf_def L_def R_def
      by (subst random_bst.simps) (simp add: bind_return_pmf bind_assoc_pmf)
    also have "\<dots> = do {
                      x \<leftarrow> pmf_of_set A;
                      (l, r) \<leftarrow> pair_pmf
                        (map_pmf bst_of_list (pmf_of_set (permutations_of_set (L x))))
                        (map_pmf bst_of_list (pmf_of_set (permutations_of_set (R x))));
                      return_pmf (Node l x r)
                    }"
     using A by (intro bind_pmf_cong refl) (simp_all add: IH)
    also have "\<dots> = do {
                     x \<leftarrow> pmf_of_set A;
                     (ls, rs) \<leftarrow> pair_pmf (pmf_of_set (permutations_of_set (L x)))
                                          (pmf_of_set (permutations_of_set (R x)));
                     return_pmf (Node (bst_of_list ls) x (bst_of_list rs))
                   }" unfolding map_pair [symmetric]
      by (simp add: map_pmf_def case_prod_unfold bind_return_pmf bind_assoc_pmf)
    also have "L = (\<lambda>x. {y \<in> A - {x}. y \<le> x})" by (auto simp: L_def)
    also have "R = (\<lambda>x. {y \<in> A - {x}. \<not>y \<le> x})" by (auto simp: R_def)
    also have "do {
                 x \<leftarrow> pmf_of_set A;
                 (ls, rs) \<leftarrow> pair_pmf (pmf_of_set (permutations_of_set {y \<in> A - {x}. y \<le> x}))
                                      (pmf_of_set (permutations_of_set {y \<in> A - {x}. \<not>y \<le> x}));
                 return_pmf (Node (bst_of_list ls) x (bst_of_list rs))
               } =
               do {
                 x \<leftarrow> pmf_of_set A;
                 (ls, rs) \<leftarrow> map_pmf (partition (\<lambda>y. y \<le> x))
                               (pmf_of_set (permutations_of_set (A - {x})));
                 return_pmf (Node (bst_of_list ls) x (bst_of_list rs))
               }" using \<open>finite A\<close>
      by (intro bind_pmf_cong refl partition_random_permutations [symmetric]) auto
    also have "\<dots> = do {
                      x \<leftarrow> pmf_of_set A;
                      (ls, rs) \<leftarrow> map_pmf (\<lambda>xs. ([y\<leftarrow>xs. y < x], [y\<leftarrow>xs. y > x]))
                                    (pmf_of_set (permutations_of_set (A - {x})));
                      return_pmf (Node (bst_of_list ls) x (bst_of_list rs))
                    }" using A
      by (intro bind_pmf_cong refl map_pmf_cong)
         (auto intro!: filter_cong dest: permutations_of_setD simp: order.strict_iff_order)
    also have "\<dots> = map_pmf bst_of_list (pmf_of_set (permutations_of_set A))"
      using A by (subst random_permutation_of_set[of A])
                 (auto simp: map_pmf_def bind_return_pmf o_def bind_assoc_pmf not_le)
    finally show ?thesis .
  qed (simp_all add: pmf_of_set_singleton)
qed

lemma finite_set_random_bst [simp, intro]:
  "finite A \<Longrightarrow> finite (set_pmf (random_bst A))"
  by (simp add: random_bst_altdef)

lemma random_bst_code [code]:
  "random_bst (set xs) = map_pmf bst_of_list (pmf_of_set (permutations_of_set (set xs)))"
  by (rule random_bst_altdef) simp_all

lemma random_bst_singleton [simp]: "random_bst {x} = return_pmf (Node Leaf x Leaf)"
  by (simp add: random_bst_altdef pmf_of_set_singleton)

lemma size_random_bst:
  assumes "t \<in> set_pmf (random_bst A)" "finite A"
  shows   "size t = card A"
proof -
  from assms obtain xs where "distinct xs" "A = set xs" "t = bst_of_list xs"
    by (auto simp: random_bst_altdef dest: permutations_of_setD)
  thus ?thesis using \<open>finite A\<close> by (simp add: distinct_card)
qed

lemma random_bst_image:
  assumes "finite A" "strict_mono_on A f"
  shows   "random_bst (f ` A) = map_pmf (map_tree f) (random_bst A)"
proof -
  from assms(2) have inj: "inj_on f A" by (rule strict_mono_on_imp_inj_on)
  with assms have "inj_on (map f) (permutations_of_set A)"
    by (intro inj_on_mapI) auto
  with assms inj have "random_bst (f ` A) =
                         map_pmf (\<lambda>x. bst_of_list (map f x)) (pmf_of_set (permutations_of_set A))"
    by (simp add: random_bst_altdef permutations_of_set_image_inj map_pmf_of_set_inj [symmetric]
                  pmf.map_comp o_def)
  also have "\<dots> = map_pmf (map_tree f) (random_bst A)"
    unfolding random_bst_altdef[OF \<open>finite A\<close>] pmf.map_comp o_def using assms
    by (intro map_pmf_cong refl bst_of_list_map[of A f]) (auto dest: permutations_of_setD)
  finally show ?thesis .
qed

text \<open>
  We can also re-phrase the non-recursive definition using the @{const fold_random_permutation}
  combinator from the HOL-Probability library, which folds over a given set in random order.
\<close>
lemma random_bst_altdef':
  assumes "finite A"
  shows   "random_bst A = fold_random_permutation Tree_Set.insert Leaf A"
proof -
  have "random_bst A = map_pmf bst_of_list (pmf_of_set (permutations_of_set A))"
    using assms by (simp add: random_bst_altdef)
  also have "\<dots> = map_pmf (\<lambda>xs. fold Tree_Set.insert xs Leaf) (pmf_of_set (permutations_of_set A))"
    using assms by (intro map_pmf_cong refl) (auto simp: bst_of_list_altdef)
  also from assms have "\<dots> = fold_random_permutation Tree_Set.insert Leaf A"
    by (simp add: fold_random_permutation_fold)
  finally show ?thesis .
qed

subsection \<open>Expected height\<close>

text \<open>
  For the purposes of the analysis of the expected height, we define the following notion
  of `expected height', which is essentially two to the power of the height (as defined
  by Cormen \textit{et al.}) with a special treatment for the empty tree, which has exponential
  height 0.

  Note that the height defined by Cormen \textit{et al.}\ differs from the @{const height}
  function here in Isabelle in that for them, the height of the empty tree is undefined and
  the height of a singleton tree is 0 etc., whereas in Isabelle, the height of the empty tree is
  0 and the height of a singleton tree is 1.
\<close>
definition eheight :: "'a tree \<Rightarrow> nat" where
  "eheight t = (if t = Leaf then 0 else 2 ^ (height t - 1))"

lemma eheight_Leaf [simp]: "eheight Leaf = 0"
  by (simp add: eheight_def)

lemma eheight_Node_singleton [simp]: "eheight (Node Leaf x Leaf) = 1"
  by (simp add: eheight_def)

lemma eheight_Node:
  "l \<noteq> Leaf \<or> r \<noteq> Leaf \<Longrightarrow> eheight (Node l x r) = 2 * max (eheight l) (eheight r)"
  by (cases l; cases r) (simp_all add: eheight_def max_power_distrib_right)

fun eheight_rbst :: "nat \<Rightarrow> nat pmf" where
  "eheight_rbst 0 = return_pmf 0"
| "eheight_rbst (Suc 0) = return_pmf 1"
| "eheight_rbst (Suc n) =
     do {
       k \<leftarrow> pmf_of_set {..n};
       h1 \<leftarrow> eheight_rbst k;
       h2 \<leftarrow> eheight_rbst (n - k);
       return_pmf (2 * max h1 h2)}"

definition eheight_exp :: "nat \<Rightarrow> real" where
  "eheight_exp n = measure_pmf.expectation (eheight_rbst n) real"

lemma eheight_rbst_reduce:
  assumes "n > 1"
  shows   "eheight_rbst n =
             do {k \<leftarrow> pmf_of_set {..<n}; h1 \<leftarrow> eheight_rbst k; h2 \<leftarrow> eheight_rbst (n - k - 1);
                 return_pmf (2 * max h1 h2)}"
  using assms by (cases n rule: eheight_rbst.cases) (simp_all add: lessThan_Suc_atMost)

lemma Leaf_in_set_random_bst_iff:
  assumes "finite A"
  shows   "Leaf \<in> set_pmf (random_bst A) \<longleftrightarrow> A = {}"
proof
  assume "Leaf \<in> set_pmf (random_bst A)"
  from size_random_bst[OF this] and assms show "A = {}" by auto
qed auto

    hence "map_pmf eheight (random_bst A) =
             do {
               x \<leftarrow> pmf_of_set A;
               l \<leftarrow> random_bst {y \<in> A. y < x};
               r \<leftarrow> random_bst {y \<in> A. y > x};
               return_pmf (eheight (Node l x r))
             }"
      using \<open>finite A\<close> by (subst random_bst.simps) (auto simp: map_bind_pmf)
    also have "\<dots> = do {
                      x \<leftarrow> pmf_of_set A;
                      l \<leftarrow> random_bst {y \<in> A. y < x};
                      r \<leftarrow> random_bst {y \<in> A. y > x};
                      return_pmf (2 * max (eheight l) (eheight r))
                    }"
      using 3 \<open>finite A\<close> exists_other
      by (intro bind_pmf_cong refl, subst eheight_Node)
         (force simp: Leaf_in_set_random_bst_iff not_less nonempty eheight_Node)+
    also have "\<dots> = do {
                      x \<leftarrow> pmf_of_set A;
                      h1 \<leftarrow> map_pmf eheight (random_bst {y \<in> A. y < x});
                      h2 \<leftarrow> map_pmf eheight (random_bst {y \<in> A. y > x});
                      return_pmf (2 * max h1 h2)
                    }"
      by (simp add: bind_map_pmf)
    also have "\<dots> = do {
                      x \<leftarrow> pmf_of_set A;
                      h1 \<leftarrow> eheight_rbst (card {y \<in> A. y < x});
                      h2 \<leftarrow> eheight_rbst (card {y \<in> A. y > x});
                      return_pmf (2 * max h1 h2)
                    }"
      using \<open>A \<noteq> {}\<close> \<open>finite A\<close> by (intro bind_pmf_cong psubset.IH [symmetric] refl) auto
    also have "\<dots> = do {
                      k \<leftarrow> map_pmf rank (pmf_of_set A);
                      h1 \<leftarrow> eheight_rbst k;
                      h2 \<leftarrow> eheight_rbst (card A - k - 1);
                      return_pmf (2 * max h1 h2)
                    }"
      unfolding bind_map_pmf
    proof (intro bind_pmf_cong refl, goal_cases)
      case (1 x)
      have "rank x = card {y\<in>A-{x}. y \<le> x}" by (simp add: rank_def linorder_rank_def)
      also have "{y\<in>A-{x}. y \<le> x} = {y\<in>A. y < x}" by auto
      finally show ?case by simp
    next
      case (2 x)
      have "A - {x} = {y\<in>A-{x}. y \<le> x} \<union> {y\<in>A. y > x}" by auto
      also have "card \<dots> = rank x + card {y\<in>A. y > x}"
        using \<open>finite A\<close> by (subst card_Un_disjoint) (auto simp: rank_def linorder_rank_def)
      finally have "card {y\<in>A. y > x} = card A - rank x - 1"
        using 2 \<open>finite A\<close> \<open>A \<noteq> {}\<close> by simp
      thus ?case by simp
    qed
    also have "map_pmf rank (pmf_of_set A) = pmf_of_set {..<card A}"
      using \<open>A \<noteq> {}\<close> \<open>finite A\<close> unfolding rank_def
      by (intro map_pmf_of_set_bij_betw bij_betw_linorder_rank[of UNIV]) auto
    also have "do {
                 k \<leftarrow> pmf_of_set {..<card A};
                 h1 \<leftarrow> eheight_rbst k;
                 h2 \<leftarrow> eheight_rbst (card A - k - 1);
                 return_pmf (2 * max h1 h2)
               } = eheight_rbst (card A)"
      by (rule eheight_rbst_reduce [symmetric]) fact+
    finally show ?thesis ..
  qed (auto simp: is_singleton_def)
qed

lemma finite_pmf_set_eheight_rbst [simp, intro]: "finite (set_pmf (eheight_rbst n))"
proof -
  have "eheight_rbst n = map_pmf eheight (random_bst {..<n})"
    by (subst eheight_rbst [symmetric]) auto
  also have "finite (set_pmf \<dots>)" by simp
  finally show ?thesis .
qed

lemma eheight_exp_1 [simp]: "eheight_exp (Suc 0) = 1"
  by (simp add: eheight_exp_def lessThan_Suc)

lemma eheight_exp_reduce_bound:
  assumes "n > 1"
  shows   "eheight_exp n \<le> 4 / n * (\<Sum>k<n. eheight_exp k)"
proof -
  have [simp]: "real (max a b) = max (real a) (real b)" for a b
    by (simp add: max_def)
  let ?f = "\<lambda>(h1,h2). max h1 h2"
  let ?p = "\<lambda>k. pair_pmf (eheight_rbst k) (eheight_rbst (n - Suc k))"
  have "eheight_exp n = measure_pmf.expectation (eheight_rbst n) real"
    by (simp add: eheight_exp_def)
  also have "\<dots> = 1 / real n * (\<Sum>k<n. measure_pmf.expectation
                                         (map_pmf (\<lambda>(h1,h2). 2 * max h1 h2) (?p k)) real)"
    (is "_ = _ * ?S") unfolding pair_pmf_def map_bind_pmf
    by (subst eheight_rbst_reduce [OF assms], subst pmf_expectation_bind_pmf_of_set)
       (insert assms, auto simp: sum_divide_distrib divide_simps)
  also have "?S = (\<Sum>k<n. measure_pmf.expectation (map_pmf (\<lambda>x. 2 * x) (map_pmf ?f (?p k))) real)"
    by (simp only: pmf.map_comp o_def case_prod_unfold)
  also have "\<dots> = 2 * (\<Sum>k<n. measure_pmf.expectation (map_pmf ?f (?p k)) real)" (is "_ = _ * ?S'")
    by (subst integral_map_pmf) (simp add: sum_distrib_left)
  also have "?S' = (\<Sum>k<n. measure_pmf.expectation (?p k) (\<lambda>(h1,h2). max (real h1) (real h2)))"
    by (simp add: case_prod_unfold)
  also have "\<dots> \<le> (\<Sum>k<n. measure_pmf.expectation (?p k) (\<lambda>(h1,h2). real h1 + real h2))"
    unfolding integral_map_pmf case_prod_unfold
    by (intro sum_mono Bochner_Integration.integral_mono integrable_measure_pmf_finite) auto
  also have "\<dots> = (\<Sum>k<n. eheight_exp k) + (\<Sum>k<n. eheight_exp (n - Suc k))"
    by (subst expectation_add_pair_pmf) (auto simp: sum.distrib eheight_exp_def)
  also have "(\<Sum>k<n. eheight_exp (n - Suc k)) = (\<Sum>k<n. eheight_exp k)"
    by (intro sum.reindex_bij_witness[of _ "\<lambda>k. n - Suc k" "\<lambda>k. n - Suc k"]) auto
  also have "1 / real n * (2 * (\<dots> + \<dots>)) = 4 / real n * \<dots>" by simp
  finally show ?thesis using assms by (simp_all add: mult_left_mono divide_right_mono)
qed

text \<open>
  We now define the following upper bound on the expected exponential height due to
  Cormen\ \textit{et\ al.}~\<^cite>\<open>"cormen"\<close>:
\<close>
lemma eheight_exp_bound: "eheight_exp n \<le> real ((n + 3) choose 3) / 4"
proof (induction n rule: less_induct)
  case (less n)
  consider "n = 0" | "n = 1" | "n > 1" by force
  thus ?case
  proof cases
    case 3
    hence "eheight_exp n \<le> 4 / n * (\<Sum>k<n. eheight_exp k)"
      by (rule eheight_exp_reduce_bound)
    also have "(\<Sum>k<n. eheight_exp k) \<le> (\<Sum>k<n. real ((k + 3) choose 3) / 4)"
      by (intro sum_mono less.IH) auto
    also have "\<dots> = real (\<Sum>k<n. ((k + 3) choose 3)) / 4"
      by (simp add: sum_divide_distrib)
    also have "(\<Sum>k<n. ((k + 3) choose 3)) = (\<Sum>k\<le>n - 1. ((k + 3) choose 3))"
      using \<open>n > 1\<close> by (intro sum.cong) auto
    also have "\<dots> = ((n + 3) choose 4)"
      using choose_rising_sum(1)[of 3 "n - 1"] and \<open>n > 1\<close> by (simp add: add_ac Suc3_eq_add_3)
    also have "4 / real n * (\<dots> / 4) = real ((n + 3) choose 3) / 4" using \<open>n > 1\<close>
      by (cases n) (simp_all add: binomial_fact fact_numeral divide_simps)
    finally show ?thesis using \<open>n > 1\<close> by (simp add: mult_left_mono divide_right_mono)
  qed (auto simp: eval_nat_numeral)
qed

text \<open>
  We then show that this is indeed an upper bound on the expected exponential height by induction
  over the set of elements. This proof mostly follows that by Cormen\ \textit{et al.}~\<^cite>\<open>"cormen"\<close>,
  and partially an answer on the Computer Science Stack Exchange~\<^cite>\<open>"sofl"\<close>.
\<close>

text \<open>
  Since the function $\uplambda x.\ 2 ^ x$ is convex, we can then easily derive a bound on the
  actual height using Jensen's inequality:
\<close>
definition height_exp_approx :: "nat \<Rightarrow> real" where
  "height_exp_approx n = log 2 (real ((n + 3) choose 3) / 4) + 1"

theorem height_expectation_bound:
  assumes "finite A" "A \<noteq> {}"
  shows   "measure_pmf.expectation (random_bst A) height
             \<le> height_exp_approx (card A)"
proof -
  have "convex_on UNIV ((powr) 2)"
    by (intro convex_on_realI[where f' = "\<lambda>x. ln 2 * 2 powr x"])
       (auto intro!: derivative_eq_intros DERIV_powr simp: powr_def [abs_def])
  hence "2 powr measure_pmf.expectation (random_bst A) (\<lambda>t. real (height t - 1)) \<le>
          measure_pmf.expectation (random_bst A) (\<lambda>t. 2 powr real (height t - 1))"
    using assms
    by (intro measure_pmf.jensens_inequality[where I = UNIV])
       (auto intro!: integrable_measure_pmf_finite)
  also have "(\<lambda>t. 2 powr real (height t - 1)) = (\<lambda>t. 2 ^ (height t - 1))"
    by (simp add: powr_realpow)
  also have "measure_pmf.expectation (random_bst A) (\<lambda>t. 2 ^ (height t - 1)) =
               measure_pmf.expectation (random_bst A) (\<lambda>t. real (eheight t))"
    using assms
    by (intro integral_cong_AE)
       (auto simp: AE_measure_pmf_iff random_bst_altdef eheight_def)
  also have "\<dots> = measure_pmf.expectation (map_pmf eheight (random_bst A)) real"
    by simp
  also have "map_pmf eheight (random_bst A) = eheight_rbst (card A)"
    by (rule eheight_rbst [symmetric]) fact+
  also have "measure_pmf.expectation \<dots> real = eheight_exp (card A)"
    by (simp add: eheight_exp_def)
  also have "\<dots> \<le> real ((card A + 3) choose 3) / 4" by (rule eheight_exp_bound)
  also have "measure_pmf.expectation (random_bst A) (\<lambda>t. real (height t - 1)) =
               measure_pmf.expectation (random_bst A) (\<lambda>t. real (height t) - 1)"
  proof (intro integral_cong_AE AE_pmfI, goal_cases)
    case (3 t)
    with \<open>A \<noteq> {}\<close> and assms show ?case
      by (subst of_nat_diff) (auto simp: Suc_le_eq random_bst_altdef)
  qed auto
  finally have "2 powr measure_pmf.expectation (random_bst A) (\<lambda>t. real (height t) - 1)
                  \<le> real ((card A + 3) choose 3) / 4" .
  hence "log 2 (2 powr measure_pmf.expectation (random_bst A) (\<lambda>t. real (height t) - 1)) \<le>
           log 2 (real ((card A + 3) choose 3) / 4)" (is "?lhs \<le> ?rhs")
    by (subst log_le_cancel_iff) auto
  also have "?lhs = measure_pmf.expectation (random_bst A) (\<lambda>t. real (height t) - 1)"
    by simp
  also have "\<dots> = measure_pmf.expectation (random_bst A) (\<lambda>t. real (height t)) - 1"
    using assms
    by (subst Bochner_Integration.integral_diff) (auto intro!: integrable_measure_pmf_finite)
  finally show ?thesis by (simp add: height_exp_approx_def)
qed

text \<open>
  This upper bound is asymptotically equivalent to $c \ln n$ with
  $c = \frac{3}{\ln 2} \approx 4.328$. This is actually a relatively tight upper bound, since
  the exact asymptotics of the expected height of a random BST is $c \ln n$ with
  $c \approx 4.311$.~\<^cite>\<open>"reed"\<close> However, the proof of these precise asymptotics is very intricate
  and we will therefore be content with the upper bound.

  In particular, we can now show that the expected height is $O(\log n)$.
\<close>
lemma ln_sum_bigo_ln: "(\<lambda>x::real. ln (x + c)) \<in> O(ln)"
proof (rule bigoI_tendsto)
  from eventually_gt_at_top[of "1::real"] show "eventually (\<lambda>x::real. ln x \<noteq> 0) at_top"
    by eventually_elim simp_all
next
  show "((\<lambda>x. ln (x + c) / ln x) \<longlongrightarrow> 1) at_top"
  proof (rule lhospital_at_top_at_top)
    show "eventually (\<lambda>x. ((\<lambda>x. ln (x + c)) has_real_derivative inverse (x + c)) (at x)) at_top"
      using eventually_gt_at_top[of "-c"]
      by eventually_elim (auto intro!: derivative_eq_intros simp: field_simps)
    show "eventually (\<lambda>x. ((\<lambda>x. ln x) has_real_derivative inverse x) (at x)) at_top"
      using eventually_gt_at_top[of 0]
      by eventually_elim (auto intro!: derivative_eq_intros simp: field_simps)
    show "((\<lambda>x. inverse (x + c) / inverse x) \<longlongrightarrow> 1) at_top"
    proof (rule Lim_transform_eventually)
      show "eventually (\<lambda>x. inverse (1 + c / x) = inverse (x + c) / inverse x) at_top"
        using eventually_gt_at_top[of "0::real"] eventually_gt_at_top[of "-c"]
        by eventually_elim (simp add: field_simps)
      have "((\<lambda>x. inverse (1 + c / x)) \<longlongrightarrow> inverse (1 + 0)) at_top"
        by (intro tendsto_inverse tendsto_add tendsto_const
              real_tendsto_divide_at_top[OF tendsto_const] filterlim_ident) simp_all
      thus "((\<lambda>x. inverse (1 + c / x)) \<longlongrightarrow> 1) at_top" by simp
    qed
  qed (auto simp: ln_at_top eventually_at_top_not_equal)
qed

corollary height_expectation_bigo: "height_exp_approx \<in> O(ln)"
proof -
  let ?T = "\<lambda>x::real. log 2 (x + 1) + log 2 (x + 2) + log 2 (x + 3) + (1 - log 2 24)"
  have "eventually (\<lambda>n. height_exp_approx n =
          log 2 (real n + 1) + log 2 (real n + 2) + log 2 (real n + 3) + (1 - log 2 24)) at_top"
    (is "eventually (\<lambda>n. _ = ?T n) at_top") using eventually_gt_at_top[of "0::nat"]
  proof eventually_elim
    case (elim n)
    have "height_exp_approx n = log 2 (real (n + 3 choose 3) / 4) + 1"
      by (simp add: height_exp_approx_def log_divide)
    also have "real ((n + 3) choose 3) = real (n + 3) gchoose 3"
      by (simp add: binomial_gbinomial)
    also have "\<dots> / 4 = (real n + 1) * (real n + 2) * (real n + 3) / 24"
      by (simp add: gbinomial_pochhammer' numeral_3_eq_3 pochhammer_Suc add_ac)
    also have "log 2 \<dots> = log 2 (real n + 1) + log 2 (real n + 2) + log 2 (real n + 3) - log 2 24"
      by (simp add: log_divide log_mult)
    finally show ?case by simp
  qed
  hence "height_exp_approx \<in> \<Theta>(?T)" by (rule bigthetaI_cong)
  also have *: "(\<lambda>x. ln (x + c) / ln 2) \<in> O(ln)" for c :: real
    by (subst landau_o.big.cdiv_in_iff') (auto intro!: ln_sum_bigo_ln)
  have "?T \<in> O(\<lambda>n. ln (real n))" unfolding log_def
    by (intro bigo_real_nat_transfer sum_in_bigo ln_sum_bigo_ln *) simp_all
  finally show ?thesis .
qed

subsection \<open>Lookup costs\<close>

text \<open>
  The following function describes the cost incurred when looking up a specific element
  in a specific BST. The cost corresponds to the number of edges traversed in the lookup.
\<close>

primrec lookup_cost :: "'a :: linorder \<Rightarrow> 'a tree \<Rightarrow> nat" where
  "lookup_cost x Leaf = 0"
| "lookup_cost x (Node l y r) =
     (if x = y then 0
      else if x < y then Suc (lookup_cost x l)
      else Suc (lookup_cost x r))"

text \<open>
  Some of the literature defines these costs as 1 in the case that the current node is
  the correct one, i.\,e.\ their costs are our costs plus 1. These alternative costs are
  exactly the number of comparisons performed in the lookup. Our cost function has the
  advantage of precisely summing up to the internal path length and therefore gives us
  slightly nicer results, and since the difference is only a ${}+1$ in the end, this
  variant seemed more reasonable.
\<close>

text \<open>
  It can be shown with a simple induction that The sum of all lookup costs in a tree is the
  internal path length of the tree.
\<close>
theorem sum_lookup_costs:
  fixes t :: "'a :: linorder tree"
  assumes "bst t"
  shows   "(\<Sum>x\<in>set_tree t. lookup_cost x t) = ipl t"
using assms
proof (induction t)
  case (Node l x r)
  from Node.prems
    have disj: "x \<notin> set_tree l" "x \<notin> set_tree r" "set_tree l \<inter> set_tree r = {}" by force+
  have "set_tree (Node l x r) = insert x (set_tree l \<union> set_tree r)" by simp
  also have "(\<Sum>y\<in>\<dots>. lookup_cost y (Node l x r)) = lookup_cost x \<langle>l, x, r\<rangle> +
               (\<Sum>y\<in>set_tree l. lookup_cost y \<langle>l, x, r\<rangle>) + (\<Sum>y\<in>set_tree r. lookup_cost y \<langle>l, x, r\<rangle>)"
    using disj by (simp add: sum.union_disjoint)
  also have "(\<Sum>y\<in>set_tree l. lookup_cost y \<langle>l, x, r\<rangle>) = (\<Sum>y\<in>set_tree l. 1 + lookup_cost y l)"
    using disj and Node by (intro sum.cong refl) auto
  also have "\<dots> = size l + ipl l" using Node
    by (subst sum.distrib) (simp_all add: card_set_tree_bst)
  also have "(\<Sum>y\<in>set_tree r. lookup_cost y \<langle>l, x, r\<rangle>) = (\<Sum>y\<in>set_tree r. 1 + lookup_cost y r)"
    using disj and Node by (intro sum.cong refl) auto
  also have "\<dots> = size r + ipl r" using Node
    by (subst sum.distrib) (simp_all add: card_set_tree_bst)
  finally show ?case by simp
qed simp_all

text \<open>
  This allows us to easily show that the expected cost of looking up a random element in a
  fixed tree is the internal path length divided by the number of elements.
\<close>
theorem expected_lookup_cost:
  assumes "bst t" "t \<noteq> Leaf"
  shows   "measure_pmf.expectation (pmf_of_set (set_tree t)) (\<lambda>x. lookup_cost x t) =
             ipl t / size t"
  using assms by (subst integral_pmf_of_set)
                 (simp_all add: sum_lookup_costs of_nat_sum [symmetric] card_set_tree_bst)

text \<open>
  Therefore, we will now turn to analysing the internal path length of a random BST. This
  then clearly related to the expected lookup costs of a random element in a random BST by
  the above result.
\<close>

subsection \<open>Average Path Length\<close>

text \<open>
  The internal path length satisfies the recursive equation @{thm ipl.simps(2)[of l x r]}.
  This is quite similar to the number of comparisons performed by QuickSort, and indeed, we can
  reduce the internal path length of a random BST to the number of comparisons performed by
  QuickSort on a randomly-ordered list relatively easily:
\<close>
theorem map_pmf_random_bst_eq_rqs_cost:
  assumes "finite A"
  shows   "map_pmf ipl (random_bst A) = rqs_cost (card A)"
using assms
proof (induction A rule: finite_psubset_induct)
  case (psubset A)
  show ?case
  proof (cases "A = {}")
    case False
    note A = \<open>finite A\<close> \<open>A \<noteq> {}\<close>
    define n where "n = card A - 1"
    define rank :: "'a \<Rightarrow> nat" where "rank = linorder_rank {(x,y). x \<le> y} A"
    from A have card: "card A = Suc n" by (cases "card A") (auto simp: n_def)
    from A have "map_pmf ipl (random_bst A) =
                   do {
                     x \<leftarrow> pmf_of_set A;
                     (l,r) \<leftarrow> pair_pmf (random_bst {y \<in> A. y < x}) (random_bst {y \<in> A. y > x});
                     return_pmf (ipl (Node l x r))
                   }"
      by (subst random_bst.simps)
         (simp_all add: pair_pmf_def card map_pmf_def bind_assoc_pmf bind_return_pmf)
    also have "\<dots> = do {
                      x \<leftarrow> pmf_of_set A;
                      (l,r) \<leftarrow> pair_pmf (random_bst {y \<in> A. y < x}) (random_bst {y \<in> A. y > x});
                      return_pmf (n + ipl l + ipl r)
                    }"
    proof (intro bind_pmf_cong refl, clarify, goal_cases)
      case (1 x l r)
      from 1 and A have "n = card (A - {x})" by (simp add: n_def)
      also have "A - {x} = {y\<in>A. y < x} \<union> {y\<in>A. y > x}" by auto
      also have "card \<dots> = card {y\<in>A. y < x} + card {y\<in>A. y > x}"
        using \<open>finite A\<close> by (intro card_Un_disjoint) auto
      also from 1 and A have "card {y\<in>A. y < x} = size l" by (auto dest: size_random_bst)
      also from 1 and A have "card {y\<in>A. y > x} = size r" by (auto dest: size_random_bst)
      finally show ?case by simp
    qed
    also have "\<dots> = do {
                      x \<leftarrow> pmf_of_set A;
                      (l,r) \<leftarrow> pair_pmf (map_pmf ipl (random_bst {y \<in> A. y < x}))
                                        (map_pmf ipl (random_bst {y \<in> A. y > x}));
                      return_pmf (n + l + r)
                    }" by (simp add: map_pair [symmetric] case_prod_unfold bind_map_pmf)
    also have "\<dots> = do {
                      i \<leftarrow> map_pmf rank (pmf_of_set A);
                      (l,r) \<leftarrow> pair_pmf (rqs_cost i) (rqs_cost (n - i));
                      return_pmf (n + l + r)
                    }" (is "_ = bind_pmf _ ?f") unfolding bind_map_pmf
    proof (intro bind_pmf_cong refl pair_pmf_cong, goal_cases)
      case (1 x)
      have "map_pmf ipl (random_bst {y \<in> A. y < x}) = rqs_cost (card {y \<in> A. y < x})"
        using 1 and A by (intro psubset.IH) auto
      also have "{y \<in> A. y < x} = {y \<in> A - {x}. y \<le> x}" by auto
      hence "card {y \<in> A. y < x} = rank x" by (simp add: rank_def linorder_rank_def)
      finally show ?case .
    next
      case (2 x)
      have "map_pmf ipl (random_bst {y \<in> A. y > x}) = rqs_cost (card {y \<in> A. y > x})"
        using 2 and A by (intro psubset.IH) auto
      also have "{y \<in> A. y > x} = A - {x} - {y \<in> A - {x}. y \<le> x}" by auto
      hence "card {y \<in> A. y > x} = card \<dots>" by (simp only:)
      also from 2 and A have "\<dots> = n - rank x"
        by (subst card_Diff_subset) (auto simp: rank_def linorder_rank_def n_def)
      finally show ?case .
    qed
    also from A have "map_pmf rank (pmf_of_set A) = pmf_of_set {..<card A}"
      unfolding rank_def by (intro map_pmf_of_set_bij_betw bij_betw_linorder_rank[of UNIV]) auto
    also have "{..<card A} = {..n}" by (auto simp: card)
    also have "pmf_of_set \<dots> \<bind> ?f = rqs_cost (card A)"
      by (simp add: pair_pmf_def bind_assoc_pmf bind_return_pmf card)
    finally show ?thesis .
  qed simp_all
qed

text \<open>
  In particular, this means that the expected values are the same:
\<close>
corollary expected_ipl_random_bst_eq:
  assumes "finite A"
  shows   "measure_pmf.expectation (random_bst A) ipl = rqs_cost_exp (card A)"
proof -
  have "measure_pmf.expectation (random_bst A) ipl =
          measure_pmf.expectation (map_pmf ipl (random_bst A)) real" by simp
  also from assms have "map_pmf ipl (random_bst A) = rqs_cost (card A)"
    by (rule map_pmf_random_bst_eq_rqs_cost)
  also have "measure_pmf.expectation \<dots> real = rqs_cost_exp (card A)"
    by (rule expectation_rqs_cost)
  finally show ?thesis .
qed

text \<open>
  Therefore, the results about the expected number of comparisons of QuickSort carry over
  to the expected internal path length:
\<close>
corollary expected_ipl_random_bst_eq':
  assumes "finite A"
  shows   "measure_pmf.expectation (random_bst A) ipl =
             2 * real (card A + 1) * harm (card A) - 4 * real (card A)"
  by (simp add: expected_ipl_random_bst_eq rqs_cost_exp_eq assms)

end


header {*Preliminaries*}

theory Measure
imports Sigma_Algebra MonConv
begin

(*We use a modified version of the simple Sigma_Algebra Theory by
Markus Wenzel here,
  which does not need an explicit definition of countable,
  changing the names according to Joe Hurd*)
text {* Now we are already set for the central concept of
  measure. The following definitions are translated as faithfully as possible
  from those in Joe Hurd's thesis \cite{hurd2002}. *}

definition
  measurable:: "'a set set \<Rightarrow> 'b set set \<Rightarrow> ('a \<Rightarrow> 'b) set" where
  "measurable F G = {f. \<forall>g\<in>G. f -` g \<in> F}"

text {*So a function is called $F$-$G$-measurable if and only if the inverse
  image of any set in $G$ is in $F$. $F$ and $G$ are usually the sets of
  measurable sets, the first component of a measure space\footnote{In
  standard mathematical notation, the universe is first in a
  measure space triple, but in our definitions, following Joe Hurd, it is always the
  whole type universe and therefore omitted.}.*}

definition
  measurable_sets:: "('a set set * ('a set \<Rightarrow> real)) \<Rightarrow> 'a set set" where
  "measurable_sets = fst"

definition
  measure:: "('a set set * ('a set \<Rightarrow> real)) \<Rightarrow> ('a set \<Rightarrow> real)" where
  "measure = snd"

text {*The other component is the measure itself. It is a function that
  assigns a nonnegative real number to every measurable set and has
  the property of being
  countably additive for disjoint sets.*}

definition
  positive:: "('a set set * ('a set \<Rightarrow> real)) \<Rightarrow> bool" where
  "positive M \<longleftrightarrow> measure M {} = 0 \<and>
  (\<forall>A. A\<in> measurable_sets M \<longrightarrow> 0 \<le> measure M A)"
  (*Remark: This definition of measure space is not minimal,
  in the sense that the containment of the UNION in the measurable sets
  is implied by the measurable sets being a sigma algebra*)

definition
  countably_additive:: "('a set set * ('a set => real)) => bool" where
  "countably_additive M \<longleftrightarrow> (\<forall>f::(nat => 'a set). range f \<subseteq> measurable_sets M
  \<and> (\<forall>m n. m \<noteq> n \<longrightarrow> f m \<inter> f n = {}) \<and>  (\<Union>i. f i) \<in> measurable_sets M
  \<longrightarrow> (\<lambda>n. measure M (f n)) sums  measure M (\<Union>i. f i))"

text {*This last property deserves some comments. The conclusion is
  usually --- also in the aforementioned source --- phrased as

  @{text "measure M (\<Union>i. f i) = (\<Sum>n. measure M (f n))"}.

  In our formal setting this is unsatisfactory, because the
  sum operator\footnote{Which is merely syntactic sugar for the
  \isa{suminf} functional from the \isa{Series} theory
  \cite{Fleuriot:2000:MNR}.}, like any HOL function, is total, although
  a series obviously need not converge. It is defined using the @{text \<epsilon>} operator, and its
  behavior is unspecified in the diverging case. Hence, the above assertion
  would give no information about the convergence of the series.

  Furthermore, the definition contains redundancy. Assuming that the
  countable union of sets is measurable is unnecessary when the
  measurable sets form a sigma algebra, which is postulated in the
  final definition\footnote{Joe Hurd inherited this practice from a very
  influential probability textbook \cite{Williams.mart}}.
  *}

definition
  measure_space:: "('a set set * ('a set \<Rightarrow> real)) \<Rightarrow> bool" where
  "measure_space M \<longleftrightarrow> sigma_algebra (measurable_sets M) \<and>
  positive M \<and> countably_additive M"

text {*Note that our definition is restricted to finite measure
  spaces --- that is, @{text "measure M UNIV < \<infinity>"} --- since the measure
  must be a real number for any measurable set. In probability, this
  is naturally the case.

  Two important theorems close this section. Both appear in
  Hurd's work as well, but are shown anyway, owing to their central
  role in measure theory. The first one is a mighty tool for proving measurability. It states
  that for a function mapping one sigma algebra into another, it is
  sufficient to be measurable regarding only a generator of the target
  sigma algebra. Formalizing the interesting proof out of Bauer's
  textbook \cite{Bauer} is relatively straightforward using rule
  induction. *}

theorem assumes sig: "sigma_algebra a" and meas: "f \<in> measurable a b" shows
  measurable_lift: "f \<in> measurable a (sigma b)"
proof -
  def Q \<equiv> "{q. f -` q \<in> a}"
  with meas have 1: "b \<subseteq> Q" by (auto simp add: measurable_def)

  { fix x assume "x\<in>sigma b"
    hence "x\<in>Q"
    proof (induct rule: sigma.induct)
      case basic
      from 1 show " \<And>a. a \<in> b \<Longrightarrow> a \<in> Q" ..
    next
      case empty
      from sig have "{}\<in>a"
        by (simp only: sigma_algebra_def)
      thus "{} \<in> Q"
        by (simp add: Q_def)
    next
      case complement
      fix r assume "r \<in> Q"
      then obtain r1 where im: "r1 = f -` r" and a: "r1 \<in> a"
        by (simp add: Q_def)
      with sig have "-r1 \<in> a"
        by (simp only: sigma_algebra_def)
      with im Q_def show "-r \<in> Q"
        by (simp add: vimage_Compl)
    next
      case Union
      fix r assume "\<And>i::nat. r i \<in> Q"
      then obtain r1 where im: "\<And>i. r1 i =  f -` r i" and a: "\<And>i. r1 i \<in>        a"
        by (simp add: Q_def)
      from a sig have "UNION UNIV r1 \<in> a"
        by (auto simp only: sigma_algebra_def)
      with im Q_def show "UNION UNIV r \<in> Q"
        by (auto simp add: vimage_UN)
    qed }

  hence "(sigma b) \<subseteq> Q" ..
  thus "f \<in> measurable a (sigma b)"
    by (auto simp add: measurable_def Q_def)
qed

text {*The case is different for the second theorem. It is only five
  lines in the book (ibid.), but almost 200 in formal text. Precision
  still pays here, gaining a detailed view of a technique that
  is often employed in measure theory --- making a sequence of sets
  disjoint. Moreover, the necessity for the above-mentioned change in the
  definition of countably additive was detected only in the
  formalization of this proof.

  To enable application of the additivity of measures, the following construction
  yields disjoint sets. We skip the justification of the lemmata for
  brevity. *}

primrec mkdisjoint:: "(nat \<Rightarrow> 'a set) \<Rightarrow> (nat \<Rightarrow> 'a set)"
where
  "mkdisjoint A 0 = A 0"
| "mkdisjoint A (Suc n) = A (Suc n) - A n"

lemma mkdisjoint_un:
  assumes up: "\<And>n. A n \<subseteq> A (Suc n)"
  shows "A n = (\<Union>i\<in>{..n}. mkdisjoint A i)"
(*<*)proof (induct n)
  case 0 show ?case by simp
next
  case (Suc n)
  hence "A n = (\<Union>i\<in>{..n}. mkdisjoint A i)" .
  moreover
  have "(\<Union>i\<in>{..(Suc n)}. mkdisjoint A i) = mkdisjoint A (Suc n) \<union>
    (\<Union>i\<in>{..n}. mkdisjoint A i)" by (simp add: atMost_Suc)
  moreover
  have "mkdisjoint A (Suc n) \<union> A n = A (Suc n) \<union> A n" by simp
  moreover
  from up have "\<dots> = A (Suc n)" by auto
  ultimately
  show ?case by simp
qed(*>*)

lemma mkdisjoint_disj:
  assumes up: "\<And>n. A n \<subseteq> A (Suc n)" and ne: "m \<noteq> n"
  shows "mkdisjoint A m \<inter> mkdisjoint A n = {}"
(*<*)proof -
  { fix m1 m2::nat assume less: "m1 < m2"
    hence "0 < m2" by simp
    then obtain n where eq: "m2 = Suc n" by (auto simp add: gr0_conv_Suc)
    with less have less2: "m1 < Suc n" by simp

    {
      fix y assume y: "y \<in> mkdisjoint A m1"
      fix x assume x: "x \<in> mkdisjoint A m2"
      with eq have"x \<notin> A n" by simp
      also from up have "A n = (\<Union>i\<in>{..n}. mkdisjoint A i)"
        by (rule mkdisjoint_un)
      also
      from less2 have "m1 \<in> {..n}" by simp
      hence "mkdisjoint A m1 \<subseteq> (\<Union>i\<in>{..n}. mkdisjoint A i)" by auto
      ultimately
      have "x \<notin> mkdisjoint A m1" by fast
      with y have "y \<noteq> x" by fast
    }
    hence "mkdisjoint A m1 \<inter> mkdisjoint A m2 = {}"
      by (simp add: disjoint_iff_not_equal)
  } hence 1: "\<And>m1 m2. m1 < m2 \<Longrightarrow>  mkdisjoint A m1 \<inter> mkdisjoint A m2 = {}" .

  show ?thesis
  proof (cases "m < n")
    case True
    thus ?thesis by (rule 1)
  next
    case False
    with ne have "n < m" by arith
    hence "mkdisjoint A n \<inter> mkdisjoint A m = {}" by (rule 1)
    thus ?thesis by fast
  qed
qed(*>*)

lemma mkdisjoint_mon_conv:
  assumes mc: "A\<up>B"
  shows "(\<Union>i. mkdisjoint A i) = B"
(*<*)proof
  { fix x assume "x \<in> (\<Union>i. mkdisjoint A i)"
    then obtain i where "x \<in> mkdisjoint A i" by auto
    hence "x \<in> A i" by (cases i) simp_all
    with mc have "x \<in> B" by (auto simp add: set_mon_conv)
  }
  thus "(\<Union>i. mkdisjoint A i) \<subseteq> B" by fast

  { fix x assume "x \<in> B"
    with mc obtain i where "x \<in> A i" by (auto simp add: set_mon_conv)
    also from mc have "\<And>n. A n \<subseteq> A (Suc n)" by (simp only: set_mon_conv)
    hence "A i = (\<Union>r\<in>{..i}. mkdisjoint A r)" by (rule mkdisjoint_un)
    also have "\<dots> \<subseteq> (\<Union>r. mkdisjoint A r)" by auto
    finally have "x \<in> (\<Union>i. mkdisjoint A i)".
  }
  thus "B \<subseteq> (\<Union>i. mkdisjoint A i)" by fast
qed(*>*)

(*This is in Joe Hurd's Thesis (p. 35) as Monotone Convergence theorem. Check the real name \<dots> .
    Also, it's not as strong as it could be,
    but we need no more.*)

text {* Joe Hurd calls the following the Monotone Convergence Theorem,
  though in mathematical literature this name is often reserved for a
  similar fact
  about integrals that we will prove in \ref{nnfis}, which depends on this
  one. The claim made here is that the measures of monotonically convergent sets
  approach the measure of their limit. A strengthened version would
  imply monotone convergence of the measures, but is not needed in the
  development.
  *}

theorem measure_mon_conv:
  assumes ms: "measure_space M" and
  Ams: "\<And>n. A n \<in> measurable_sets M" and AB: "A\<up>B"
  shows "(\<lambda>n. measure M (A n)) ----> measure M B"
proof -

  from AB have up: "\<And>n. A n \<subseteq> A (Suc n)"
    by (simp only: set_mon_conv)

  { fix i
    have "mkdisjoint A i \<in> measurable_sets M"
    proof (cases i)
      case 0 with Ams show ?thesis by simp
    next
      case (Suc i)
      have "A (Suc i) - A i = A (Suc i) \<inter> - A i" by blast
      with Suc ms Ams show ?thesis
        by (auto simp add: measure_space_def sigma_algebra_def sigma_algebra_inter)
    qed
  }
  hence i: "\<And>i. mkdisjoint A i \<in> measurable_sets M" .

  with ms have un: "(\<Union>i. mkdisjoint A i) \<in> measurable_sets M"
    by (simp add: measure_space_def sigma_algebra_def)
  moreover
  from i have range: "range (mkdisjoint A) \<subseteq> measurable_sets M"
    by fast
  moreover
  from up have "\<forall>i j. i \<noteq> j \<longrightarrow>  mkdisjoint A i \<inter> mkdisjoint A j = {}"
    by (simp add: mkdisjoint_disj)
  moreover note ms
  ultimately
  have sums:
    "(\<lambda>i. measure M (mkdisjoint A i)) sums (measure M (\<Union>i. mkdisjoint A i))"
    by (simp add: measure_space_def countably_additive_def)
  hence "(\<Sum>i. measure M (mkdisjoint A i)) = (measure M (\<Union>i. mkdisjoint A i))"
    by (rule sums_unique[THEN sym])

  also
  from sums have "summable (\<lambda>i. measure M (mkdisjoint A i))"
    by (rule sums_summable)

  hence "(\<lambda>n. \<Sum>i<n. measure M (mkdisjoint A i))
    ----> (\<Sum>i. measure M (mkdisjoint A i))"
    by (rule summable_LIMSEQ)

  hence "(\<lambda>n. \<Sum>i<Suc n. measure M (mkdisjoint A i)) ----> (\<Sum>i. measure M (mkdisjoint A i))"
    by (rule LIMSEQ_Suc)

  ultimately have "(\<lambda>n. \<Sum>i<Suc n. measure M (mkdisjoint A i))
    ----> (measure M (\<Union>i. mkdisjoint A i))" by simp

  also
  { fix n
    from up have "A n = (\<Union>i\<in>{..n}. mkdisjoint A i)"
      by (rule mkdisjoint_un)
    hence "measure M (A n) = measure M (\<Union>i\<in>{..n}. mkdisjoint A i)"
      by simp

    also have
      "(\<Union>i\<in>{..n}. mkdisjoint A i) = (\<Union>i. if i\<le>n then mkdisjoint A i else {})"
    proof -
      have "UNIV = {..n} \<union> {n<..}" by auto
      hence "(\<Union>i. if i\<le>n then mkdisjoint A i else {}) =
        (\<Union>i\<in>{..n}. if i\<le>n then mkdisjoint A i else {})
        \<union>  (\<Union>i\<in>{n<..}. if i\<le>n then mkdisjoint A i else {})"
        by (auto split: if_splits)
      moreover
      { have "(\<Union>i\<in>{n<..}. if i\<le>n then mkdisjoint A i else {}) = {}"
          by force }
      hence "\<dots> = (\<Union>i\<in>{..n}. mkdisjoint A i)"
        by auto
      ultimately show
        "(\<Union>i\<in>{..n}. mkdisjoint A i) = (\<Union>i. if i\<le>n then mkdisjoint A i else {})" by simp
    qed

    ultimately have
      "measure M (A n) = measure M (\<Union>i. if i\<le>n then mkdisjoint A i else {})"
      by simp

    also
    from i ms have
      un: "(\<Union>i. if i\<le>n then mkdisjoint A i else {}) \<in> measurable_sets M"
      by (simp add: measure_space_def sigma_algebra_def)
    moreover
    from i ms have
      "range (\<lambda>i. if i\<le>n then mkdisjoint A i else {}) \<subseteq> measurable_sets M"
      by (auto simp add: measure_space_def sigma_algebra_def)
    moreover
    from up have "\<forall>i j. i \<noteq> j \<longrightarrow>
      (if i\<le>n then mkdisjoint A i else {}) \<inter>
      (if j\<le>n then mkdisjoint A j else {}) = {}"
      by (simp add: mkdisjoint_disj)
    moreover note ms
    ultimately have
      "measure M (A n) = (\<Sum>i. measure M (if i \<le> n then mkdisjoint A i else {}))"
      by (simp add: measure_space_def countably_additive_def sums_unique)

    also
    from ms have
      "\<forall>i. (Suc n)\<le>i \<longrightarrow> measure M (if i \<le> n then mkdisjoint A i else {}) = 0"
      by (simp add: measure_space_def positive_def)
    hence "(\<lambda>i. measure M (if i \<le> n then mkdisjoint A i else {})) sums
      (\<Sum>i<Suc n. measure M (if i \<le> n then mkdisjoint A i else {}))"
      by (intro sums_finite) auto
    hence "(\<Sum>i. measure M (if i \<le> n then mkdisjoint A i else {})) =
      (\<Sum>i<Suc n. measure M (if i \<le> n then mkdisjoint A i else {}))"
      by (rule sums_unique[THEN sym])
    also
    have "\<dots> = (\<Sum>i<Suc n. measure M (mkdisjoint A i))"
      by simp
    finally have
      "measure M (A n) = (\<Sum>i<Suc n. measure M (mkdisjoint A i))" .
  }

  ultimately have
    "(\<lambda>n. measure M (A n)) ----> (measure M (\<Union>i. mkdisjoint A i))"
    by simp

  with AB show ?thesis
    by (simp add: mkdisjoint_mon_conv)
qed(*>*)

(*<*)
primrec trivial_series2:: "'a set \<Rightarrow> 'a set \<Rightarrow> (nat \<Rightarrow> 'a set)"
where
  "trivial_series2 a b 0 = a"
| "trivial_series2 a b (Suc n) = (if (n=0) then b else {})"

lemma measure_additive: assumes ms: "measure_space M"
  and disj: "a \<inter> b = {}" and a: "a \<in> measurable_sets M"
  and b:"b \<in> measurable_sets M"
  shows "measure M (a \<union> b) = measure M a + measure M b"
(*<*)proof -
  have "(a \<union> b) = (\<Union>i. trivial_series2 a b i)"
  proof (rule set_eqI)
    fix x
    {
      assume "x \<in> a \<union> b"
      hence "\<exists>i. x \<in> trivial_series2 a b i"
      proof
        assume "x \<in> a"
        hence "x \<in> trivial_series2 a b 0"
          by simp
        thus "\<exists>i. x \<in> trivial_series2 a b i"
          by fast
      next
        assume "x \<in> b"
        hence "x \<in> trivial_series2 a b 1"
          by simp
        thus "\<exists>i. x \<in> trivial_series2 a b i"
          by fast
      qed
    }
    hence "(x \<in> a \<union> b) \<Longrightarrow> (x \<in> (\<Union>i. trivial_series2 a b i))"
      by simp
    also
    {
      assume "x \<in> (\<Union>i. trivial_series2 a b i)"
      then obtain i where x: "x \<in> trivial_series2 a b i"
        by auto
      hence "x \<in> a \<union> b"
      proof (cases i)
        case 0
        with x show ?thesis by simp
      next
        case (Suc n)
        with x show ?thesis
          by (cases n) auto
      qed
    }
    ultimately show "(x \<in> a \<union> b) = (x \<in> (\<Union>i. trivial_series2 a b i))"
      by fast
  qed
  also
  { fix i
    from a b ms have "trivial_series2 a b i \<in> measurable_sets M"
      by (cases i) (auto simp add: measure_space_def sigma_algebra_def)
  }
  hence m1: "range (trivial_series2 a b) \<subseteq> measurable_sets M"
    and m2: "(\<Union>i. trivial_series2 a b i) \<in> measurable_sets M"
    using ms
    by (auto simp add: measure_space_def sigma_algebra_def)

  { fix i j::nat
    assume "i \<noteq> j"
    hence "trivial_series2 a b i \<inter> trivial_series2 a b j = {}"
      using disj
      by (cases i, cases j, auto)(cases j, auto)
  }
  with m1 m2 have "(\<lambda>n. measure M (trivial_series2 a b n)) sums  measure M (\<Union>i. trivial_series2 a b i)"
    using ms
    by (simp add: measure_space_def countably_additive_def)
  moreover
  from ms have "\<forall>m. Suc(Suc 0) \<le> m \<longrightarrow> measure M (trivial_series2 a b m) = 0"
  proof (clarify)
    fix m
    assume "Suc (Suc 0) \<le> m"
    thus "measure M (trivial_series2 a b m) = 0"
      using ms
      by (cases m) (auto simp add: measure_space_def positive_def)
  qed
  hence "(\<lambda>n. measure M (trivial_series2 a b n)) sums (\<Sum>n<Suc(Suc 0). measure M (trivial_series2 a b n))"
    by (intro sums_finite) auto
  moreover
  have "(\<Sum>n=0..<Suc(Suc 0). measure M (trivial_series2 a b n)) =
    measure M a + measure M b"
    by simp
  ultimately
  have "measure M (a \<union> b) = (\<Sum>n. measure M (trivial_series2 a b n))"
    and "Measure.measure M a + Measure.measure M b = (\<Sum>n. measure M (trivial_series2 a b n))"
    by (simp_all add: sums_unique)
  thus ?thesis by simp
qed
(*>*)
end


(*
  File:    Efficient_Discrete_Sqrt.thy
  Author:  Markus Großer, Manuel Eberl

  A reasonably efficient algorithm to compute the square root of a natural number (rounded down)
  and to test if a natural number is a perfect square.
*)
theory Efficient_Discrete_Sqrt
imports
  Complex_Main
  "HOL-Computational_Algebra.Computational_Algebra"
  "HOL-Library.Discrete"
  "HOL-Library.Tree"
  "HOL-Library.IArray"
begin
section \<open>Efficient Algorithms for the Square Root on \<open>\<nat>\<close>\<close>

(*
  TODO: This could perhaps be moved somewhere else. Thre is also probably some overlap
  with Sqrt_Babylonian
*)

subsection \<open>A Discrete Variant of Heron's Algorithm\<close>

text \<open>
  An algorithm for calculating the discrete square root, taken from
  Cohen~\cite{cohen2010algebraic}. This algorithm is essentially a discretised variant of
  Heron's method or Newton's method specialised to the square root function.
\<close>

lemma sqrt_eq_floor_sqrt: "Discrete.sqrt n = nat \<lfloor>sqrt n\<rfloor>"
proof -
  have "real ((nat \<lfloor>sqrt n\<rfloor>)\<^sup>2) = (real (nat \<lfloor>sqrt n\<rfloor>))\<^sup>2"
    by simp
  also have "\<dots> \<le> sqrt (real n) ^ 2"
    by (intro power_mono) auto
  also have "\<dots> = real n" by simp
  finally have "(nat \<lfloor>sqrt n\<rfloor>)\<^sup>2 \<le> n"
    by (simp only: of_nat_le_iff)
  moreover have "n < (Suc (nat \<lfloor>sqrt n\<rfloor>))\<^sup>2" proof -
    have "(1 + \<lfloor>sqrt n\<rfloor>)\<^sup>2 > n"
      using floor_correct[of "sqrt n"] real_le_rsqrt[of "1 + \<lfloor>sqrt n\<rfloor>" n]
        of_int_less_iff[of n "(1 + \<lfloor>sqrt n\<rfloor>)\<^sup>2"] not_le
      by fastforce
    then show ?thesis
      using le_nat_floor[of "Suc (nat \<lfloor>sqrt n\<rfloor>)" "sqrt n"]
        of_nat_le_iff[of "(Suc (nat \<lfloor>sqrt n\<rfloor>))\<^sup>2" n] real_le_rsqrt[of _ n] not_le
      by fastforce
  qed
  ultimately show ?thesis using sqrt_unique by fast
qed

fun newton_sqrt_aux :: "nat \<Rightarrow> nat \<Rightarrow> nat" where
  "newton_sqrt_aux x n =
     (let y = (x + n div x) div 2
      in if y < x then newton_sqrt_aux y n else x)"

declare newton_sqrt_aux.simps [simp del]

lemma newton_sqrt_aux_simps:
  "(x + n div x) div 2 < x \<Longrightarrow> newton_sqrt_aux x n = newton_sqrt_aux ((x + n div x) div 2) n"
  "(x + n div x) div 2 \<ge> x \<Longrightarrow> newton_sqrt_aux x n = x"
  by (subst newton_sqrt_aux.simps; simp add: Let_def)+

lemma heron_step_real: "\<lbrakk>t > 0; n \<ge> 0\<rbrakk> \<Longrightarrow> (t + n/t) / 2 \<ge> sqrt n"
  using arith_geo_mean_sqrt[of t "n/t"] by simp

lemma heron_step_div_eq_floored:
  "(t::nat) > 0 \<Longrightarrow> (t + (n::nat) div t) div 2 = nat \<lfloor>(t + n/t) / 2\<rfloor>"
proof -
  assume "t > 0"
  then have "\<lfloor>(t + n/t) / 2\<rfloor> = \<lfloor>(t*t + n) / (2*t)\<rfloor>"
    by (simp add: mult_divide_mult_cancel_right[of t "t + n/t" 2, symmetric]
        algebra_simps)
  also have "\<dots> = (t*t + n) div (2*t)"
    using floor_divide_of_nat_eq by blast
  also have "\<dots> = (t*t + n) div t div 2"
    by (simp add: Divides.div_mult2_eq mult.commute)
  also have "\<dots> = (t + n div t) div 2"
    by (simp add: \<open>0 < t\<close> power2_eq_square)
  finally show ?thesis by simp
qed

lemma heron_step: "t > 0 \<Longrightarrow> (t + n div t) div 2 \<ge> Discrete.sqrt n"
proof -
  assume "t > 0"
  have "Discrete.sqrt n = nat \<lfloor>sqrt n\<rfloor>" by (rule sqrt_eq_floor_sqrt)
  also have "\<dots> \<le> nat \<lfloor>(t + n/t) / 2\<rfloor>"
    using heron_step_real[of t n] \<open>t > 0\<close> by linarith
  also have "\<dots> = (t + n div t) div 2"
    using heron_step_div_eq_floored[OF \<open>t > 0\<close>] by simp
  finally show ?thesis .
qed

lemma newton_sqrt_aux_correct:
  assumes "x \<ge> Discrete.sqrt n"
  shows   "newton_sqrt_aux x n = Discrete.sqrt n"
  using assms
proof (induction x n rule: newton_sqrt_aux.induct)
  case (1 x n)
  show ?case
  proof (cases "x = Discrete.sqrt n")
    case True
    then have "(x ^ 2) div x \<le> n div x" by (intro div_le_mono) simp_all
    also have "(x ^ 2) div x = x" by (simp add: power2_eq_square)
    finally have "(x + n div x) div 2 \<ge> x" by linarith
    with True show ?thesis by (auto simp: newton_sqrt_aux_simps)
  next
    case False
    with "1.prems" have x_gt_sqrt: "x > Discrete.sqrt n" by auto
    with Discrete.le_sqrt_iff[of x n] have "n < x ^ 2" by simp
    have "x * (n div x) \<le> n" using mult_div_mod_eq[of x n] by linarith
    also have "\<dots> < x ^ 2" using Discrete.le_sqrt_iff[of x n] and x_gt_sqrt by simp
    also have "\<dots> = x * x" by (simp add: power2_eq_square)
    finally have "n div x < x" by (subst (asm) mult_less_cancel1) auto
    then have step_decreasing: "(x + n div x) div 2 < x" by linarith
    with x_gt_sqrt have step_ge_sqrt: "(x + n div x) div 2 \<ge> Discrete.sqrt n"
      by (simp add: heron_step)
    from step_decreasing have "newton_sqrt_aux x n = newton_sqrt_aux ((x + n div x) div 2) n"
      by (simp add: newton_sqrt_aux_simps)
    also have "\<dots> = Discrete.sqrt n"
      by (intro "1.IH" step_decreasing step_ge_sqrt) simp_all
    finally show ?thesis .
  qed
qed

definition newton_sqrt :: "nat \<Rightarrow> nat" where
  "newton_sqrt n = newton_sqrt_aux n n"

declare Discrete.sqrt_code [code del]

theorem Discrete_sqrt_eq_newton_sqrt [code]: "Discrete.sqrt n = newton_sqrt n"
  unfolding newton_sqrt_def by (simp add: newton_sqrt_aux_correct Discrete.sqrt_le)

subsection \<open>Square Testing\<close>

text \<open>
  Next, we implement an algorithm to determine whether a given natural number is a perfect square,
  as described by Cohen~\cite{cohen2010algebraic}. Essentially, the number first determines whether
  the number is a square. Essentially
\<close>

definition q11 :: "nat set"
  where "q11 = {0, 1, 3, 4, 5, 9}"
definition q63 :: "nat set"
  where "q63 = {0, 1, 4, 7, 9, 16, 28, 18, 22, 25, 36, 58, 46, 49, 37, 43}"
definition q64 :: "nat set"
  where "q64 = {0, 1, 4, 9, 16, 17, 25, 36, 33, 49, 41, 57}"
definition q65 :: "nat set"
  where "q65 = {0, 1, 4, 10, 14, 9, 16, 26, 30, 25, 29, 40, 56, 36, 49, 61, 35, 51, 39, 55, 64}"

definition q11_array where
  "q11_array = IArray [True,True,False,True,True,True,False,False,False,True,False]"

definition q63_array where
  "q63_array = IArray [True,True,False,False,True,False,False,True,False,True,False,False,
     False,False,False,False,True,False,True,False,False,False,True,False,False,True,False,
     False,True,False,False,False,False,False,False,False,True,True,False,False,False,False,
     False,True,False,False,True,False,False,True,False,False,False,False,False,False,False,
     False,True,False,False,False,False,False]"

definition q64_array where
  "q64_array = IArray [True,True,False,False,True,False,False,False,False,True,False,False,
     False,False,False,False,True,True,False,False,False,False,False,False,False,True,False,
     False,False,False,False,False,False,True,False,False,True,False,False,False,False,True,
     False,False,False,False,False,False,False,True,False,False,False,False,False,False,
     False,True,False,False,False,False,False,False, False]"

definition q65_array where
  "q65_array = IArray [True,True,False,False,True,False,False,False,False,True,True,False,
     False,False,True,False,True,False,False,False,False,False,False,False,False,True,True,
     False,False,True,True,False,False,False,False,True,True,False,False,True,True,False,
     False,False,False,False,False,False,False,True,False,True,False,False,False,True,True
     ,False,False,False,False,True,False,False,True,False]"

lemma sub_q11_array: "i \<in> {..<11} \<Longrightarrow> IArray.sub q11_array i \<longleftrightarrow> i \<in> q11"
  by (simp add: lessThan_nat_numeral lessThan_Suc q11_def q11_array_def, elim disjE; simp)

lemma sub_q63_array: "i \<in> {..<63} \<Longrightarrow> IArray.sub q63_array i \<longleftrightarrow> i \<in> q63"
  by (simp add: lessThan_nat_numeral lessThan_Suc q63_def q63_array_def, elim disjE; simp)

lemma sub_q64_array: "i \<in> {..<64} \<Longrightarrow> IArray.sub q64_array i \<longleftrightarrow> i \<in> q64"
  by (simp add: lessThan_nat_numeral lessThan_Suc q64_def q64_array_def, elim disjE; simp)

lemma sub_q65_array: "i \<in> {..<65} \<Longrightarrow> IArray.sub q65_array i \<longleftrightarrow> i \<in> q65"
  by (simp add: lessThan_nat_numeral lessThan_Suc q65_def q65_array_def, elim disjE; simp)

lemma in_q11_code: "x mod 11 \<in> q11 \<longleftrightarrow> IArray.sub q11_array (x mod 11)"
  by (subst sub_q11_array) auto

lemma in_q63_code: "x mod 63 \<in> q63 \<longleftrightarrow> IArray.sub q63_array (x mod 63)"
  by (subst sub_q63_array) auto

lemma in_q64_code: "x mod 64 \<in> q64 \<longleftrightarrow> IArray.sub q64_array (x mod 64)"
  by (subst sub_q64_array) auto

lemma in_q65_code: "x mod 65 \<in> q65 \<longleftrightarrow> IArray.sub q65_array (x mod 65)"
  by (subst sub_q65_array) auto

definition square_test :: "nat \<Rightarrow> bool" where
  "square_test n =
    (n mod 64 \<in> q64 \<and> (let r = n mod 45045 in
      r mod 63 \<in> q63 \<and> r mod 65 \<in> q65 \<and> r mod 11 \<in> q11 \<and> n = (Discrete.sqrt n)\<^sup>2))"

lemma square_test_code [code]:
  "square_test n =
    (IArray.sub q64_array (n mod 64) \<and> (let r = n mod 45045 in
           IArray.sub q63_array (r mod 63) \<and>
           IArray.sub q65_array (r mod 65) \<and>
           IArray.sub q11_array (r mod 11) \<and> n = (Discrete.sqrt n)\<^sup>2))"
    using in_q11_code [symmetric] in_q63_code [symmetric]
          in_q64_code [symmetric] in_q65_code [symmetric]
  by (simp add: Let_def square_test_def)

lemma square_mod_lower: "m > 0 \<Longrightarrow> (q\<^sup>2 :: nat) mod m = a \<Longrightarrow> \<exists>q' < m. q'\<^sup>2 mod m = a"
  using mod_less_divisor mod_mod_trivial power_mod by blast

lemma q11_upto_def: "q11 = (\<lambda>k. k\<^sup>2 mod 11) ` {..<11}"
  by (simp add: q11_def lessThan_nat_numeral lessThan_Suc insert_commute)

lemma q11_infinite_def: "q11 = (\<lambda>k. k\<^sup>2 mod 11) ` {0..}"
  unfolding q11_upto_def image_def proof (auto, goal_cases)
  case (1 xa)
  show ?case
    using square_mod_lower[of 11 xa "xa\<^sup>2 mod 11"]
      ex_nat_less_eq[of 11 "\<lambda>x. xa\<^sup>2 mod 11 = x\<^sup>2 mod 11"]
    by auto
qed

lemma q63_upto_def: "q63 = (\<lambda>k. k\<^sup>2 mod 63) ` {..<63}"
  by (simp add: q63_def lessThan_nat_numeral lessThan_Suc insert_commute)

lemma q63_infinite_def: "q63 = (\<lambda>k. k\<^sup>2 mod 63) ` {0..}"
  unfolding q63_upto_def image_def proof (auto, goal_cases)
  case (1 xa)
  show ?case
    using square_mod_lower[of 63 xa "xa\<^sup>2 mod 63"]
      ex_nat_less_eq[of 63 "\<lambda>x. xa\<^sup>2 mod 63 = x\<^sup>2 mod 63"]
    by auto
qed

lemma q64_upto_def: "q64 = (\<lambda>k. k\<^sup>2 mod 64) ` {..<64}"
  by (simp add: q64_def lessThan_nat_numeral lessThan_Suc insert_commute)

lemma q64_infinite_def: "q64 = (\<lambda>k. k\<^sup>2 mod 64) ` {0..}"
  unfolding q64_upto_def image_def proof (auto, goal_cases)
  case (1 xa)
  show ?case
    using square_mod_lower[of 64 xa "xa\<^sup>2 mod 64"]
      ex_nat_less_eq[of 64 "\<lambda>x. xa\<^sup>2 mod 64 = x\<^sup>2 mod 64"]
    by auto
qed

lemma q65_upto_def: "q65 = (\<lambda>k. k\<^sup>2 mod 65) ` {..<65}"
  by (simp add: q65_def lessThan_nat_numeral lessThan_Suc insert_commute)

lemma q65_infinite_def: "q65 = (\<lambda>k. k\<^sup>2 mod 65) ` {0..}"
  unfolding q65_upto_def image_def proof (auto, goal_cases)
  case (1 xa)
  show ?case
    using square_mod_lower[of 65 xa "xa\<^sup>2 mod 65"]
      ex_nat_less_eq[of 65 "\<lambda>x. xa\<^sup>2 mod 65 = x\<^sup>2 mod 65"]
    by auto
qed

lemma square_mod_existence:
  fixes n k :: nat
  assumes "\<exists>q. q\<^sup>2 = n"
  shows "\<exists>q. n mod k = q\<^sup>2 mod k"
  using assms by auto

theorem square_test_correct: "square_test n \<longleftrightarrow> is_square n"
proof cases
  assume "is_square n"
  hence  rhs: "\<exists>q. q\<^sup>2 = n" by (auto elim: is_nth_powerE)
  note sq_mod = square_mod_existence[OF this]
  have q64_member: "n mod 64 \<in> q64" using sq_mod[of 64]
    unfolding q64_infinite_def image_def by simp
  let ?r = "n mod 45045"
  have "11 dvd (45045::nat)" "63 dvd (45045::nat)" "65 dvd (45045::nat)" by force+
  then have mod_45045: "?r mod 11 = n mod 11" "?r mod 63 = n mod 63" "?r mod 65 = n mod 65"
    using mod_mod_cancel[of _ 45045 n] by presburger+
  then have "?r mod 11 \<in> q11" "?r mod 63 \<in> q63" "?r mod 65 \<in> q65"
    using sq_mod[of 11] sq_mod[of 63] sq_mod[of 65]
    unfolding q11_infinite_def q63_infinite_def q65_infinite_def image_def mod_45045
    by fast+
  then show ?thesis unfolding square_test_def Let_def using q64_member rhs by auto
next
  assume not_rhs: "\<not>is_square n"
  hence "\<nexists>q. q\<^sup>2 = n" by auto
  then have "(Discrete.sqrt n)\<^sup>2 \<noteq> n" by simp
  then show ?thesis unfolding square_test_def by (auto simp: is_nth_power_def)
qed

definition get_nat_sqrt :: "nat \<Rightarrow> nat option"
  where "get_nat_sqrt n = (if is_square n then Some (Discrete.sqrt n) else None)"

lemma get_nat_sqrt_code [code]:
  "get_nat_sqrt n =
    (if IArray.sub q64_array (n mod 64) \<and> (let r = n mod 45045 in
           IArray.sub q63_array (r mod 63) \<and>
           IArray.sub q65_array (r mod 65) \<and>
           IArray.sub q11_array (r mod 11)) then
       (let x = Discrete.sqrt n in if x\<^sup>2 = n then Some x else None) else None)"
  unfolding get_nat_sqrt_def square_test_correct [symmetric] square_test_def
  using in_q11_code [symmetric] in_q63_code [symmetric]
        in_q64_code [symmetric] in_q65_code [symmetric]
  by (auto split: if_splits simp: Let_def )

end