(*  Title:      HOL/Probability/PMF_Impl.thy
    Author:     Manuel Eberl, TU München

    An implementation of PMFs using Mappings, which are implemented with association lists
    by default. Also includes Quickcheck setup for PMFs.
*)

section \<open>Code generation for PMFs\<close>

theory PMF_Impl
imports Probability_Mass_Function "~~/src/HOL/Library/AList_Mapping"
begin

subsection \<open>General code generation setup\<close>

definition pmf_of_mapping :: "('a, real) mapping \<Rightarrow> 'a pmf" where
  "pmf_of_mapping m = embed_pmf (Mapping.lookup_default 0 m)"

lemma nn_integral_lookup_default:
  fixes m :: "('a, real) mapping"
  assumes "finite (Mapping.keys m)" "All_mapping m (\<lambda>_ x. x \<ge> 0)"
  shows   "nn_integral (count_space UNIV) (\<lambda>k. ennreal (Mapping.lookup_default 0 m k)) =
             ennreal (\<Sum>k\<in>Mapping.keys m. Mapping.lookup_default 0 m k)"
proof -
  have "nn_integral (count_space UNIV) (\<lambda>k. ennreal (Mapping.lookup_default 0 m k)) =
          (\<Sum>x\<in>Mapping.keys m. ennreal (Mapping.lookup_default 0 m x))" using assms
    by (subst nn_integral_count_space'[of "Mapping.keys m"])
       (auto simp: Mapping.lookup_default_def keys_is_none_rep Option.is_none_def)
  also from assms have "\<dots> = ennreal (\<Sum>k\<in>Mapping.keys m. Mapping.lookup_default 0 m k)"
    by (intro sum_ennreal)
       (auto simp: Mapping.lookup_default_def All_mapping_def split: option.splits)
  finally show ?thesis .
qed

lemma pmf_of_mapping:
  assumes "finite (Mapping.keys m)" "All_mapping m (\<lambda>_ p. p \<ge> 0)"
  assumes "(\<Sum>x\<in>Mapping.keys m. Mapping.lookup_default 0 m x) = 1"
  shows   "pmf (pmf_of_mapping m) x = Mapping.lookup_default 0 m x"
  unfolding pmf_of_mapping_def
proof (intro pmf_embed_pmf)
  from assms show "(\<integral>\<^sup>+x. ennreal (Mapping.lookup_default 0 m x) \<partial>count_space UNIV) = 1"
    by (subst nn_integral_lookup_default) (simp_all)
qed (insert assms, simp add: All_mapping_def Mapping.lookup_default_def split: option.splits)

lemma pmf_of_set_pmf_of_mapping:
  assumes "A \<noteq> {}" "set xs = A" "distinct xs"
  shows   "pmf_of_set A = pmf_of_mapping (Mapping.tabulate xs (\<lambda>_. 1 / real (length xs)))"
           (is "?lhs = ?rhs")
  by (rule pmf_eqI, subst pmf_of_mapping)
     (insert assms, auto intro!: All_mapping_tabulate
                         simp: Mapping.lookup_default_def lookup_tabulate distinct_card)

lift_definition mapping_of_pmf :: "'a pmf \<Rightarrow> ('a, real) mapping" is
  "\<lambda>p x. if pmf p x = 0 then None else Some (pmf p x)" .

lemma lookup_default_mapping_of_pmf:
  "Mapping.lookup_default 0 (mapping_of_pmf p) x = pmf p x"
  by (simp add: mapping_of_pmf.abs_eq lookup_default_def Mapping.lookup.abs_eq)

context
begin

interpretation pmf_as_function .

lemma nn_integral_pmf_eq_1: "(\<integral>\<^sup>+ x. ennreal (pmf p x) \<partial>count_space UNIV) = 1"
  by transfer simp_all
end

lemma pmf_of_mapping_mapping_of_pmf [code abstype]:
    "pmf_of_mapping (mapping_of_pmf p) = p"
  unfolding pmf_of_mapping_def
  by (rule pmf_eqI, subst pmf_embed_pmf)
     (insert nn_integral_pmf_eq_1[of p],
      auto simp: lookup_default_mapping_of_pmf split: option.splits)

lemma mapping_of_pmfI:
  assumes "\<And>x. x \<in> Mapping.keys m \<Longrightarrow> Mapping.lookup m x = Some (pmf p x)"
  assumes "Mapping.keys m = set_pmf p"
  shows   "mapping_of_pmf p = m"
  using assms by transfer (rule ext, auto simp: set_pmf_eq)

lemma mapping_of_pmfI':
  assumes "\<And>x. x \<in> Mapping.keys m \<Longrightarrow> Mapping.lookup_default 0 m x = pmf p x"
  assumes "Mapping.keys m = set_pmf p"
  shows   "mapping_of_pmf p = m"
  using assms unfolding Mapping.lookup_default_def
  by transfer (rule ext, force simp: set_pmf_eq)

lemma return_pmf_code [code abstract]:
  "mapping_of_pmf (return_pmf x) = Mapping.update x 1 Mapping.empty"
  by (intro mapping_of_pmfI) (auto simp: lookup_update')

lemma pmf_of_set_code_aux:
  assumes "A \<noteq> {}" "set xs = A" "distinct xs"
  shows   "mapping_of_pmf (pmf_of_set A) = Mapping.tabulate xs (\<lambda>_. 1 / real (length xs))"
  using assms
  by (intro mapping_of_pmfI, subst pmf_of_set)
     (auto simp: lookup_tabulate distinct_card)

definition pmf_of_set_impl where
  "pmf_of_set_impl A = mapping_of_pmf (pmf_of_set A)"

(* This equation can be used to easily implement pmf_of_set for other set implementations *)
lemma pmf_of_set_impl_code_alt:
  assumes "A \<noteq> {}" "finite A"
  shows   "pmf_of_set_impl A =
             (let p = 1 / real (card A)
              in  Finite_Set.fold (\<lambda>x. Mapping.update x p) Mapping.empty A)"
proof -
  define p where "p = 1 / real (card A)"
  let ?m = "Finite_Set.fold (\<lambda>x. Mapping.update x p) Mapping.empty A"
  interpret comp_fun_idem "\<lambda>x. Mapping.update x p"
    by standard (transfer, force simp: fun_eq_iff)+
  have keys: "Mapping.keys ?m = A"
    using assms(2) by (induction A rule: finite_induct) simp_all
  have lookup: "Mapping.lookup ?m x = Some p" if "x \<in> A" for x
    using assms(2) that by (induction A rule: finite_induct) (auto simp: lookup_update')
  from keys lookup assms show ?thesis unfolding pmf_of_set_impl_def
    by (intro mapping_of_pmfI) (simp_all add: Let_def p_def)
qed

lemma pmf_of_set_impl_code [code]:
  "pmf_of_set_impl (set xs) =
    (if xs = [] then
         Code.abort (STR ''pmf_of_set of empty set'') (\<lambda>_. mapping_of_pmf (pmf_of_set (set xs)))
      else let xs' = remdups xs; p = 1 / real (length xs') in
         Mapping.tabulate xs' (\<lambda>_. p))"
  unfolding pmf_of_set_impl_def
  using pmf_of_set_code_aux[of "set xs" "remdups xs"] by (simp add: Let_def)

lemma pmf_of_set_code [code abstract]:
  "mapping_of_pmf (pmf_of_set A) = pmf_of_set_impl A"
  by (simp add: pmf_of_set_impl_def)

lemma pmf_of_multiset_pmf_of_mapping:
  assumes "A \<noteq> {#}" "set xs = set_mset A" "distinct xs"
  shows   "mapping_of_pmf (pmf_of_multiset A) = Mapping.tabulate xs (\<lambda>x. count A x / real (size A))"
  using assms by (intro mapping_of_pmfI) (auto simp: lookup_tabulate)

definition pmf_of_multiset_impl where
  "pmf_of_multiset_impl A = mapping_of_pmf (pmf_of_multiset A)"

lemma pmf_of_multiset_impl_code_alt:
  assumes "A \<noteq> {#}"
  shows   "pmf_of_multiset_impl A =
             (let p = 1 / real (size A)
              in  fold_mset (\<lambda>x. Mapping.map_default x 0 (op + p)) Mapping.empty A)"
proof -
  define p where "p = 1 / real (size A)"
  interpret comp_fun_commute "\<lambda>x. Mapping.map_default x 0 (op + p)"
    unfolding Mapping.map_default_def [abs_def]
    by (standard, intro mapping_eqI ext)
       (simp_all add: o_def lookup_map_entry' lookup_default' lookup_default_def)
  let ?m = "fold_mset (\<lambda>x. Mapping.map_default x 0 (op + p)) Mapping.empty A"
  have keys: "Mapping.keys ?m = set_mset A" by (induction A) simp_all
  have lookup: "Mapping.lookup_default 0 ?m x = real (count A x) * p" for x
    by (induction A)
       (simp_all add: lookup_map_default' lookup_default_def lookup_empty ring_distribs)
  from keys lookup assms show ?thesis unfolding pmf_of_multiset_impl_def
    by (intro mapping_of_pmfI') (simp_all add: Let_def p_def)
qed

lemma pmf_of_multiset_impl_code [code]:
  "pmf_of_multiset_impl (mset xs) =
     (if xs = [] then
        Code.abort (STR ''pmf_of_multiset of empty multiset'')
          (\<lambda>_. mapping_of_pmf (pmf_of_multiset (mset xs)))
      else let xs' = remdups xs; p = 1 / real (length xs) in
         Mapping.tabulate xs' (\<lambda>x. real (count (mset xs) x) * p))"
  using pmf_of_multiset_pmf_of_mapping[of "mset xs" "remdups xs"]
  by (simp add: pmf_of_multiset_impl_def)

lemma pmf_of_multiset_code [code abstract]:
  "mapping_of_pmf (pmf_of_multiset A) = pmf_of_multiset_impl A"
  by (simp add: pmf_of_multiset_impl_def)

lemma bernoulli_pmf_code [code abstract]:
  "mapping_of_pmf (bernoulli_pmf p) =
     (if p \<le> 0 then Mapping.update False 1 Mapping.empty
      else if p \<ge> 1 then Mapping.update True 1 Mapping.empty
      else Mapping.update False (1 - p) (Mapping.update True p Mapping.empty))"
  by (intro mapping_of_pmfI) (auto simp: bernoulli_pmf.rep_eq lookup_update' set_pmf_eq)

lemma pmf_code [code]: "pmf p x = Mapping.lookup_default 0 (mapping_of_pmf p) x"
  unfolding mapping_of_pmf_def Mapping.lookup_default_def
  by (auto split: option.splits simp: id_def Mapping.lookup.abs_eq)

lemma set_pmf_code [code]: "set_pmf p = Mapping.keys (mapping_of_pmf p)"
  by transfer (auto simp: dom_def set_pmf_eq)

lemma keys_mapping_of_pmf [simp]: "Mapping.keys (mapping_of_pmf p) = set_pmf p"
  by transfer (auto simp: dom_def set_pmf_eq)

definition fold_combine_plus where
  "fold_combine_plus = comm_monoid_set.F (Mapping.combine (op + :: real \<Rightarrow> _)) Mapping.empty"

context
begin

interpretation fold_combine_plus: combine_mapping_abel_semigroup "op + :: real \<Rightarrow> _"
  by unfold_locales (simp_all add: add_ac)

qualified lemma lookup_default_fold_combine_plus:
  fixes A :: "'b set" and f :: "'b \<Rightarrow> ('a, real) mapping"
  assumes "finite A"
  shows   "Mapping.lookup_default 0 (fold_combine_plus f A) x =
             (\<Sum>y\<in>A. Mapping.lookup_default 0 (f y) x)"
  unfolding fold_combine_plus_def using assms
    by (induction A rule: finite_induct)
       (simp_all add: lookup_default_empty lookup_default_neutral_combine)

qualified lemma keys_fold_combine_plus:
  "finite A \<Longrightarrow> Mapping.keys (fold_combine_plus f A) = (\<Union>x\<in>A. Mapping.keys (f x))"
  by (simp add: fold_combine_plus_def fold_combine_plus.keys_fold_combine)

qualified lemma fold_combine_plus_code [code]:
  "fold_combine_plus g (set xs) = foldr (\<lambda>x. Mapping.combine op+ (g x)) (remdups xs) Mapping.empty"
  by (simp add: fold_combine_plus_def fold_combine_plus.fold_combine_code)

private lemma lookup_default_0_map_values:
  assumes "f x 0 = 0"
  shows   "Mapping.lookup_default 0 (Mapping.map_values f m) x = f x (Mapping.lookup_default 0 m x)"
  unfolding Mapping.lookup_default_def
  using assms by transfer (auto split: option.splits)

qualified lemma mapping_of_bind_pmf:
  assumes "finite (set_pmf p)"
  shows   "mapping_of_pmf (bind_pmf p f) =
             fold_combine_plus (\<lambda>x. Mapping.map_values (\<lambda>_. op * (pmf p x))
               (mapping_of_pmf (f x))) (set_pmf p)"
  using assms
  by (intro mapping_of_pmfI')
     (auto simp: keys_fold_combine_plus lookup_default_fold_combine_plus
                 pmf_bind integral_measure_pmf lookup_default_0_map_values
                 lookup_default_mapping_of_pmf mult_ac)

lift_definition bind_pmf_aux :: "'a pmf \<Rightarrow> ('a \<Rightarrow> 'b pmf) \<Rightarrow> 'a set \<Rightarrow> ('b, real) mapping" is
  "\<lambda>(p :: 'a pmf) (f :: 'a \<Rightarrow> 'b pmf) (A::'a set) (x::'b).
     if x \<in> (\<Union>y\<in>A. set_pmf (f y)) then
       Some (measure_pmf.expectation p (\<lambda>y. indicator A y * pmf (f y) x))
     else None" .

lemma keys_bind_pmf_aux [simp]:
  "Mapping.keys (bind_pmf_aux p f A) = (\<Union>x\<in>A. set_pmf (f x))"
  by transfer (auto split: if_splits)

lemma lookup_default_bind_pmf_aux:
  "Mapping.lookup_default 0 (bind_pmf_aux p f A) x =
     (if x \<in> (\<Union>y\<in>A. set_pmf (f y)) then
        measure_pmf.expectation p (\<lambda>y. indicator A y * pmf (f y) x) else 0)"
  unfolding lookup_default_def by transfer' simp_all

lemma lookup_default_bind_pmf_aux' [simp]:
  "Mapping.lookup_default 0 (bind_pmf_aux p f (set_pmf p)) x = pmf (bind_pmf p f) x"
  unfolding lookup_default_def
  by transfer (auto simp: pmf_bind AE_measure_pmf_iff set_pmf_eq
                    intro!: integral_cong_AE integral_eq_zero_AE)

lemma bind_pmf_aux_correct:
  "mapping_of_pmf (bind_pmf p f) = bind_pmf_aux p f (set_pmf p)"
  by (intro mapping_of_pmfI') simp_all

lemma bind_pmf_aux_code_aux:
  assumes "finite A"
  shows   "bind_pmf_aux p f A =
             fold_combine_plus (\<lambda>x. Mapping.map_values (\<lambda>_. op * (pmf p x))
               (mapping_of_pmf (f x))) A" (is "?lhs = ?rhs")
proof (intro mapping_eqI'[where d = 0])
  fix x assume "x \<in> Mapping.keys ?lhs"
  then obtain y where y: "y \<in> A" "x \<in> set_pmf (f y)" by auto
  hence "Mapping.lookup_default 0 ?lhs x =
           measure_pmf.expectation p (\<lambda>y. indicator A y * pmf (f y) x)"
    by (auto simp: lookup_default_bind_pmf_aux)
  also from assms have "\<dots> = (\<Sum>y\<in>A. pmf p y * pmf (f y) x)"
    by (subst integral_measure_pmf [of A])
       (auto simp: set_pmf_eq indicator_def mult_ac split: if_splits)
  also from assms have "\<dots> = Mapping.lookup_default 0 ?rhs x"
    by (simp add: lookup_default_fold_combine_plus lookup_default_0_map_values
          lookup_default_mapping_of_pmf)
  finally show "Mapping.lookup_default 0 ?lhs x = Mapping.lookup_default 0 ?rhs x" .
qed (insert assms, simp_all add: keys_fold_combine_plus)

lemma bind_pmf_aux_code [code]:
  "bind_pmf_aux p f (set xs) =
     fold_combine_plus (\<lambda>x. Mapping.map_values (\<lambda>_. op * (pmf p x))
               (mapping_of_pmf (f x))) (set xs)"
  by (rule bind_pmf_aux_code_aux) simp_all

lemmas bind_pmf_code [code abstract] = bind_pmf_aux_correct

end

hide_const (open) fold_combine_plus

lift_definition cond_pmf_impl :: "'a pmf \<Rightarrow> 'a set \<Rightarrow> ('a, real) mapping option" is
  "\<lambda>p A. if A \<inter> set_pmf p = {} then None else
     Some (\<lambda>x. if x \<in> A \<inter> set_pmf p then Some (pmf p x / measure_pmf.prob p A) else None)" .

lemma cond_pmf_impl_code_alt:
  assumes "finite A"
  shows   "cond_pmf_impl p A = (
             let C = A \<inter> set_pmf p;
                 prob = (\<Sum>x\<in>C. pmf p x)
             in  if prob = 0 then
                   None
                 else
                   Some (Mapping.map_values (\<lambda>_ y. y / prob)
                     (Mapping.filter (\<lambda>k _. k \<in> C) (mapping_of_pmf p))))"
proof -
  define C where "C = A \<inter> set_pmf p"
  define prob where "prob = (\<Sum>x\<in>C. pmf p x)"
  also note C_def
  also from assms have "(\<Sum>x\<in>A \<inter> set_pmf p. pmf p x) = (\<Sum>x\<in>A. pmf p x)"
    by (intro sum.mono_neutral_left) (auto simp: set_pmf_eq)
  finally have prob1: "prob = (\<Sum>x\<in>A. pmf p x)" .
  hence prob2: "prob = measure_pmf.prob p A"
    using assms by (subst measure_measure_pmf_finite) simp_all
  have prob3: "prob = 0 \<longleftrightarrow> A \<inter> set_pmf p = {}"
    by (subst prob1, subst sum_nonneg_eq_0_iff) (auto simp: set_pmf_eq assms)
  from assms have prob4: "prob = measure_pmf.prob p C"
    unfolding prob_def by (intro measure_measure_pmf_finite [symmetric]) (simp_all add: C_def)

  show ?thesis
  proof (cases "prob = 0")
    case True
    hence "A \<inter> set_pmf p = {}" by (subst (asm) prob3)
    with True show ?thesis by (simp add: Let_def prob_def C_def cond_pmf_impl.abs_eq)
  next
    case False
    hence A: "C \<noteq> {}" unfolding C_def by (subst (asm) prob3) auto
    with prob3 have prob_nz: "prob \<noteq> 0" by (auto simp: C_def)
    fix x
    have "cond_pmf_impl p A =
            Some (mapping.Mapping (\<lambda>x. if x \<in> C then
              Some (pmf p x / measure_pmf.prob p C) else None))"
         (is "_ = Some ?m")
      using A prob2 prob4 unfolding C_def by transfer (auto simp: fun_eq_iff)
    also have "?m = Mapping.map_values (\<lambda>_ y. y / prob)
                 (Mapping.filter (\<lambda>k _. k \<in> C) (mapping_of_pmf p))"
      using prob_nz prob4 assms unfolding C_def
      by transfer (auto simp: fun_eq_iff set_pmf_eq)
    finally show ?thesis using False by (simp add: Let_def prob_def C_def)
  qed
qed

lemma cond_pmf_impl_code [code]:
  "cond_pmf_impl p (set xs) = (
     let C = set xs \<inter> set_pmf p;
         prob = (\<Sum>x\<in>C. pmf p x)
     in  if prob = 0 then
           None
         else
           Some (Mapping.map_values (\<lambda>_ y. y / prob)
             (Mapping.filter (\<lambda>k _. k \<in> C) (mapping_of_pmf p))))"
  by (rule cond_pmf_impl_code_alt) simp_all

lemma cond_pmf_code [code abstract]:
  "mapping_of_pmf (cond_pmf p A) =
     (case cond_pmf_impl p A of
        None \<Rightarrow> Code.abort (STR ''cond_pmf with set of probability 0'')
                  (\<lambda>_. mapping_of_pmf (cond_pmf p A))
      | Some m \<Rightarrow> m)"
proof (cases "cond_pmf_impl p A")
  case (Some m)
  hence A: "set_pmf p \<inter> A \<noteq> {}" by transfer (auto split: if_splits)
  from Some have B: "Mapping.keys m = set_pmf (cond_pmf p A)"
    by (subst set_cond_pmf[OF A], transfer) (auto split: if_splits)
  with Some A have "mapping_of_pmf (cond_pmf p A) = m"
    by (intro mapping_of_pmfI[OF _ B], transfer) (auto split: if_splits simp: pmf_cond)
  with Some show ?thesis by simp
qed simp_all

lemma binomial_pmf_code [code abstract]:
  "mapping_of_pmf (binomial_pmf n p) = (
     if p < 0 \<or> p > 1 then
       Code.abort (STR ''binomial_pmf with invalid probability'')
         (\<lambda>_. mapping_of_pmf (binomial_pmf n p))
     else if p = 0 then Mapping.update 0 1 Mapping.empty
     else if p = 1 then Mapping.update n 1 Mapping.empty
     else Mapping.tabulate [0..<Suc n] (\<lambda>k. real (n choose k) * p ^ k * (1 - p) ^ (n - k)))"
  by (cases "p < 0 \<or> p > 1")
     (simp, intro mapping_of_pmfI,
      auto simp: lookup_update' lookup_empty set_pmf_binomial_eq lookup_tabulate split: if_splits)

lemma pred_pmf_code [code]:
  "pred_pmf P p = (\<forall>x\<in>set_pmf p. P x)"
  by (auto simp: pred_pmf_def)

lemma mapping_of_pmf_pmf_of_list:
  assumes "\<And>x. x \<in> snd ` set xs \<Longrightarrow> x > 0" "sum_list (map snd xs) = 1"
  shows   "mapping_of_pmf (pmf_of_list xs) =
             Mapping.tabulate (remdups (map fst xs))
               (\<lambda>x. sum_list (map snd (filter (\<lambda>z. fst z = x) xs)))"
proof -
  from assms have wf: "pmf_of_list_wf xs" by (intro pmf_of_list_wfI) force
  with assms have "set_pmf (pmf_of_list xs) = fst ` set xs"
    by (intro set_pmf_of_list_eq) auto
  with wf show ?thesis
    by (intro mapping_of_pmfI) (auto simp: lookup_tabulate pmf_pmf_of_list)
qed

lemma mapping_of_pmf_pmf_of_list':
  assumes "pmf_of_list_wf xs"
  defines "xs' \<equiv> filter (\<lambda>z. snd z \<noteq> 0) xs"
  shows   "mapping_of_pmf (pmf_of_list xs) =
             Mapping.tabulate (remdups (map fst xs'))
               (\<lambda>x. sum_list (map snd (filter (\<lambda>z. fst z = x) xs')))" (is "_ = ?rhs")
proof -
  have wf: "pmf_of_list_wf xs'" unfolding xs'_def by (rule pmf_of_list_remove_zeros) fact
  have pos: "\<forall>x\<in>snd`set xs'. x > 0" using assms(1) unfolding xs'_def
    by (force simp: pmf_of_list_wf_def)
  from assms have "pmf_of_list xs = pmf_of_list xs'"
    unfolding xs'_def by (subst pmf_of_list_remove_zeros) simp_all
  also from wf pos have "mapping_of_pmf \<dots> = ?rhs"
    by (intro mapping_of_pmf_pmf_of_list) (auto simp: pmf_of_list_wf_def)
  finally show ?thesis .
qed

lemma pmf_of_list_wf_code [code]:
  "pmf_of_list_wf xs \<longleftrightarrow> list_all (\<lambda>z. snd z \<ge> 0) xs \<and> sum_list (map snd xs) = 1"
  by (auto simp add: pmf_of_list_wf_def list_all_def)

lemma pmf_of_list_code [code abstract]:
  "mapping_of_pmf (pmf_of_list xs) = (
     if pmf_of_list_wf xs then
       let xs' = filter (\<lambda>z. snd z \<noteq> 0) xs
       in  Mapping.tabulate (remdups (map fst xs'))
             (\<lambda>x. sum_list (map snd (filter (\<lambda>z. fst z = x) xs')))
     else
       Code.abort (STR ''Invalid list for pmf_of_list'') (\<lambda>_. mapping_of_pmf (pmf_of_list xs)))"
  using mapping_of_pmf_pmf_of_list'[of xs] by (simp add: Let_def)

lemma mapping_of_pmf_eq_iff [simp]:
  "mapping_of_pmf p = mapping_of_pmf q \<longleftrightarrow> p = (q :: 'a pmf)"
proof (transfer, intro iffI pmf_eqI)
  fix p q :: "'a pmf" and x :: 'a
  assume "(\<lambda>x. if pmf p x = 0 then None else Some (pmf p x)) =
            (\<lambda>x. if pmf q x = 0 then None else Some (pmf q x))"
  hence "(if pmf p x = 0 then None else Some (pmf p x)) =
           (if pmf q x = 0 then None else Some (pmf q x))" for x
    by (simp add: fun_eq_iff)
  from this[of x] show "pmf p x = pmf q x" by (auto split: if_splits)
qed (simp_all cong: if_cong)

subsection \<open>Code abbreviations for integrals and probabilities\<close>

text \<open>
  Integrals and probabilities are defined for general measures, so we cannot give any
  code equations directly. We can, however, specialise these constants them to PMFs,
  give code equations for these specialised constants, and tell the code generator
  to unfold the original constants to the specialised ones whenever possible.
\<close>

definition pmf_integral where
  "pmf_integral p f = lebesgue_integral (measure_pmf p) (f :: _ \<Rightarrow> real)"

definition pmf_set_integral where
  "pmf_set_integral p f A = lebesgue_integral (measure_pmf p) (\<lambda>x. indicator A x * f x :: real)"

definition pmf_prob where
  "pmf_prob p A = measure_pmf.prob p A"

lemma pmf_prob_compl: "pmf_prob p (-A) = 1 - pmf_prob p A"
  using measure_pmf.prob_compl[of A p] by (simp add: pmf_prob_def Compl_eq_Diff_UNIV)

lemma pmf_integral_pmf_set_integral [code]:
  "pmf_integral p f = pmf_set_integral p f (set_pmf p)"
  unfolding pmf_integral_def pmf_set_integral_def
  by (intro integral_cong_AE) (simp_all add: AE_measure_pmf_iff)

lemma pmf_prob_pmf_set_integral:
  "pmf_prob p A = pmf_set_integral p (\<lambda>_. 1) A"
  by (simp add: pmf_prob_def pmf_set_integral_def)

lemma pmf_set_integral_code_alt_finite:
  "finite A \<Longrightarrow> pmf_set_integral p f A = (\<Sum>x\<in>A. pmf p x * f x)"
  unfolding pmf_set_integral_def
  by (subst integral_measure_pmf[of A]) (auto simp: indicator_def mult_ac split: if_splits)

lemma pmf_set_integral_code [code]:
  "pmf_set_integral p f (set xs) = (\<Sum>x\<in>set xs. pmf p x * f x)"
  by (rule pmf_set_integral_code_alt_finite) simp_all

lemma pmf_prob_code_alt_finite:
  "finite A \<Longrightarrow> pmf_prob p A = (\<Sum>x\<in>A. pmf p x)"
  by (simp add: pmf_prob_pmf_set_integral pmf_set_integral_code_alt_finite)

lemma pmf_prob_code [code]:
  "pmf_prob p (set xs) = (\<Sum>x\<in>set xs. pmf p x)"
  "pmf_prob p (List.coset xs) = 1 - (\<Sum>x\<in>set xs. pmf p x)"
  by (simp_all add: pmf_prob_code_alt_finite pmf_prob_compl)

lemma pmf_prob_code_unfold [code_abbrev]: "pmf_prob p = measure_pmf.prob p"
  by (intro ext) (simp add: pmf_prob_def)

(* FIXME: Why does this not work without parameters? *)
lemma pmf_integral_code_unfold [code_abbrev]: "pmf_integral p = measure_pmf.expectation p"
  by (intro ext) (simp add: pmf_integral_def)

definition "pmf_of_alist xs = embed_pmf (\<lambda>x. case map_of xs x of Some p \<Rightarrow> p | None \<Rightarrow> 0)"

lemma pmf_of_mapping_Mapping [code_post]:
    "pmf_of_mapping (Mapping xs) = pmf_of_alist xs"
  unfolding pmf_of_mapping_def Mapping.lookup_default_def [abs_def] pmf_of_alist_def
  by transfer simp_all

instantiation pmf :: (equal) equal
begin

definition "equal_pmf p q = (mapping_of_pmf p = mapping_of_pmf (q :: 'a pmf))"

instance by standard (simp add: equal_pmf_def)
end

definition single :: "'a \<Rightarrow> 'a multiset" where
"single s = {#s#}"

definition (in term_syntax)
  pmfify :: "('a::typerep multiset \<times> (unit \<Rightarrow> Code_Evaluation.term)) \<Rightarrow>
             'a \<times> (unit \<Rightarrow> Code_Evaluation.term) \<Rightarrow>
             'a pmf \<times> (unit \<Rightarrow> Code_Evaluation.term)" where
  [code_unfold]: "pmfify A x =
    Code_Evaluation.valtermify pmf_of_multiset {\<cdot>}
      (Code_Evaluation.valtermify (op +) {\<cdot>} A {\<cdot>}
       (Code_Evaluation.valtermify single {\<cdot>} x))"

notation fcomp (infixl "\<circ>>" 60)
notation scomp (infixl "\<circ>\<rightarrow>" 60)

instantiation pmf :: (random) random
begin

definition
  "Quickcheck_Random.random i =
     Quickcheck_Random.random i \<circ>\<rightarrow> (\<lambda>A.
       Quickcheck_Random.random i \<circ>\<rightarrow> (\<lambda>x. Pair (pmfify A x)))"

instance ..

end

no_notation fcomp (infixl "\<circ>>" 60)
no_notation scomp (infixl "\<circ>\<rightarrow>" 60)

instantiation pmf :: (full_exhaustive) full_exhaustive
begin

definition full_exhaustive_pmf :: "('a pmf \<times> (unit \<Rightarrow> term) \<Rightarrow> (bool \<times> term list) option) \<Rightarrow> natural \<Rightarrow> (bool \<times> term list) option"
where
  "full_exhaustive_pmf f i =
     Quickcheck_Exhaustive.full_exhaustive (\<lambda>A.
       Quickcheck_Exhaustive.full_exhaustive (\<lambda>x. f (pmfify A x)) i) i"

instance ..

end

end


(* Author: Tobias Nipkow *)

section \<open>Leftist Heap\<close>

theory Leftist_Heap
imports
  "HOL-Library.Pattern_Aliases"
  Tree2
  Priority_Queue_Specs
  Complex_Main
begin

fun mset_tree :: "('a*'b) tree \<Rightarrow> 'a multiset" where
"mset_tree Leaf = {#}" |
"mset_tree (Node l (a, _) r) = {#a#} + mset_tree l + mset_tree r"

type_synonym 'a lheap = "('a*nat)tree"

fun rank :: "'a lheap \<Rightarrow> nat" where
"rank Leaf = 0" |
"rank (Node _ _ r) = rank r + 1"

fun rk :: "'a lheap \<Rightarrow> nat" where
"rk Leaf = 0" |
"rk (Node _ (_, n) _) = n"

text\<open>The invariants:\<close>

fun (in linorder) heap :: "('a*'b) tree \<Rightarrow> bool" where
"heap Leaf = True" |
"heap (Node l (m, _) r) =
  (heap l \<and> heap r \<and> (\<forall>x \<in> set_tree l \<union> set_tree r. m \<le> x))"

fun ltree :: "'a lheap \<Rightarrow> bool" where
"ltree Leaf = True" |
"ltree (Node l (a, n) r) =
 (n = rank r + 1 \<and> rank l \<ge> rank r \<and> ltree l & ltree r)"

definition empty :: "'a lheap" where
"empty = Leaf"

definition node :: "'a lheap \<Rightarrow> 'a \<Rightarrow> 'a lheap \<Rightarrow> 'a lheap" where
"node l a r =
 (let rl = rk l; rr = rk r
  in if rl \<ge> rr then Node l (a,rr+1) r else Node r (a,rl+1) l)"

fun get_min :: "'a lheap \<Rightarrow> 'a" where
"get_min(Node l (a, n) r) = a"

text \<open>For function \<open>merge\<close>:\<close>
unbundle pattern_aliases

fun merge :: "'a::ord lheap \<Rightarrow> 'a lheap \<Rightarrow> 'a lheap" where
"merge Leaf t = t" |
"merge t Leaf = t" |
"merge (Node l1 (a1, n1) r1 =: t1) (Node l2 (a2, n2) r2 =: t2) =
   (if a1 \<le> a2 then node l1 a1 (merge r1 t2)
    else node l2 a2 (merge t1 r2))"

text \<open>Termination of @{const merge}: by sum or lexicographic product of the sizes
of the two arguments. Isabelle uses a lexicographic product.\<close>

lemma merge_code: "merge t1 t2 = (case (t1,t2) of
  (Leaf, _) \<Rightarrow> t2 |
  (_, Leaf) \<Rightarrow> t1 |
  (Node l1 (a1, n1) r1, Node l2 (a2, n2) r2) \<Rightarrow>
    if a1 \<le> a2 then node l1 a1 (merge r1 t2) else node l2 a2 (merge t1 r2))"
by(induction t1 t2 rule: merge.induct) (simp_all split: tree.split)

hide_const (open) insert

definition insert :: "'a::ord \<Rightarrow> 'a lheap \<Rightarrow> 'a lheap" where
"insert x t = merge (Node Leaf (x,1) Leaf) t"

fun del_min :: "'a::ord lheap \<Rightarrow> 'a lheap" where
"del_min Leaf = Leaf" |
"del_min (Node l _ r) = merge l r"

subsection "Lemmas"

lemma mset_tree_empty: "mset_tree t = {#} \<longleftrightarrow> t = Leaf"
by(cases t) auto

lemma rk_eq_rank[simp]: "ltree t \<Longrightarrow> rk t = rank t"
by(cases t) auto

lemma ltree_node: "ltree (node l a r) \<longleftrightarrow> ltree l \<and> ltree r"
by(auto simp add: node_def)

lemma heap_node: "heap (node l a r) \<longleftrightarrow>
  heap l \<and> heap r \<and> (\<forall>x \<in> set_tree l \<union> set_tree r. a \<le> x)"
by(auto simp add: node_def)

lemma set_tree_mset: "set_tree t = set_mset(mset_tree t)"
by(induction t) auto

subsection "Functional Correctness"

lemma mset_merge: "mset_tree (merge h1 h2) = mset_tree h1 + mset_tree h2"
by (induction h1 h2 rule: merge.induct) (auto simp add: node_def ac_simps)

lemma mset_insert: "mset_tree (insert x t) = mset_tree t + {#x#}"
by (auto simp add: insert_def mset_merge)

lemma get_min: "\<lbrakk> heap h;  h \<noteq> Leaf \<rbrakk> \<Longrightarrow> get_min h = Min(set_tree h)"
by (induction h) (auto simp add: eq_Min_iff)

lemma mset_del_min: "mset_tree (del_min h) = mset_tree h - {# get_min h #}"
by (cases h) (auto simp: mset_merge)

lemma ltree_merge: "\<lbrakk> ltree l; ltree r \<rbrakk> \<Longrightarrow> ltree (merge l r)"
proof(induction l r rule: merge.induct)
  case (3 l1 a1 n1 r1 l2 a2 n2 r2)
  show ?case (is "ltree(merge ?t1 ?t2)")
  proof cases
    assume "a1 \<le> a2"
    hence "ltree (merge ?t1 ?t2) = ltree (node l1 a1 (merge r1 ?t2))" by simp
    also have "\<dots> = (ltree l1 \<and> ltree(merge r1 ?t2))"
      by(simp add: ltree_node)
    also have "..." using "3.prems" "3.IH"(1)[OF \<open>a1 \<le> a2\<close>] by (simp)
    finally show ?thesis .
  next (* analogous but automatic *)
    assume "\<not> a1 \<le> a2"
    thus ?thesis using 3 by(simp)(auto simp: ltree_node)
  qed
qed simp_all

lemma heap_merge: "\<lbrakk> heap l; heap r \<rbrakk> \<Longrightarrow> heap (merge l r)"
proof(induction l r rule: merge.induct)
  case 3 thus ?case by(auto simp: heap_node mset_merge ball_Un set_tree_mset)
qed simp_all

lemma ltree_insert: "ltree t \<Longrightarrow> ltree(insert x t)"
by(simp add: insert_def ltree_merge del: merge.simps split: tree.split)

lemma heap_insert: "heap t \<Longrightarrow> heap(insert x t)"
by(simp add: insert_def heap_merge del: merge.simps split: tree.split)

lemma ltree_del_min: "ltree t \<Longrightarrow> ltree(del_min t)"
by(cases t)(auto simp add: ltree_merge simp del: merge.simps)

lemma heap_del_min: "heap t \<Longrightarrow> heap(del_min t)"
by(cases t)(auto simp add: heap_merge simp del: merge.simps)

text \<open>Last step of functional correctness proof: combine all the above lemmas
to show that leftist heaps satisfy the specification of priority queues with merge.\<close>

interpretation lheap: Priority_Queue_Merge
where empty = empty and is_empty = "\<lambda>h. h = Leaf"
and insert = insert and del_min = del_min
and get_min = get_min and merge = merge
and invar = "\<lambda>h. heap h \<and> ltree h" and mset = mset_tree
proof(standard, goal_cases)
  case 1 show ?case by (simp add: empty_def)
next
  case (2 q) show ?case by (cases q) auto
next
  case 3 show ?case by(rule mset_insert)
next
  case 4 show ?case by(rule mset_del_min)
next
  case 5 thus ?case by(simp add: get_min mset_tree_empty set_tree_mset)
next
  case 6 thus ?case by(simp add: empty_def)
next
  case 7 thus ?case by(simp add: heap_insert ltree_insert)
next
  case 8 thus ?case by(simp add: heap_del_min ltree_del_min)
next
  case 9 thus ?case by (simp add: mset_merge)
next
  case 10 thus ?case by (simp add: heap_merge ltree_merge)
qed

subsection "Complexity"

lemma pow2_rank_size1: "ltree t \<Longrightarrow> 2 ^ rank t \<le> size1 t"
proof(induction t rule: tree2_induct)
  case Leaf show ?case by simp
next
  case (Node l a n r)
  hence "rank r \<le> rank l" by simp
  hence *: "(2::nat) ^ rank r \<le> 2 ^ rank l" by simp
  have "(2::nat) ^ rank \<langle>l, (a, n), r\<rangle> = 2 ^ rank r + 2 ^ rank r"
    by(simp add: mult_2)
  also have "\<dots> \<le> size1 l + size1 r"
    using Node * by (simp del: power_increasing_iff)
  also have "\<dots> = size1 \<langle>l, (a, n), r\<rangle>" by simp
  finally show ?case .
qed

text\<open>Explicit termination argument: sum of sizes\<close>

fun t_merge :: "'a::ord lheap \<Rightarrow> 'a lheap \<Rightarrow> nat" where
"t_merge Leaf t = 1" |
"t_merge t Leaf = 1" |
"t_merge (Node l1 (a1, n1) r1 =: t1) (Node l2 (a2, n2) r2 =: t2) =
  (if a1 \<le> a2 then 1 + t_merge r1 t2
   else 1 + t_merge t1 r2)"

definition t_insert :: "'a::ord \<Rightarrow> 'a lheap \<Rightarrow> nat" where
"t_insert x t = t_merge (Node Leaf (x, 1) Leaf) t"

fun t_del_min :: "'a::ord lheap \<Rightarrow> nat" where
"t_del_min Leaf = 1" |
"t_del_min (Node l _ r) = t_merge l r"

lemma t_merge_rank: "t_merge l r \<le> rank l + rank r + 1"
proof(induction l r rule: merge.induct)
  case 3 thus ?case by(simp)
qed simp_all

corollary t_merge_log: assumes "ltree l" "ltree r"
  shows "t_merge l r \<le> log 2 (size1 l) + log 2 (size1 r) + 1"
using le_log2_of_power[OF pow2_rank_size1[OF assms(1)]]
  le_log2_of_power[OF pow2_rank_size1[OF assms(2)]] t_merge_rank[of l r]
by linarith

corollary t_insert_log: "ltree t \<Longrightarrow> t_insert x t \<le> log 2 (size1 t) + 2"
using t_merge_log[of "Node Leaf (x, 1) Leaf" t]
by(simp add: t_insert_def split: tree.split)

(* FIXME mv ? *)
lemma ld_ld_1_less:
  assumes "x > 0" "y > 0" shows "log 2 x + log 2 y + 1 < 2 * log 2 (x+y)"
proof -
  have "2 powr (log 2 x + log 2 y + 1) = 2*x*y"
    using assms by(simp add: powr_add)
  also have "\<dots> < (x+y)^2" using assms
    by(simp add: numeral_eq_Suc algebra_simps add_pos_pos)
  also have "\<dots> = 2 powr (2 * log 2 (x+y))"
    using assms by(simp add: powr_add log_powr[symmetric])
  finally show ?thesis by simp
qed

corollary t_del_min_log: assumes "ltree t"
  shows "t_del_min t \<le> 2 * log 2 (size1 t) + 1"
proof(cases t rule: tree2_cases)
  case Leaf thus ?thesis using assms by simp
next
  case [simp]: (Node t1 _ _ t2)
  have "t_del_min t = t_merge t1 t2" by simp
  also have "\<dots> \<le> log 2 (size1 t1) + log 2 (size1 t2) + 1"
    using \<open>ltree t\<close> by (auto simp: t_merge_log simp del: t_merge.simps)
  also have "\<dots> \<le> 2 * log 2 (size1 t) + 1"
    using ld_ld_1_less[of "size1 t1" "size1 t2"] by (simp)
  finally show ?thesis .
qed

end


section \<open>Matrices\<close>
theory Matrices
  imports
    Nat
    (* TSFunctions *)
begin

unbundle no_HOL_groups_syntax
unbundle no_HOL_order_syntax

definition "matrices A m n \<equiv> ([0,\<dots>,m[ :: set) \<rightarrow>s ([0,\<dots>,n[ :: set) \<rightarrow>s A"

definition [typedef]: "Matrix A m n \<equiv> Element [0,\<dots>,m[ \<rightarrow>cs Element [0,\<dots>,n[ \<rightarrow>cs A"

lemma mem_matrices_iff_Matrix:
  "M \<in> matrices A m n \<longleftrightarrow> M : Matrix (Element A) m n"
  (is "?lhs \<longleftrightarrow> ?rhs")
proof -
  have
    "?lhs \<longleftrightarrow> M : Element [0,\<dots>,m[ \<rightarrow>cs Element (([0,\<dots>,n[ :: set) \<rightarrow>s A)"
    unfolding matrices_def by auto
  also have "...  \<longleftrightarrow> ?rhs"
  proof -
    have "\<And>N. N : Element (([0,\<dots>,n[ :: set) \<rightarrow>s A) \<longleftrightarrow> N : (Element [0,\<dots>,n[ \<rightarrow>cs Element A)"
      by (subst mem_iff_Element[symmetric]) auto
    then show ?thesis
      unfolding Matrix_def by (auto elim: CDep_Function_covariant_codom)
  qed
  finally show ?thesis .
qed

soft_type_translation
  "M \<in> matrices A m n" \<rightleftharpoons> "M : Matrix (Element A) m n"
  using mem_matrices_iff_Matrix by auto

lemma Matrix_if_CDep_Function [derive]:
  assumes "M : Element [0,\<dots>,m[ \<rightarrow>cs Element [0,\<dots>,n[ \<rightarrow>cs A"
  shows "M : Matrix A m n"
  using assms Dep_Function_covariant_codom unfolding Matrix_def by auto

lemma eval_type [type]:
  "eval : Matrix A m n \<Rightarrow> (x : Element [0,\<dots>,m[) \<Rightarrow> Element [0,\<dots>,n[ \<rightarrow>cs A"
  unfolding Matrix_def by discharge_types

lemma Matrix_eval_typeI:
  assumes "M : Matrix A m n" "i : Nat" "j : Nat"
  and i_lt_m: "i < m"
  and j_lt_n: "j < n"
  shows "M`i`j : A"
proof -
  have "i \<in> [0,\<dots>,m[" "j \<in> [0,\<dots>,n[" by auto
  then show ?thesis
    using assms [[type_derivation_depth=4]] unfolding Matrix_def
    by discharge_types
qed

subsection \<open>Zero\<close>

definition "Matrix_zero Z m n \<equiv> \<lambda>i \<in> [0,\<dots>,m[. \<lambda>j \<in> [0,\<dots>,n[. zero Z"

lemma Matrix_zero_type [type]:
  "Matrix_zero : Zero A \<Rightarrow> (m : Nat) \<Rightarrow> (n : Nat) \<Rightarrow> Matrix A m n"
  unfolding Matrix_zero_def by discharge_types

lemma Matrix_zero_eq_zero [simp]:
  assumes "i : Nat" "j : Nat"
  and i_lt_m: "i < m"
  and j_lt_n: "j < n"
  shows "(Matrix_zero Z m n)`i`j = zero Z"
proof -
(*
Note Kevin: The simplifier tries to apply eval_lambda_eq. We need to discharge
i \<in> [0,\<dots>,m[.  This goal gets transformed to i : Element [0,\<dots>,m[.
Now, the type-derivator cannot solve this as there's no good rule for this type.
We might think about tagging {@thm in_range_excl_rightI} with "backward_derive".
But then the type derivator gets called with 0 \<le> i as a goal, which is no good.

Maybe there's a good way to incorporate auto/simp calls for non-type
premises in typing rules without making everything blow up.
*)
  have "i \<in> [0,\<dots>,m[" and "j \<in> [0,\<dots>,n[" by auto
  then show ?thesis unfolding Matrix_zero_def by auto
qed

definition "Matrix_Zero Z m n \<equiv> object {
    \<langle>@zero, Matrix_zero Z m n\<rangle>
  }"

lemma Matrix_Zero_type: assumes "Z : Zero A" "m : Nat" "n : Nat"
  shows "Matrix_Zero Z m n : Zero (Matrix A m n)"
  unfolding Matrix_Zero_def by (rule ZeroI) auto

subsection \<open>One\<close>

definition "Matrix_one Z O m n \<equiv>
  \<lambda>i \<in> [0,\<dots>,m[. \<lambda>j \<in> [0,\<dots>,n[. if i = j then one O else zero Z"

lemma Matrix_one_type [type]:
  "Matrix_one : Zero A \<Rightarrow> One A \<Rightarrow> (m : Nat) \<Rightarrow> (n : Nat) \<Rightarrow> Matrix A m n"
  unfolding Matrix_one_def by discharge_types

lemma Matrix_one_eq_one [simp]:
  assumes "i : Nat"
  and i_lt_m: "i < m"
  and i_lt_n: "i < n"
  shows "(Matrix_one Z O m n)`i`i = one O"
proof -
  have "i \<in> [0,\<dots>,m[" and "i \<in> [0,\<dots>,n[" by auto
  then show ?thesis unfolding Matrix_one_def by auto
qed

lemma Matrix_one_eq_zero [simp]:
  assumes "i : Nat" "j : Nat"
  and i_lt_m: "i < m"
  and j_lt_n: "j < n"
  and i_ne_j: "i \<noteq> j"
  shows "(Matrix_one Z O m n)`i`j = zero Z"
proof -
  have "i \<in> [0,\<dots>,m[" and "j \<in> [0,\<dots>,n[" by auto
  with i_ne_j show ?thesis unfolding Matrix_one_def by auto
qed

definition "Matrix_One Z O m n \<equiv> object {
    \<langle>@one, Matrix_one Z O m n\<rangle>
  }"

lemma Matrix_One_type:
  assumes "Z : Zero A" "O : One A" "m : Nat" "n : Nat"
  shows "Matrix_One Z O m n : One (Matrix A m n)"
  unfolding Matrix_One_def by (rule OneI) simp

subsection \<open>Addition\<close>

definition "Matrix_add A m n M N \<equiv>
  \<lambda>i \<in> [0,\<dots>,m[. \<lambda>j \<in> [0,\<dots>,n[. add A (M`i`j) (N`i`j)"

lemma Matrix_add_type [type]:
  "Matrix_add : Add C \<Rightarrow> (m : Nat) \<Rightarrow> (n : Nat) \<Rightarrow> Matrix C m n \<Rightarrow> Matrix C m n
    \<Rightarrow> Matrix C m n"
  unfolding Matrix_add_def using [[type_derivation_depth=4]]
  by discharge_types

lemma Matrix_add_eq_add [simp]:
  assumes "i : Nat" "j : Nat"
  and i_lt_m: "i < m"
  and j_lt_n: "j < n"
  shows "(Matrix_add A m n M N)`i`j = add A (M`i`j) (N`i`j)"
proof -
  have "i \<in> [0,\<dots>,m[" and "j \<in> [0,\<dots>,n[" by auto
  then show ?thesis unfolding Matrix_add_def by auto
qed

(*Note Kevin: or one could do the following:*)
(* declare [[coercion_enabled]] [[coercion "eval"]]

definition "Matrix_add' a m n (M :: set) (N :: set) \<equiv>
  \<lambda>i \<in> [0,\<dots>,m[. \<lambda>j \<in> [0,\<dots>,n[. add a (M i j) (N i j)"

declare [[coercion "Element"]]

lemma Matrix_add'_type [type]: "Matrix_add' : Add A \<Rightarrow> (m : Nat) \<Rightarrow> (n : Nat) \<Rightarrow>
  Matrix A m n \<Rightarrow> Matrix A m n \<Rightarrow> Matrix A m n"
  using [[type_derivation_depth=5]]
  unfolding Matrix_def Matrix_add'_def by discharge_types *)

definition "Matrix_Add C A m n \<equiv> object {
    \<langle>@add, \<lambda>M N \<in> matrices C m n. Matrix_add A m n M N\<rangle>
  }"

lemma Matrix_Add_type : assumes "A : Add (Element C)" "m : Nat" "n : Nat"
  shows "Matrix_Add C A m n : Add (Matrix (Element C) m n)"
proof -
  have [derive]: "\<And>M. M : Element (matrices C m n) \<Longrightarrow> M : Matrix (Element C) m n"
    by (drule ElementD) discharge_types
  have 2 [derive]: "\<And>M. M : Matrix (Element C) m n \<Longrightarrow> M : Element (matrices C m n)"
    by (rule ElementI) discharge_types
  then have "\<lambda>M N \<in> matrices C m n. Matrix_add A m n M N :
    Element (matrices C m n) \<rightarrow>cs Element (matrices C m n) \<rightarrow>cs
      Matrix (Element C) m n"
    by discharge_types
  then have "\<lambda>M N \<in> matrices C m n. Matrix_add A m n M N :
    Element (matrices C m n) \<rightarrow>s Element (matrices C m n) \<rightarrow>cs Matrix (Element C) m n"
    by (elim Dep_Function_if_CDep_Function)
  then have "\<lambda>M N \<in> matrices C m n. Matrix_add A m n M N :
    Element (matrices C m n) \<rightarrow>s Element (matrices C m n) \<rightarrow>s Matrix (Element C) m n"
    by (rule Dep_Function_covariant_codom, intro Dep_Function_if_CDep_Function)
  then have "\<lambda>M N \<in> matrices C m n. Matrix_add A m n M N :
    Matrix (Element C) m n \<rightarrow>s Element (matrices C m n) \<rightarrow>s Matrix (Element C) m n"
    by (elim Dep_Function_contravariant_dom) discharge_types
  then have "\<lambda>M N \<in> matrices C m n. Matrix_add A m n M N :
    Matrix (Element C) m n \<rightarrow>s Matrix (Element C) m n \<rightarrow>s Matrix (Element C) m n"
    by (elim Dep_Function_covariant_codom[OF Dep_Function_contravariant_dom])
      simp_all
  (* TODO Kevin: why is this selector not simplified automatically? *)
  have sel_simp:
    "(Matrix_Add C A m n)@@add = \<lambda>M N \<in> matrices C m n. Matrix_add A m n M N"
    unfolding Matrix_Add_def by simp
  show ?thesis by (intro AddI, subst sel_simp) fact
qed

subsection \<open>Additive Monoid\<close>

lemma Matrix_add_zero:
  assumes "M : Monoid (Element C)" "N : Matrix (Element C) m n"
  shows "Matrix_add M m n N (Matrix_zero M m n) = N"
  using assms using [[type_derivation_depth=4]]
  unfolding Matrix_add_def Matrix_zero_def
  apply (intro lambda_ext)
  apply (subst mem_dep_functions_iff_CDep_Function)
  apply (subst CDep_Function_covariant_codom)
  apply (simp only: Matrix_def)
  apply (rule ElementI)
  apply (subst mem_dep_functions_iff_CDep_Function)
  apply (auto simp only: eval_lambda_eq Monoid_add_zero_eq)
  done

lemma Matrix_zero_add:
  assumes "M : Monoid (Element C)" "N : Matrix (Element C) m n"
  shows "Matrix_add M m n (Matrix_zero M m n) N = N"
  using assms using [[type_derivation_depth=4]]
  unfolding Matrix_add_def Matrix_zero_def
  apply (intro lambda_ext)
  apply (subst mem_dep_functions_iff_CDep_Function)
  apply (subst CDep_Function_covariant_codom)
  apply (simp only: Matrix_def)
  apply (rule ElementI)
  apply (subst mem_dep_functions_iff_CDep_Function)
  apply (auto simp only: eval_lambda_eq Monoid_zero_add_eq)
  done

(*FIXME*)
lemma Matrix_add_assoc:
  assumes "M : Monoid (Element C)" "N : Matrix (Element C) m n"
    "O : Matrix (Element C) m n" "P : Matrix (Element C) m n"
  shows "Matrix_add M m n (Matrix_add M m n N O) P =
    Matrix_add M m n N (Matrix_add M m n O P)"
  using [[quick_and_dirty]]
  sorry
  (* using [[type_derivation_depth=4]] assms
  unfolding Matrix_add_def Matrix_def
  by (auto 0 0 intro!: lambda_ext simp: Nat_add_assoc) *)

definition "Matrix_Monoid C M m n \<equiv> object {
  \<langle>@zero, Matrix_zero M m n\<rangle>,
  \<langle>@add, \<lambda>N O \<in> matrices C m n. Matrix_add M m n N O\<rangle>
}"

(*TODO Kevin: Create object extension method so that one can re-use the proofs
from Matrix_Add_type and Matrix_Zero_type instead of unfolding and
proving everything again (cf branch object_extend).*)
lemma assumes "M : Monoid (Element C)" "m : Nat" "n : Nat"
  shows "Matrix_Monoid C M m n : Monoid (Matrix (Element C) m n)"
  using [[quick_and_dirty]]
  sorry
(* proof -
  have
    sel_add: "(Matrix_Monoid C M m n)@@add = \<lambda>N O \<in> matrices C m n. Matrix_add M m n N O"
    unfolding Matrix_Monoid_def by simp
  show ?thesis
    (* by (intro MonoidI ZeroI AddI; (subst sel_add)?)
    (auto simp: Matrix_Monoid_def Matrix_add_zero Matrix_zero_add
    Matrix_add_assoc add_def zero_def mem_dep_functions_iff_CDep_Function
    intro!: Dep_fun_typeI) *)
  (* proof (rule MonoidI)
    have "(Matrix_Monoid C M m n)@@add = (Matrix_Add C M m n)@@add"
      unfolding Matrix_Monoid_def Matrix_Add_def by simp
  qed *)
qed *)

subsection \<open>Multiplication\<close>

text \<open>Multiplying an l \<times> 0 with a 0 \<times> n Matrix returns the l \<times> n zero Matrix.\<close>
definition "Matrix_mul A M l m n N O \<equiv> \<lambda>i \<in> [0,\<dots>,l[. \<lambda>j \<in> [0,\<dots>,n[. nat_rec'
  m (zero A) (\<lambda>k acc. add A acc (mul M (N`i`(pred k)) (O `(pred k) `j)))"

(*Note Kevin: TODO: type derivator is not able to handle this automatically
yet*)
lemma Matrix_mul_type [type]: "Matrix_mul : Monoid C \<Rightarrow> Mul C \<Rightarrow> (l : Nat) \<Rightarrow>
  (m : Nat) \<Rightarrow> (n : Nat) \<Rightarrow> Matrix C l m \<Rightarrow> Matrix C m n \<Rightarrow> Matrix C l n"
proof (intro Dep_fun_typeI)
  fix AddM M l m n N O
  assume "AddM : Monoid C" "M : Mul C" "l : Nat" "m : Nat" "n : Nat"
    "N : Matrix C l m" "O : Matrix C m n"
  {
    fix i j assume "i : Element [0,\<dots>,l[" "j : Element [0,\<dots>,n["
    have "pred : Element [1,\<dots>,m] \<Rightarrow> Element [0,\<dots>,m["
    proof unfold_types
      fix n assume "n \<in> [1,\<dots>,m]"
      then have "n \<in> \<nat>" by (fact mem_nat_if_mem_range_incl_excl)
      have "pred n < m"
      proof -
        from \<open>n \<in> [1,\<dots>,m]\<close> have "0 \<noteq> n" "n \<le> m"
          unfolding nat_one_def by (auto elim: mem_rangeE)
        then show ?thesis by (auto intro: Nat_pred_lt_if_le_if_ne_zero)
      qed
      then show "pred n \<in> [0,\<dots>,m[" by auto
    qed
    have "(\<lambda>k acc. add AddM acc (mul M (N`i`(pred k)) (O`(pred k)`j)))
      : Element [1,\<dots>,m] \<Rightarrow> C \<Rightarrow> C"
      using [[type_derivation_depth=6]] by discharge_types
  }
  show "Matrix_mul AddM M l m n N O : Matrix C l n"
    using [[type_derivation_depth=4]]
    unfolding Matrix_mul_def by discharge_types
qed

definition "Matrix_Mul C A M l m n \<equiv> object {
    \<langle>@mul, \<lambda>N \<in> matrices C l m. (\<lambda>O \<in> matrices C m n. Matrix_mul A M l m n N O)\<rangle>
  }"

lemma Matrix_Mul_type:
  assumes "A : Monoid (Element C)" "M : Mul (Element C)" "n : Nat"
  shows "Matrix_Mul C A M n n n : Mul (Matrix (Element C) n n)"
proof -
  have [derive]: "\<And>M. M : Element (matrices C n n) \<Longrightarrow> M : Matrix (Element C) n n"
    by (drule ElementD) discharge_types
  have 2 [derive]: "\<And>M. M : Matrix (Element C) n n \<Longrightarrow> M : Element (matrices C n n)"
    by (rule ElementI) discharge_types
  let ?f = "\<lambda>N \<in> matrices C n n. (\<lambda>O \<in> matrices C n n. Matrix_mul A M n n n N O)"
  have "?f : Element (matrices C n n) \<rightarrow>cs Element (matrices C n n) \<rightarrow>cs
      Matrix (Element C) n n"
    by discharge_types
  then have "?f : Element (matrices C n n) \<rightarrow>s Element (matrices C n n) \<rightarrow>cs
    Matrix (Element C) n n"
    by (elim Dep_Function_if_CDep_Function)
  then have "?f : Element (matrices C n n) \<rightarrow>s Element (matrices C n n) \<rightarrow>s
    Matrix (Element C) n n"
    by (rule Dep_Function_covariant_codom, intro Dep_Function_if_CDep_Function)
  then have "?f : Matrix (Element C) n n \<rightarrow>s Element (matrices C n n) \<rightarrow>s
    Matrix (Element C) n n"
    by (elim Dep_Function_contravariant_dom) discharge_types
  then have "?f : Matrix (Element C) n n \<rightarrow>s Matrix (Element C) n n \<rightarrow>s
    Matrix (Element C) n n"
    by (elim Dep_Function_covariant_codom[OF Dep_Function_contravariant_dom])
      simp_all
  (* TODO Kevin: why is this selector not simplified automatically? *)
  have sel_simp: "(Matrix_Mul C A M n n n)@@mul =
    \<lambda>N \<in> matrices C n n. (\<lambda>O \<in> matrices C n n. Matrix_mul A M n n n N O)"
    unfolding Matrix_Mul_def by simp
  show ?thesis by (intro MulI, subst sel_simp) fact
qed

subsection \<open>Multiplicative Monoid\<close>

(*Note: This could be generalised to non-square matrices, but we do not need
that for now. *)
lemma
  assumes "A : Monoid C" "M : Mul_Monoid C" "n : Nat" "N : Matrix C n n"
  and mul_zero: "\<And>c. c : C \<Longrightarrow> mul M c (zero A) = zero A"
  and mul_one: "\<And>c. c : C \<Longrightarrow> mul M c (one M) = c"
  shows Matrix_mul_one: "Matrix_mul A M n n n N (Matrix_one A M n n) = N"
using mem_range_incl_exclE[elim]
(*FIXME*) [[quick_and_dirty]]
unfolding Matrix_mul_def
proof (intro lambda_ext)
  fix i j assume i_mem: "i \<in> [0,\<dots>,n[" and j_mem: "j \<in> [0,\<dots>,n["
  let ?f = "\<lambda>k acc. add A acc (mul M (N`i`(pred k)) (Matrix_one A M n n `(pred k) `j))"
  {
    fix m assume lassms: "m : Nat" "m < n"
    with i_mem have "N`i`m : C" by (intro Matrix_eval_typeI) auto
    with lassms have
      "nat_rec' (succ m) (zero A) ?f = (if m < j then zero A else N`i`j)"
    proof (induction m rule: Nat_induct)
      case zero
      then show ?case
      proof (cases "0 < j")
        case True
        with j_mem have "Matrix_one A M n n `0 `j = zero A"
          by (intro Matrix_one_eq_zero) auto
        with mul_zero show ?thesis using mul_zero True by auto
      next
        case False
        with j_mem have "j = 0" by (auto elim: leE)
        moreover with j_mem have "Matrix_one A M n n `0 `0 = one M"
          by (intro Matrix_one_eq_one) auto
        ultimately show ?thesis using mul_one by auto
      qed
    next
      case (succ m)
      then have "m < n" by (auto intro: Nat_lt_if_succ_lt)
      with i_mem have "N`i`m : C" by (intro Matrix_eval_typeI) auto
      with succ.IH have
        IH: "nat_rec' (succ m) (zero A) ?f = (if m < j then zero A else N`i`j)"
        by auto
      show ?case
      proof (cases "succ m < j")
        case True
        (* Note Kevin: this is BAD *)
        from j_mem have "j : Nat" by auto
        then have "m < j" using Nat_lt_if_succ_lt[OF _ \<open>succ m < j\<close>] by blast
        moreover with True j_mem have "Matrix_one A M n n `(succ m) `j = zero A"
          by (intro Matrix_one_eq_zero) auto
        ultimately show ?thesis using IH mul_zero True by auto
      next
        case False
        from j_mem have "j : Nat" "m : Nat" by auto
        then have f: "m < j \<and> j < succ m \<Longrightarrow> False" using
          Nat_succ_le_if_lt[of m j] Nat_not_lt_iff_le sorry
        have "j \<le> succ m" by (rule Nat_le_if_not_lt) simp
        then show ?thesis
        proof (cases rule: leE)
          case lt
          from i_mem j_mem have "N`i`j : C" by (intro Matrix_eval_typeI) auto
          from lt j_mem have "Matrix_one A M n n `(succ m) `j = zero A"
            by (intro Matrix_one_eq_zero) auto
          with f lt IH mul_zero show ?thesis using lt_asym by auto
        next
          case eq
          with IH mul_one show ?thesis by auto
        qed
      qed
    qed
  } note rec = this
  moreover from \<open>i \<in> [0,\<dots>,n[\<close> have "i < n" "i \<in> \<nat>" by auto
  (* case n=0 needed *)
  ultimately show "nat_rec' n (zero A) ?f = N`i`j"
  proof (cases "n = 0")
    case True
    with \<open>i < n\<close> show ?thesis by auto
  next
    case False
    then obtain m where n_eq_succ_m: "n = succ m" by (auto intro: mem_natE[of n])
    then have "m : Nat" "m < n" using \<open>n : Nat\<close> sorry
    then have
      lem: "nat_rec' (succ m) (zero A) ?f = (if m < j then zero A else N`i`j)"
      by (fact rec)
    have "j \<le> pred (succ m)"
      by (rule Nat_le_pred_if_lt) (insert j_mem n_eq_succ_m[symmetric], auto)
    then have "j \<le> m" by simp
    moreover from j_mem have "j : Nat" by auto
    ultimately have "\<not> m < j" using Nat_lt_if_lt_if_le[of j j m] by auto
    then show ?thesis using n_eq_succ_m lem by auto
  qed
  (*remaining type assumptions as in Matrix_Add_type and Matrix_Mul_Type*)
oops

end


theory Llist_of_stream_unfold_stream2
  imports Main "$HIPSTER_HOME/IsaHipster"
begin

setup Tactic_Data.set_coinduct_sledgehammer

codatatype (lset: 'a) Llist =
    lnull: LNil
    | LCons (lhd: 'a) (ltl: "'a Llist")
where
 "ltl LNil = LNil"

codatatype (sset: 'a) Stream =
  SCons (shd: 'a) (stl: "'a Stream")

primcorec llist_of_stream2 :: "'a Stream \<Rightarrow> 'a Llist"
  where "llist_of_stream2 s = LCons (shd s) (llist_of_stream2 (stl s))"

primcorec stream_of_llist2 :: "'a Llist \<Rightarrow> 'a Stream"
where "stream_of_llist2 l = SCons (lhd l) (stream_of_llist2 (ltl l))"
primcorec unfold_llist :: "('a \<Rightarrow> bool) \<Rightarrow> ('a \<Rightarrow> 'b) \<Rightarrow> ('a \<Rightarrow> 'a) \<Rightarrow> 'a \<Rightarrow> 'b Llist" where
  "p a \<Longrightarrow> unfold_llist p f g a = LNil" |
  "_ \<Longrightarrow> unfold_llist p f g a = LCons (f a) (unfold_llist p f g (g a))"

primcorec unfold_stream :: "('a \<Rightarrow> 'b) \<Rightarrow> ('a \<Rightarrow> 'a) \<Rightarrow> 'a \<Rightarrow> 'b Stream" where
  "unfold_stream g1 g2 a = SCons (g1 a) (unfold_stream g1 g2 (g2 a))"

datatype 'a Lst =
  Emp
  | Cons "'a" "'a Lst"

fun obsStream :: "int \<Rightarrow> 'a Stream \<Rightarrow> 'a Lst" where
"obsStream n s = (if (n \<le> 0) then Emp else Cons (shd s) (obsStream (n - 1) (stl s)))"
fun obsLlist :: "int \<Rightarrow> 'a Llist \<Rightarrow> 'a Lst" where
"obsLlist n s = (if (n \<le> 0) then Emp else Cons (lhd s) (obsLlist (n - 1) (ltl s)))"
(*hipster_obs Stream Lst obsStream stream_of_llist llist_of_stream
Times out, would need an observer for the llist also but can only
have one at a time now*)
(* hipster_obs Stream Lst obsStream stream_of_llist2
Finds nothing ? *)
(* hipster_obs Llist Lst obsLlist llist_of_stream2
This ran for a long time and returned nothing, made my computer v. slow and wonky
*)
theorem llist_of_stream_unfold_stream2:
  "llist_of_stream2 (unfold_stream shd stl x) = unfold_llist (\<lambda>_. False) shd stl x"
(*by hipster_coinduct_sledgehammer
Metis: Failed to replay Metis proof
resynchronize: Out of sync
Failed to apply initial proof method\<here>*)
  apply(coinduction arbitrary: x rule: Llist.coinduct_strong)
  by auto
end