subsection \<open>Binary orthogonality is Diophantine\<close>

theory Binary_Orthogonal
  imports Binomial_Coefficient "Digit_Expansions.Binary_Operations" "Lucas_Theorem.Lucas_Theorem"
begin

lemma equiv_with_lucas: "nth_digit = Lucas_Theorem.nth_digit_general"
  unfolding nth_digit_def Lucas_Theorem.nth_digit_general_def by simp

lemma lm0241_ortho_binom_equiv:"(a \<bottom> b) \<longleftrightarrow> odd ((a + b) choose b)" (is "?P \<longleftrightarrow> ?Q")
proof
  assume "?P"
  hence dig0:"(\<forall>i. (nth_bit a i) * (nth_bit b i) = 0)"
    using ortho_mult_equiv
    by auto
  hence "(\<forall>i. (nth_bit a i) * (nth_bit b i) \<noteq> 1)"
    by presburger
  hence dcons:"(\<forall>i. \<not>(((nth_bit a i) = 1) \<and> ((nth_bit b i) = 1)))"
    by auto
  hence "(\<forall>i. (bin_carry a b i) = 0)" using no_carry_mult_equiv dig0
    by blast
  hence dsum:"(\<forall>i. (nth_bit (a + b) i) = (nth_bit a i) + (nth_bit b i))"
    by (metis One_nat_def add.commute add_cancel_right_left add_self_mod_2
        dig0 mult_is_0 not_mod2_eq_Suc_0_eq_0 nth_bit_def one_mod_two_eq_one
        sum_digit_formula)

  have bdd_ab_exists:"(\<exists>p. (a + b) < 2^(Suc p))"
    using aux1_lm0241_pow2_up_bound by auto
  then obtain p where bdd_ab:"(a + b) < 2^(Suc p)" by blast
  moreover from bdd_ab have "b < 2^(Suc p)" by auto

  ultimately have "((a + b) choose b) mod 2 =
        (\<Prod>i\<le>p. ((nth_digit (a + b) i 2) choose (nth_digit b i 2))) mod 2"
    using lucas_theorem[of "a+b" 2 p b] bdd_ab two_is_prime_nat
    by (simp add: equiv_with_lucas)

  then have a_choose_b_digit_prod: "((a + b) choose b) mod 2 =
        (\<Prod>i\<le>p. ((nth_bit (a + b) i) choose (nth_bit b i))) mod 2"
    using nth_digit_base2_equiv
    by (auto cong: prod.cong)

  have "(\<forall>i. ((nth_bit (a + b) i) choose (nth_bit b i) = 1))"
    using aux2_lm0241_single_digit_binom[where ?a="nth_bit (a + b) i"
                                               and ?b="nth_bit b i"]
    by (metis add.commute add.right_neutral binomial_n_0 binomial_n_n dig0 dsum
        mult_is_0)
  hence f0:"1 = (\<Prod>i<p. (nth_bit (a + b) i) choose (nth_bit b i))"
    by simp
  hence f1:"... = ... mod 2" by simp
  hence f2:"... = ((a + b) choose b) mod 2"
    using a_choose_b_digit_prod by (simp add: \<open>\<forall>i. (a + b) \<exclamdown> i choose b \<exclamdown> i = 1\<close>)
  then show "?Q" using f0 by fastforce
next
  assume "?Q"
  hence a_choose_b_odd:"((a + b) choose b) mod 2 = 1"
    using odd_iff_mod_2_eq_one by blast

  have bdd_ab_exists:"(\<exists>p. (a + b) < 2^(Suc p))"
    using aux1_lm0241_pow2_up_bound by auto
  then obtain p where bdd_ab:"(a + b) < 2^(Suc p)" by blast
  moreover from bdd_ab have bdd_b: "b < 2^(Suc p)" by auto

  ultimately have "((a + b) choose b) mod 2 =
        (\<Prod>i\<le>p. ((nth_digit (a + b) i 2) choose (nth_digit b i 2))) mod 2"
    using lucas_theorem[of "a+b" 2 p b] bdd_ab two_is_prime_nat
    by (simp add: equiv_with_lucas)

  then have a_choose_b_digit_prod: "((a + b) choose b) mod 2 =
        (\<Prod>i\<le>p. ((nth_bit (a + b) i) choose (nth_bit b i))) mod 2"
    using nth_digit_base2_equiv nth_digit_def
    by (auto cong: prod.cong)
  hence all_prod_one_mod2:"... = 1" using a_choose_b_odd by linarith

  have choose_bdd:"(\<forall>i. 1 \<ge> (nth_bit (a + b) i) choose (nth_bit b i))"
    using nth_bit_bounded
    by (metis One_nat_def binomial_n_0 choose_one not_mod2_eq_Suc_0_eq_0
        nth_bit_def order_refl)
  hence "1 \<ge> (\<Prod>i\<le>p. ((nth_bit (a + b) i) choose (nth_bit b i)))"
    using all_prod_one_mod2 by (meson prod_le_1 zero_le)
  hence "(\<Prod>i\<le>p. ((nth_bit (a + b) i) choose (nth_bit b i))) mod 2 =
             (\<Prod>i\<le>p. ((nth_bit (a + b) i) choose (nth_bit b i)))"
    using all_prod_one_mod2 by linarith
  hence "... = 1"
    using all_prod_one_mod2 by linarith
  hence sub_pq_one:"\<forall>i\<le>p. (nth_bit (a + b) i) choose (nth_bit b i) = 1"
    using
      aux4_lm0241_prod_one[where ?n="p" and ?f="(\<lambda>i. (nth_bit (a + b) i) choose (nth_bit b i))"]
      choose_bdd by blast

  have "\<forall>r > p. (a + b) < 2^r" using bdd_ab
    by(metis One_nat_def Suc_lessI lessI less_imp_add_positive numeral_2_eq_2
       power_strict_increasing_iff trans_less_add1)
  hence "\<forall>r > p. r \<ge> p \<longrightarrow> (a + b) < 2^r" by auto
  hence ab_digit_bdd:"\<forall>r > p. r \<ge> p \<longrightarrow> nth_bit (a + b) r = 0"
    using nth_bit_def by simp

  have "\<forall>k > p. b < 2 ^ k" using bdd_b
    by(metis One_nat_def Suc_lessI lessI less_imp_add_positive numeral_2_eq_2
       power_strict_increasing_iff trans_less_add1)
  hence b_digit_bdd:"\<forall>k > p. k \<ge> p \<longrightarrow> nth_bit b k = 0"
    using nth_bit_def
    by (simp add: \<open>\<forall>k>p. b < 2 ^ k\<close>)

  from b_digit_bdd ab_digit_bdd aux3_lm0241_binom_bounds
  have "\<forall>i. i \<ge> p \<longrightarrow> (nth_bit (a + b) i) choose (nth_bit b i) = 1"
    by (simp add: le_less sub_pq_one)

  hence "\<forall>i. ((nth_bit (a + b) i) choose (nth_bit b i)) = 1"
    using sub_pq_one not_less by (metis linear)
  hence "\<forall>i. \<not>(nth_bit a i = 1 \<and> nth_bit b i = 1)" using aux5_lm0241 by blast
  hence "\<forall>i. ((nth_bit a i = 0 \<and> nth_bit b i = 1) \<or>
                  (nth_bit a i = 1 \<and> nth_bit b i = 0) \<or>
                  (nth_bit a i = 0 \<and> nth_bit b i = 0))"
    by (auto simp add:nth_bit_def nth_digit_bounded; metis nat.simps(3))
  hence "\<forall>i. (nth_bit a i) * (nth_bit b i) = 0" using mult_is_0 by blast
  then show "?P" using ortho_mult_equiv by blast
qed

definition orthogonal (infix "[\<bottom>]" 50)
  where "P [\<bottom>] Q \<equiv> (BINARY (\<lambda>a b. a \<bottom> b) P Q)"

lemma orthogonal_dioph[dioph]:
  fixes P Q
  defines "DR \<equiv> (P [\<bottom>] Q)"
  shows "is_dioph_rel DR"
proof -
  define P' Q' where pushed_def: "P' \<equiv> push_param P 1" "Q' \<equiv> push_param Q 1"

  define DS where "DS \<equiv> [\<exists>] [Param 0 = (P' [+] Q') choose Q'] [\<and>] ODD (Param 0)"

  have "eval DS a = eval DR a" for a
    unfolding DS_def DR_def orthogonal_def pushed_def defs
    using push_push1 lm0241_ortho_binom_equiv by (simp add: push0)

  moreover have "is_dioph_rel DS"
    unfolding DS_def by (simp add: dioph)

  ultimately show ?thesis
    by (auto simp: is_dioph_rel_def)
qed

declare orthogonal_def[defs]

end


(*
  Title:      QR_Efficient.thy
  Author:     Jose Divasón <jose.divasonm at unirioja.es>
  Author:     Jesús Aransay <jesus-maria.aransay at unirioja.es>
*)

section\<open>Improvements to get better performance of the algorithm\<close>

theory QR_Efficient
imports QR_Decomposition_IArrays
begin

subsection\<open>Improvements for computing the Gram Schmidt algorithm and QR decomposition using vecs\<close>

text\<open>Essentialy, we try to avoid removing duplicates in each iteration.
  They will not affect the @{const sum_list} since the duplicates will be the vector zero.\<close>

subsubsection\<open>New definitions\<close>

definition "Gram_Schmidt_column_k_efficient A k
  = (\<chi> a b. (if b = from_nat k
  then column b A - sum_list (map (\<lambda>x. ((column b A \<bullet> x) / (x \<bullet> x)) *\<^sub>R x)
  ((map (\<lambda>n. column (from_nat n) A) [0..<to_nat b]))) else column b A) $ a)"

subsubsection\<open>General properties about @{const sum_list}\<close>

lemma sum_list_remdups_2:
  fixes f:: "'a::{zero, monoid_add}\<Rightarrow>'a"
  assumes "!!i j. i<length xs \<and> j<length xs \<and> i \<noteq> j \<and> (xs ! i) = (xs ! j)
    \<longrightarrow> f (xs ! i) = 0 \<and> f (xs ! j) = 0"
  shows "sum_list (map f (remdups xs)) = sum_list (map f xs)"
  using assms
proof (induct xs)
  case Nil
  show ?case by auto
next
  case (Cons a xs)
  show ?case
  proof (cases "a \<in> set xs")
    case False
    hence "sum_list (map f (remdups (a # xs))) =  sum_list (map f (a # (remdups xs)))"
      by simp
    also have "... = sum_list (f a # (map f (remdups xs)))" by auto
    also have "... = f a + sum_list (map f (remdups xs))" by auto
    also have "... = f a + sum_list (map f xs)" using Cons.prems Cons.hyps
      using id_apply by fastforce
    also have "... =  sum_list (map f (a # xs))" by auto
    finally show ?thesis .
  next
    case True
    have fa_0: "f a = 0" using Cons.hyps Cons.prems True
      by (metis Suc_less_eq add.right_neutral add_Suc_right add_gr_0
        in_set_conv_nth lessI list.size(4) nth_Cons_0 nth_Cons_Suc)
    have "sum_list (map f (remdups (a # xs))) =  sum_list (map f (remdups xs))"
      using True by simp
    also have "... = sum_list (map f xs)" using Cons.prems Cons.hyps
      using id_apply by fastforce
    also have "... = f a + sum_list (map f xs)" using fa_0 by simp
    also have "... = sum_list (map f (a # xs))" by auto
    finally show ?thesis .
  qed
qed

subsubsection\<open>Proving a code equation to improve the performance\<close>

lemma set_map_column:
  "set (map (\<lambda>n. column (from_nat n) G) [0..<to_nat b]) =  {column i G|i. i<b}"
proof (auto)
  fix n assume "n < to_nat b"
  hence "from_nat n < b" using from_nat_mono to_nat_less_card by fastforce
  thus "\<exists>i. column (from_nat n) G = column i G\<and> i < b" by auto
next
  fix i assume "i < b" hence ib: "to_nat i < to_nat b" by (simp add: to_nat_le)
  show "column i G \<in> (\<lambda>n. column (from_nat n) G) ` {0..<to_nat b}"
    unfolding image_def
    by (auto, rule bexI[of _ "to_nat i"], auto simp add: ib)
qed

lemma column_Gram_Schmidt_column_k_repeated_0:
  fixes A::"'a::{real_inner}^'n::{mod_type}^'m::{mod_type}"
  assumes i_not_k: "i\<noteq>k" and ik: "i<k"
  and c_eq: "column k (Gram_Schmidt_column_k A (to_nat k))
  = column i (Gram_Schmidt_column_k A (to_nat k))"
  and o: "pairwise orthogonal {column i A|i. i<k}"
  shows "column k (Gram_Schmidt_column_k A (to_nat k)) = 0"
  and "column i (Gram_Schmidt_column_k A (to_nat k)) = 0"
proof -
  have "column k (Gram_Schmidt_column_k A (to_nat k))
    = column k A - (\<Sum>x\<in>{column i A |i. i < k}. (x \<bullet> column k A / (x \<bullet> x)) *\<^sub>R x)"
    by (rule column_Gram_Schmidt_column_k)
  also have "... = column k A - proj_onto (column k A) {column i A |i. i < k}"
    unfolding proj_onto_def proj_def[abs_def]
    by (metis (no_types, lifting) inner_commute)
  finally have col_k_rw: "column k (Gram_Schmidt_column_k A (to_nat k))
    = column k A - proj_onto (column k A) {column i A |i. i < k}" .
  have "orthogonal (column k (Gram_Schmidt_column_k A (to_nat k)))
    (column i (Gram_Schmidt_column_k A (to_nat k)))"
    unfolding col_k_rw
  proof (rule orthogonal_proj_set[OF _ _ o])
    have "column i (Gram_Schmidt_column_k A (to_nat k)) = column i A"
      using column_Gram_Schmidt_column_k'[OF i_not_k] .
    also have "...  \<in> {column i A |i. i < k}" using assms(2) by blast
    finally show "column i (Gram_Schmidt_column_k A (to_nat k)) \<in> {column i A |i. i < k}" .
    show "finite {column i A |i. i < k}" by auto
  qed
  thus "column k (Gram_Schmidt_column_k A (to_nat k)) = 0"
    and "column i (Gram_Schmidt_column_k A (to_nat k)) = 0"
    unfolding orthogonal_def c_eq inner_eq_zero_iff by auto
qed

lemma column_Gram_Schmidt_upt_k_repeated_0':
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  assumes i_not_k: "i\<noteq>j" and ij: "i<j" and j: "j\<le>from_nat k"
  and c_eq: "column j (Gram_Schmidt_upt_k A k)
  = column i (Gram_Schmidt_upt_k A k)"
  and k: "k<ncols A"
  shows "column j (Gram_Schmidt_upt_k A k) = 0"
  using j c_eq k
proof (induct k)
  case 0
  thus ?case
    using ij least_mod_type
    unfolding from_nat_0
    by (metis (no_types) dual_order.strict_trans1 ij least_mod_type not_less)
next
  case (Suc k)
  have k: "k<ncols A" using Suc.prems unfolding ncols_def by auto
  have i_not_k: "i\<noteq>from_nat (Suc k)" using ij Suc.prems by auto
  have col_i_rw: "column i (Gram_Schmidt_upt_k A (Suc k)) = column i (Gram_Schmidt_upt_k A k)"
    by (simp add: i_not_k Gram_Schmidt_column_k_def Gram_Schmidt_upt_k_suc column_def)
  have tn_fn_suc: "to_nat (from_nat (Suc k)::'n) = Suc k"
    using to_nat_from_nat_id Suc.prems
    unfolding ncols_def by blast
  show ?case
  proof (cases "j=from_nat (Suc k)")
    case False
    have jk: "j \<le> from_nat k"
      by (metis False One_nat_def Suc.prems(1) add.right_neutral add_Suc_right
        from_nat_suc le_Suc less_le linorder_not_le)
    have col_j_rw: "column j (Gram_Schmidt_upt_k A (Suc k)) = column j (Gram_Schmidt_upt_k A k)"
      by (simp add: False Gram_Schmidt_column_k_def Gram_Schmidt_upt_k_suc column_def)
    have col_j_eq_col_i_k: "column j (Gram_Schmidt_upt_k A k) = column i (Gram_Schmidt_upt_k A k)"
      using Suc.prems unfolding col_j_rw col_i_rw by simp
    show ?thesis using Suc.hyps col_j_eq_col_i_k k jk unfolding col_j_rw by blast
  next
    case True
    show ?thesis unfolding True unfolding Gram_Schmidt_upt_k_suc
    proof (rule column_Gram_Schmidt_column_k_repeated_0(1)
        [of i "from_nat (Suc k)" "Gram_Schmidt_upt_k A k", unfolded tn_fn_suc])
      have set_rw: "{column i (Gram_Schmidt_upt_k A k) |i. i < from_nat (Suc k)}
        = {column i (Gram_Schmidt_upt_k A k) |i. to_nat i \<le> k}"
        by (metis (mono_tags, hide_lams) less_Suc_eq_le less_le not_less tn_fn_suc to_nat_mono)
      show "i \<noteq> from_nat (Suc k)" using i_not_k .
      show "i < from_nat (Suc k)" using True ij by blast
      show "column (from_nat (Suc k)) (Gram_Schmidt_column_k (Gram_Schmidt_upt_k A k) (Suc k)) =
        column i (Gram_Schmidt_column_k (Gram_Schmidt_upt_k A k) (Suc k))"
        using Suc.prems True by (simp add: Gram_Schmidt_upt_k_suc)
      show "pairwise orthogonal {column i (Gram_Schmidt_upt_k A k) |i. i < from_nat (Suc k)}"
        unfolding set_rw by (rule orthogonal_Gram_Schmidt_upt_k[OF k])
    qed
  qed
qed

lemma column_Gram_Schmidt_upt_k_repeated_0:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  assumes i_not_k: "i\<noteq>j" and ij: "i<j" and j: "j\<le>k"
  and c_eq: "column j (Gram_Schmidt_upt_k A (to_nat k))
  = column i (Gram_Schmidt_upt_k A (to_nat k))"
  shows "column j (Gram_Schmidt_upt_k A (to_nat k)) = 0"
  using assms column_Gram_Schmidt_upt_k_repeated_0' to_nat_less_card ncols_def
  by (metis c_eq column_Gram_Schmidt_upt_k_repeated_0'
     from_nat_to_nat_id i_not_k ij j ncols_def to_nat_less_card)

corollary column_Gram_Schmidt_upt_k_repeated:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  assumes i_not_k: "i\<noteq>j" and ij: "i\<le>k" and "j\<le>k"
  and c_eq: "column j (Gram_Schmidt_upt_k A (to_nat k))
    = column i (Gram_Schmidt_upt_k A (to_nat k))"
  shows "column j (Gram_Schmidt_upt_k A (to_nat k)) = 0"
  and "column i (Gram_Schmidt_upt_k A (to_nat k)) = 0"
proof -
  show "column j (Gram_Schmidt_upt_k A (to_nat k)) = 0"
  proof (cases "i<j")
    case True
    thus ?thesis using assms column_Gram_Schmidt_upt_k_repeated_0 by metis
  next
    case False hence ji: "j<i" using i_not_k by auto
    thus ?thesis using assms column_Gram_Schmidt_upt_k_repeated_0 by metis
  qed
  show "column i (Gram_Schmidt_upt_k A (to_nat k)) = 0"
  proof (cases "i<j")
    case True
    thus ?thesis using assms column_Gram_Schmidt_upt_k_repeated_0 by metis
  next
    case False hence ji: "j<i" using i_not_k by auto
    thus ?thesis using assms column_Gram_Schmidt_upt_k_repeated_0 by metis
  qed
qed

lemma column_Gram_Schmidt_column_k_eq_efficient:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  assumes "Gram_Schmidt_upt_k A k = foldl Gram_Schmidt_column_k_efficient A [0..<Suc k]"
  and suc_k: "Suc k < ncols A"
  shows "column b (Gram_Schmidt_column_k (Gram_Schmidt_upt_k A k) (Suc k))
  = column b (Gram_Schmidt_column_k_efficient (Gram_Schmidt_upt_k A k) (Suc k))"
proof (cases "b = from_nat (Suc k)")
  case False thus ?thesis
    unfolding Gram_Schmidt_column_k_efficient_def Gram_Schmidt_column_k_def column_def by auto
next
  case True
  have tn_fn_suc: "to_nat (from_nat (Suc k)::'n) = Suc k"
    using suc_k to_nat_from_nat_id unfolding ncols_def by auto
  define G where"G = Gram_Schmidt_upt_k A k"
  let ?f="(\<lambda>x. (column b G \<bullet> x / (x \<bullet> x)) *\<^sub>R x)"
  let ?g="(\<lambda>n. column (from_nat n) G) "
  have proj_eq: "proj_onto (column b G) {column i G |i. i < b}
    = sum_list (map ?f (map ?g [0..<to_nat b]))"
  proof -
    have "proj_onto (column b G) {column i G |i. i < b} = sum ?f {column i G |i. i < b}"
      unfolding proj_onto_def proj_def[abs_def] by simp
    also have "... = sum ?f (set (map ?g [0..<to_nat b]))"
      by (rule sum.cong, auto simp add: set_map_column[symmetric])
    also have "... = sum_list (map ?f (remdups (map ?g [0..<to_nat b])))" unfolding sum_code ..
    also have "... = sum_list ((map ?f ((map ?g [0..<to_nat b]))))"
    proof (rule sum_list_remdups_2, auto)
      fix i j assume i: "i < to_nat b"
        and j: "j < to_nat b" and ij: "i \<noteq> j"
        and col_eq: "column (from_nat i) G = column (from_nat j) G"
        and col_0: "column (from_nat j) G \<noteq> 0"
      have k: "to_nat (from_nat k::'n) = k"
        by (metis Suc_lessD ncols_def suc_k to_nat_from_nat_id)
      have "column (from_nat j) G = 0"
      proof (unfold G_def, rule column_Gram_Schmidt_upt_k_repeated(1)
          [of "(from_nat i)::'n" "from_nat j" "from_nat k" A, unfolded k])
        have "from_nat i < (from_nat (Suc k)::'n)"
          using from_nat_mono[of i "Suc k"] suc_k i
          unfolding True tn_fn_suc ncols_def by simp
        thus "from_nat i \<le> (from_nat k::'n)"
          by (metis Suc_lessD True from_nat_mono' i less_Suc_eq_le ncols_def suc_k tn_fn_suc)
        have "from_nat j < (from_nat (Suc k)::'n)"
          using from_nat_mono[of j "Suc k"] suc_k j
          unfolding True tn_fn_suc ncols_def by simp
        thus "from_nat j \<le> (from_nat k::'n)"
          by (metis Suc_lessD True from_nat_mono' j less_Suc_eq_le ncols_def suc_k tn_fn_suc)
        show "from_nat i \<noteq> (from_nat j::'n)" using ij i j True suc_k
          by (metis (no_types, lifting) dual_order.strict_trans from_nat_eq_imp_eq ncols_def tn_fn_suc)
        show "column (from_nat j) (Gram_Schmidt_upt_k A k)
          = column (from_nat i) (Gram_Schmidt_upt_k A k)" using G_def col_eq by auto
      qed
      thus "column b G \<bullet> column (from_nat j) G = 0" using col_0 by contradiction
    qed
    finally show ?thesis .
  qed
  have "column b (Gram_Schmidt_column_k G (Suc k))
    = column b G - proj_onto (column b G) {column i G |i. i < b}"
    unfolding True
    unfolding Gram_Schmidt_column_k_def G_def column_def by vector
  also have "... = column b G
    - (\<Sum>x\<leftarrow>map (\<lambda>n. column (from_nat n) G) [0..<to_nat b]. (column b G \<bullet> x / (x \<bullet> x)) *\<^sub>R x)"
    unfolding proj_eq ..
  also have "... = column b (Gram_Schmidt_column_k_efficient G (Suc k))"
    unfolding True Gram_Schmidt_column_k_efficient_def G_def column_def by vector
  finally show ?thesis unfolding G_def .
qed

lemma Gram_Schmidt_upt_k_efficient_induction:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  assumes "Gram_Schmidt_upt_k A k = foldl Gram_Schmidt_column_k_efficient A [0..<Suc k]"
  and suc_k: "Suc k < ncols A"
  shows "Gram_Schmidt_column_k (Gram_Schmidt_upt_k A k) (Suc k)
  = Gram_Schmidt_column_k_efficient (Gram_Schmidt_upt_k A k) (Suc k)"
  using column_Gram_Schmidt_column_k_eq_efficient[OF assms]
  unfolding column_def vec_eq_iff by vector

lemma Gram_Schmidt_upt_k_efficient:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  assumes k: "k<ncols A"
  shows "Gram_Schmidt_upt_k A k = foldl Gram_Schmidt_column_k_efficient A [0..<Suc k]"
  using k
proof (induct k)
  case 0
  have " {column i A |i. i < 0} = {}" using least_mod_type using leD by auto
  thus ?case
    by (simp, auto simp add: Gram_Schmidt_column_k_efficient_def
      Gram_Schmidt_upt_k_def Gram_Schmidt_column_k_def
      proj_onto_def proj_def vec_eq_iff from_nat_0 to_nat_0)
next
  case (Suc k)
  have "Gram_Schmidt_upt_k A (Suc k) = Gram_Schmidt_column_k (Gram_Schmidt_upt_k A k) (Suc k)"
    by (rule Gram_Schmidt_upt_k_suc)
  also have "... = Gram_Schmidt_column_k_efficient (Gram_Schmidt_upt_k A k) (Suc k)"
  proof (rule Gram_Schmidt_upt_k_efficient_induction)
    show "Gram_Schmidt_upt_k A k = foldl Gram_Schmidt_column_k_efficient A [0..<Suc k]"
      using Suc.hyps Suc.prems by auto
    show "Suc k < ncols A" using Suc.prems by auto
  qed
  also have "... = Gram_Schmidt_column_k_efficient
    (foldl Gram_Schmidt_column_k_efficient A [0..<Suc k]) (Suc k)"
    using Suc.hyps Suc.prems by auto
  also have "... = (foldl Gram_Schmidt_column_k_efficient A [0..<Suc (Suc k)])" by auto
  finally show ?case .
qed

text\<open>This equation is now more efficient than the original definition of the algoritm, since it is not
  removing duplicates in each iteration, which is more expensive in time than adding zeros (if there appear
  duplicates while applying the algorithm, they are zeros and then the @{const sum_list} is the same in each step).\<close>

lemma Gram_Schmidt_matrix_efficient[code_unfold]:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  shows "Gram_Schmidt_matrix A = foldl Gram_Schmidt_column_k_efficient A [0..<ncols A]"
proof -
  have n: "(ncols A - 1) < ncols A" unfolding ncols_def by auto
  have "Gram_Schmidt_matrix A = Gram_Schmidt_upt_k A (ncols A - 1)"
    unfolding Gram_Schmidt_matrix_def ..
  also have "... = foldl Gram_Schmidt_column_k_efficient A [0..<ncols A]"
    using Gram_Schmidt_upt_k_efficient[OF n] unfolding ncols_def by auto
  finally show ?thesis .
qed

subsection\<open>Improvements for computing the Gram Schmidt algorithm and QR decomposition
  using immutable arrays\<close>

subsubsection\<open>New definitions\<close>

definition "Gram_Schmidt_column_k_iarrays_efficient A k =
  tabulate2 (nrows_iarray A) (ncols_iarray A) (\<lambda>a b. let column_b_A = column_iarray b A in
  (if b = k then (column_b_A - sum_list (map (\<lambda>x. ((column_b_A \<bullet>i x) / (x \<bullet>i x)) *\<^sub>R x)
  ((List.map (\<lambda>n. column_iarray n A) [0..<b]))))
  else column_b_A) !! a)"

definition "Gram_Schmidt_matrix_iarrays_efficient A
  = foldl Gram_Schmidt_column_k_iarrays_efficient A [0..<ncols_iarray A]"

definition "QR_decomposition_iarrays_efficient A =
  (let Q = divide_by_norm_iarray (Gram_Schmidt_matrix_iarrays_efficient A)
  in (Q, transpose_iarray Q **i A))"

subsubsection\<open>General properties\<close>

lemma tabulate2_nth:
  assumes i: "i<nr" and j: "j<nc"
  shows "(tabulate2 nr nc f) !! i !! j = f i j"
  unfolding tabulate2_def using i j nth_map by auto

lemma vec_to_iarray_minus[code_unfold]:
  "vec_to_iarray (a - b) =  (vec_to_iarray a) - (vec_to_iarray b)"
  unfolding vec_to_iarray_def
  unfolding minus_iarray_def Let_def by auto

lemma vec_to_iarray_minus_nth:
  assumes A: "i<IArray.length (vec_to_iarray A)"
  and B: "i<IArray.length (vec_to_iarray B)"
  shows "(vec_to_iarray A - vec_to_iarray B) !! i
  = vec_to_iarray A !! i - vec_to_iarray B !! i"
proof -
  have i: "i<CARD('b)" using A unfolding vec_to_iarray_def by auto
  have i2: "i<CARD('c)" using B unfolding vec_to_iarray_def by auto
  have i_length: "i < length [0..<max CARD('b) CARD('c)]" using i i2 by auto
  have i_nth: "[0..<max CARD('b) CARD('c)] ! i = i" using i_length by auto
  let ?f="(\<lambda>a. map (\<lambda>a. if a < CARD('b) then IArray
    (map (\<lambda>i. A $ from_nat i) [0..<CARD('b)]) !! a else 0) [0..<max CARD('b) CARD('c)] !
    a - map (\<lambda>a. if a < CARD('c) then
    IArray (map (\<lambda>i. B $ from_nat i) [0..<CARD('c)]) !! a else 0) [0..<max CARD('b) CARD('c)] ! a)"
  have "(vec_to_iarray A - vec_to_iarray B) = (IArray (map (\<lambda>i. A $ from_nat i) [0..<CARD('b)])
    - IArray (map (\<lambda>i. B $ from_nat i) [0..<CARD('c)]))"
    unfolding vec_to_iarray_def by auto
  also have "... = IArray (map ?f [0..<max CARD('b) CARD('c)])"
    unfolding minus_iarray_def Let_def by simp
  also have "... !! i =  A $ from_nat i - B $ from_nat i"
    using i_length using nth_map i i2 by auto
  also have "... = vec_to_iarray A !! i - vec_to_iarray B !! i"
   by (metis i i2 vec_to_iarray_nth)
  finally show ?thesis .
qed

lemma sum_list_map_vec_to_iarray:
  assumes "xs \<noteq> []" (*If I remove this assumption, I have to prove
  vec_to_iarray 0 = IArray [] which is false.*)
  shows "sum_list (map (vec_to_iarray \<circ> f) xs) = vec_to_iarray (sum_list (map f xs))"
  using assms
proof (induct xs)
  case Nil
  thus ?case unfolding o_def by auto
next
  case (Cons a xs)
  show ?case
  proof (cases "xs=[]")
    case True
    have l_rw: "sum_list (map (vec_to_iarray \<circ> f) xs) = IArray[]"
      unfolding True by (simp add: zero_iarray_def)
    have "sum_list (map (vec_to_iarray \<circ> f) (a # xs))
      = sum_list ((vec_to_iarray \<circ> f) a # map (vec_to_iarray \<circ> f) xs)"
      by simp
    also have "... = (vec_to_iarray \<circ> f) a + sum_list (map (vec_to_iarray \<circ> f) xs)" by simp
    also have "... = vec_to_iarray (f a) + IArray[]" unfolding l_rw by auto
    also have "... = vec_to_iarray (f a) + vec_to_iarray (0::('b,'c) vec)"
      unfolding plus_iarray_def Let_def vec_to_iarray_def by auto
    also have "... = vec_to_iarray (sum_list (map (f) (a # xs)))"
      unfolding True unfolding plus_iarray_def Let_def vec_to_iarray_def by auto
    finally show ?thesis .
  next
    case False
    have "sum_list (map (vec_to_iarray \<circ> f) (a # xs))
      = sum_list ((vec_to_iarray \<circ> f) a # map (vec_to_iarray \<circ> f) xs)"
      by simp
    also have "... = (vec_to_iarray \<circ> f) a + sum_list (map (vec_to_iarray \<circ> f) xs)" by simp
    also have "... = (vec_to_iarray \<circ> f) a + vec_to_iarray (sum_list (map f xs))"
      using Cons.prems Cons.hyps False by presburger
    also have "... = vec_to_iarray (f a) + vec_to_iarray (sum_list (map f xs))" by auto
    also have "... = vec_to_iarray (f a + (sum_list (map f xs)))"
      by (simp add: vec_to_iarray_plus)
    also have "... = vec_to_iarray (sum_list (map (f) (a # xs)))" by simp
    finally show ?thesis .
  qed
qed

subsubsection\<open>Proving the equivalence\<close>

lemma matrix_to_iarray_Gram_Schmidt_column_k_efficient:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  assumes k: "k<ncols A"
  shows "matrix_to_iarray (Gram_Schmidt_column_k_efficient A k)
  = Gram_Schmidt_column_k_iarrays_efficient (matrix_to_iarray A) k"
proof (unfold iarray_exhaust2 list_eq_iff_nth_eq, rule conjI, auto,
    unfold IArray.sub_def[symmetric] IArray.length_def[symmetric])
  show "IArray.length (matrix_to_iarray (Gram_Schmidt_column_k_efficient A k))
    = IArray.length (Gram_Schmidt_column_k_iarrays_efficient (matrix_to_iarray A) k)"
    unfolding matrix_to_iarray_def Gram_Schmidt_column_k_iarrays_efficient_def tabulate2_def
    unfolding nrows_iarray_def by auto
  fix i
  show "i < IArray.length (matrix_to_iarray (Gram_Schmidt_column_k_efficient A k)) \<Longrightarrow>
    IArray.length (matrix_to_iarray (Gram_Schmidt_column_k_efficient A k) !! i) =
    IArray.length (Gram_Schmidt_column_k_iarrays_efficient (matrix_to_iarray A) k !! i)"
    by (simp add: matrix_to_iarray_def Gram_Schmidt_column_k_iarrays_efficient_def
      Gram_Schmidt_column_k_efficient_def tabulate2_def ncols_iarray_def
      nrows_iarray_def vec_to_iarray_def)
  fix i ia
  assume i:"i < IArray.length (matrix_to_iarray (Gram_Schmidt_column_k_efficient A k))"
    and ia: "ia < IArray.length (matrix_to_iarray (Gram_Schmidt_column_k_efficient A k) !! i)"
  show "matrix_to_iarray (Gram_Schmidt_column_k_efficient A k) !! i !! ia
    = Gram_Schmidt_column_k_iarrays_efficient (matrix_to_iarray A) k !! i !! ia"
  proof -
    let ?f="(\<lambda>a b. let column_b_A = column_iarray b (matrix_to_iarray A)
      in (if b = k then column_b_A
      - (\<Sum>x\<leftarrow>map (\<lambda>n. column_iarray n (matrix_to_iarray A)) [0..<b].
      (column_b_A \<bullet>i x / (x \<bullet>i x)) *\<^sub>R x) else column_b_A) !! a)"
    have i2: "i<nrows A" using i unfolding nrows_def matrix_to_iarray_def by auto
    have ia2: "ia<ncols A"
    using ia unfolding ncols_def matrix_to_iarray_def o_def vec_to_iarray_def
      by (metis i2 ia length_vec_to_iarray nrows_def to_nat_from_nat_id vec_matrix)
    have "Gram_Schmidt_column_k_iarrays_efficient (matrix_to_iarray A) k !! i !! ia = ?f i ia"
      unfolding Gram_Schmidt_column_k_iarrays_efficient_def
    proof (rule tabulate2_nth)
      show "i < nrows_iarray (matrix_to_iarray A)"
        using i2 unfolding matrix_to_iarray_nrows .
      show "ia < ncols_iarray (matrix_to_iarray A)"
        using ia2 unfolding matrix_to_iarray_ncols .
    qed
    also have "... = (Gram_Schmidt_column_k_efficient A k) $ (from_nat i) $ (from_nat ia)"
      unfolding Gram_Schmidt_column_k_efficient_def Let_def
    proof (auto)
      assume ia_neq_k: "ia \<noteq> k" and f_eq: "(from_nat ia::'n) = from_nat k "
      have "ia = k" using f_eq by (metis assms from_nat_eq_imp_eq ia2 ncols_def)
      thus "IArray.list_of (column_iarray ia (matrix_to_iarray A)) ! i =
        column (from_nat k) A $ from_nat i -
        sum_list (map ((\<lambda>x. (column (from_nat k) A \<bullet> x / (x \<bullet> x)) *\<^sub>R x)
        \<circ> (\<lambda>n. column (from_nat n) A)) [0..<to_nat (from_nat k)]) $ from_nat i"
        using ia_neq_k by contradiction
    next
      assume "ia \<noteq> k"
      thus "IArray.list_of (column_iarray ia (matrix_to_iarray A)) ! i
        = column (from_nat ia) A $ from_nat i"
        by (metis IArray.sub_def i ia2 length_eq_card_rows to_nat_from_nat_id
          vec_to_iarray_column' vec_to_iarray_nth')
    next
      assume "ia = k"
      let ?f="\<lambda>x. ((column (from_nat k) A) \<bullet> (column (from_nat x) A) /
        ((column (from_nat x) A) \<bullet> (column (from_nat x) A))) *\<^sub>R (column (from_nat x) A)"
      let ?l="sum_list (map ?f [0..<k])"
      show "IArray.list_of
        (column_iarray k (matrix_to_iarray A) -
        sum_list (map ((\<lambda>x. (column_iarray k (matrix_to_iarray A) \<bullet>i x / (x \<bullet>i x)) *\<^sub>R x)
        \<circ> (\<lambda>n. column_iarray n (matrix_to_iarray A))) [0..<k])) ! i =
        column (from_nat k) A $ from_nat i -
        sum_list (map ((\<lambda>x. (column (from_nat k) A \<bullet> x / (x \<bullet> x)) *\<^sub>R x)
        \<circ> (\<lambda>n. column (from_nat n) A)) [0..<to_nat (from_nat k::'n)]) $ from_nat i"
      proof (cases "k=0")
        case True
        show ?thesis
          unfolding vec_to_iarray_column'[OF k, symmetric]
          unfolding True from_nat_0 to_nat_0
          by (auto, metis IArray.sub_def i2 minus_zero_iarray nrows_def vec_to_iarray_nth)
      next
        case False
        have tn_fn_k: "to_nat (from_nat k::'n) = k"
          by (metis assms from_nat_to_nat ncols_def)
        have column_rw: "column_iarray k (matrix_to_iarray A)
          = vec_to_iarray (column (from_nat k) A)"
          by (rule vec_to_iarray_column'[symmetric, OF k])
        have sum_list_rw: "(\<Sum>x\<leftarrow>[0..<k]. (column_iarray k (matrix_to_iarray A)
        \<bullet>i column_iarray x (matrix_to_iarray A) / (column_iarray x (matrix_to_iarray A)
        \<bullet>i column_iarray x (matrix_to_iarray A))) *\<^sub>R column_iarray x (matrix_to_iarray A))
        = vec_to_iarray ?l"
        proof -
          have "(\<Sum>x\<leftarrow>[0..<k]. (column_iarray k (matrix_to_iarray A)
          \<bullet>i column_iarray x (matrix_to_iarray A) / (column_iarray x (matrix_to_iarray A)
          \<bullet>i column_iarray x (matrix_to_iarray A))) *\<^sub>R column_iarray x (matrix_to_iarray A))
          = sum_list (map (vec_to_iarray \<circ> ?f) [0..<k])"
          proof (unfold interv_sum_list_conv_sum_set_nat, rule sum.cong, auto)
            fix x assume "x<k"
            hence x: "x<ncols A" using k by auto
            show "(column_iarray k (matrix_to_iarray A) \<bullet>i column_iarray x (matrix_to_iarray A) /
              (column_iarray x (matrix_to_iarray A) \<bullet>i column_iarray x (matrix_to_iarray A))) *\<^sub>R
              column_iarray x (matrix_to_iarray A) =
              vec_to_iarray ((column (from_nat k) A \<bullet> column (from_nat x) A /
              (column (from_nat x) A \<bullet> column (from_nat x) A)) *\<^sub>R column (from_nat x) A)"
              unfolding vec_to_iarray_scaleR vec_to_iarray_inner
              unfolding column_rw unfolding vec_to_iarray_column'[OF x, symmetric] ..
          qed
          also have "... = vec_to_iarray (sum_list (map ?f [0..<k]))"
            by (rule sum_list_map_vec_to_iarray, auto simp add: False)
          finally show ?thesis .
        qed

        have "IArray.list_of
          (column_iarray k (matrix_to_iarray A) -
          sum_list (map ((\<lambda>x. (column_iarray k (matrix_to_iarray A) \<bullet>i x / (x \<bullet>i x)) *\<^sub>R x)
          \<circ> (\<lambda>n. column_iarray n (matrix_to_iarray A))) [0..<k])) ! i =
          (column_iarray k (matrix_to_iarray A) -
          (\<Sum>x\<leftarrow>[0..<k]. (column_iarray k (matrix_to_iarray A) \<bullet>i column_iarray x (matrix_to_iarray A) /
          (column_iarray x (matrix_to_iarray A) \<bullet>i column_iarray x (matrix_to_iarray A))) *\<^sub>R
          column_iarray x (matrix_to_iarray A))) !! i"
          unfolding vec_to_iarray_inner tn_fn_k o_def
          unfolding IArray.sub_def[symmetric] ..
        also have "... = (vec_to_iarray (column (from_nat k) A) - vec_to_iarray ?l) !! i"
          unfolding sum_list_rw unfolding column_rw  ..
        also have "... = ((vec_to_iarray (column (from_nat k) A)) !! i) - (vec_to_iarray ?l !! i)"
        proof (rule vec_to_iarray_minus_nth)
          show " i < IArray.length (vec_to_iarray (column (from_nat k) A))"
            by (metis i2 length_vec_to_iarray nrows_def)
          show "i < IArray.length (vec_to_iarray ?l)"
            by (metis (no_types, lifting) i2 length_vec_to_iarray nrows_def)
        qed
        also have "... = column (from_nat k) A $ from_nat i - ?l $ from_nat i"
          unfolding column_rw
          by (metis (no_types, lifting) i2 nrows_def vec_to_iarray_nth)
        also have "... = column (from_nat k) A $ from_nat i -
          sum_list (map ((\<lambda>x. (column (from_nat k) A \<bullet> x / (x \<bullet> x)) *\<^sub>R x)
          \<circ> (\<lambda>n. column (from_nat n) A)) [0..<to_nat (from_nat k::'n)]) $ from_nat i"
          unfolding o_def tn_fn_k ..
        finally show "IArray.list_of
          (column_iarray k (matrix_to_iarray A) -
          sum_list (map ((\<lambda>x. (column_iarray k (matrix_to_iarray A) \<bullet>i x / (x \<bullet>i x)) *\<^sub>R x)
          \<circ> (\<lambda>n. column_iarray n (matrix_to_iarray A))) [0..<k])) ! i =
          column (from_nat k) A $ from_nat i -
          sum_list (map ((\<lambda>x. (column (from_nat k) A \<bullet> x / (x \<bullet> x)) *\<^sub>R x)
          \<circ> (\<lambda>n. column (from_nat n) A)) [0..<to_nat (from_nat k::'n)]) $ from_nat i" .
      qed
    qed
    also have "... = matrix_to_iarray (Gram_Schmidt_column_k_efficient A k) !! i !! ia"
      using matrix_to_iarray_nth[of "(Gram_Schmidt_column_k_efficient A k)" "from_nat i" "from_nat ia"]
      using ia2 i2
      unfolding to_nat_from_nat_id[OF i2[unfolded nrows_def]]
      unfolding to_nat_from_nat_id[OF ia2[unfolded ncols_def]] by simp
    finally show ?thesis ..
  qed
qed

lemma matrix_to_iarray_Gram_Schmidt_upt_k_efficient:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  assumes k: "k<ncols A"
  shows "matrix_to_iarray (Gram_Schmidt_upt_k A k)
  = foldl Gram_Schmidt_column_k_iarrays_efficient (matrix_to_iarray A) [0..<Suc k]"
  using assms
proof (induct k)
  case 0
  have zero_le: "0<ncols A" unfolding ncols_def by simp
  thus ?case  unfolding Gram_Schmidt_upt_k_efficient[OF zero_le] unfolding Gram_Schmidt_upt_k_efficient
    by (simp add: matrix_to_iarray_Gram_Schmidt_column_k_efficient[OF "0.prems"])
next
  case (Suc k)
  let ?G="foldl Gram_Schmidt_column_k_iarrays_efficient (matrix_to_iarray A)"
  have k: "k<ncols (Gram_Schmidt_upt_k A k)" using Suc.prems unfolding ncols_def by simp
  have k2: "Suc k < ncols (Gram_Schmidt_upt_k A k)" using Suc.prems unfolding ncols_def .
  have list_rw: "[0..<Suc (Suc k)] = [0..<Suc k] @ [(Suc k)]" by simp
  have hyp: "matrix_to_iarray (Gram_Schmidt_upt_k A k) = ?G [0..<Suc k]"
    by (metis Suc.hyps Suc.prems Suc_lessD)
  show "matrix_to_iarray (Gram_Schmidt_upt_k A (Suc k)) = ?G [0..<Suc (Suc k)]"
    unfolding Gram_Schmidt_upt_k_def
    unfolding list_rw
    unfolding foldl_append
    unfolding foldl.simps
    unfolding Gram_Schmidt_upt_k_def[symmetric]
    unfolding hyp[symmetric]
    using matrix_to_iarray_Gram_Schmidt_column_k_efficient
    by (metis (no_types) Gram_Schmidt_upt_k_efficient Gram_Schmidt_upt_k_efficient_induction
      Suc.prems k matrix_to_iarray_Gram_Schmidt_column_k_efficient ncols_def)
qed

lemma matrix_to_iarray_Gram_Schmidt_matrix_efficient[code_unfold]:
  fixes A::"real^'n::{mod_type}^'m::{mod_type}"
  shows "matrix_to_iarray (Gram_Schmidt_matrix A)
  = Gram_Schmidt_matrix_iarrays_efficient (matrix_to_iarray A)"
proof -
  have n: "ncols A - 1 < ncols A" unfolding ncols_def by auto
  thus ?thesis
    unfolding Gram_Schmidt_matrix_iarrays_efficient_def Gram_Schmidt_matrix_def
    using matrix_to_iarray_Gram_Schmidt_upt_k_efficient[OF n]
    unfolding matrix_to_iarray_ncols by auto
qed

lemma QR_decomposition_iarrays_efficient[code]:
  "QR_decomposition_iarrays (matrix_to_iarray A)
  = QR_decomposition_iarrays_efficient (matrix_to_iarray A)"
  unfolding QR_decomposition_iarrays_def QR_decomposition_iarrays_efficient_def Let_def
  unfolding matrix_to_iarray_Gram_Schmidt_matrix_efficient[symmetric]
  unfolding matrix_to_iarray_Gram_Schmidt_matrix ..

subsection\<open>Other code equations that improve the performance\<close>

lemma inner_iarray_code[code]:
  "inner_iarray A B = sum_list (map (\<lambda>n. A !! n * B !! n) [0..<IArray.length A])"
proof -
  have set_Eq: "{0..<IArray.length A} = set ([0..<IArray.length A])" using atLeastLessThan_upt by blast
  have "inner_iarray A B = sum (\<lambda>n. A !! n * B !! n) {0..<IArray.length A}"
    unfolding inner_iarray_def ..
  also have "... = sum (\<lambda>n. A !! n * B !! n) (set [0..<IArray.length A])"
    unfolding set_Eq ..
  also have "... = sum_list (map (\<lambda>n. A !! n * B !! n) [0..<IArray.length A])"
    unfolding sum_set_upt_conv_sum_list_nat ..
  finally show ?thesis .
qed

definition "Gram_Schmidt_column_k_iarrays_efficient2 A k =
  tabulate2 (nrows_iarray A) (ncols_iarray A)
  (let col_k = column_iarray k A;
       col = (col_k - sum_list (map (\<lambda>x. ((col_k \<bullet>i x) / (x \<bullet>i x)) *\<^sub>R x)
              ((List.map (\<lambda>n. column_iarray n A) [0..<k]))))
  in (\<lambda>a b. (if b = k then col else column_iarray b A) !! a))"

lemma Gram_Schmidt_column_k_iarrays_efficient_eq[code]: "Gram_Schmidt_column_k_iarrays_efficient A k
  = Gram_Schmidt_column_k_iarrays_efficient2 A k"
  unfolding Gram_Schmidt_column_k_iarrays_efficient_def
  unfolding Gram_Schmidt_column_k_iarrays_efficient2_def
  unfolding Let_def tabulate2_def
  by simp

end


section \<open>Hidden Markov Models\<close>

theory Hidden_Markov_Model
  imports
    Markov_Models.Discrete_Time_Markov_Chain Auxiliary
    "HOL-Library.IArray"
begin

subsection \<open>Definitions\<close>

text \<open>Definition of Markov Kernels that are closed w.r.t. to a set of states.\<close>
locale Closed_Kernel =
  fixes K :: "'s \<Rightarrow> 't pmf" and S :: "'t set"
  assumes finite: "finite S"
      and wellformed: "S \<noteq> {}"
      and closed: "\<forall> s. K s \<subseteq> S"

text \<open>
  An HMM is parameterized by a Markov kernel for the transition probabilites between internal states,
  a Markov kernel for the output probabilities of observations,
  and a fixed set of observations.
\<close>
locale HMM_defs =
  fixes \<K> :: "'s \<Rightarrow> 's pmf" and \<O> :: "'s \<Rightarrow> 't pmf" and \<O>\<^sub>s :: "'t set"

locale HMM =
  HMM_defs + O: Closed_Kernel \<O> \<O>\<^sub>s
begin

lemma observations_finite: "finite \<O>\<^sub>s"
  and observations_wellformed: "\<O>\<^sub>s \<noteq> {}"
  and observations_closed: "\<forall> s. \<O> s \<subseteq> \<O>\<^sub>s"
  using O.finite O.wellformed O.closed by -

end

text \<open>Fixed set of internal states.\<close>
locale HMM2_defs = HMM_defs \<K> \<O> for \<K> :: "'s \<Rightarrow> 's pmf" and \<O> :: "'s \<Rightarrow> 't pmf" +
  fixes \<S> :: "'s set"

locale HMM2 = HMM2_defs + HMM + K: Closed_Kernel \<K> \<S>
begin

lemma states_finite: "finite \<S>"
  and states_wellformed: "\<S> \<noteq> {}"
  and states_closed: "\<forall> s. \<K> s \<subseteq> \<S>"
  using K.finite K.wellformed K.closed by -

end

text \<open>
  The set of internal states is now given as a list to iterate over.
  This is needed for the computations on HMMs.
\<close>
locale HMM3_defs = HMM2_defs \<O>\<^sub>s \<K> for \<O>\<^sub>s :: "'t set" and \<K> :: "'s \<Rightarrow> 's pmf" +
  fixes state_list :: "'s list"

locale HMM3 = HMM3_defs _ _ \<O>\<^sub>s \<K> + HMM2 \<O>\<^sub>s \<K> for \<O>\<^sub>s :: "'t set" and \<K> :: "'s \<Rightarrow> 's pmf" +
  assumes state_list_\<S>: "set state_list = \<S>"

context HMM_defs
begin

no_notation (ASCII) comp  (infixl "o" 55)

text \<open>The ``default'' observation.\<close>
definition
  "obs \<equiv> SOME x. x \<in> \<O>\<^sub>s"

lemma (in HMM) obs:
  "obs \<in> \<O>\<^sub>s"
  unfolding obs_def using observations_wellformed by (auto intro: someI_ex)

text \<open>
  The HMM is encoded as a Markov chain over pairs of states and observations.
  This is the Markov chain's defining Markov kernel.
\<close>
definition
  "K \<equiv> \<lambda> (s\<^sub>1, o\<^sub>1 :: 't). bind_pmf (\<K> s\<^sub>1) (\<lambda> s\<^sub>2. map_pmf (\<lambda> o\<^sub>2. (s\<^sub>2, o\<^sub>2)) (\<O> s\<^sub>2))"

sublocale MC_syntax K .

text \<open>
  Uniform distribution of the pairs \<open>(s, o)\<close> for a fixed state \<open>s\<close>.
\<close>
definition "I (s :: 's) = map_pmf (\<lambda> x. (s, x)) (pmf_of_set \<O>\<^sub>s)"

text \<open>
  The likelihood of an observation sequence given a starting state \<open>s\<close> is defined in terms of
  the trace space of the Markov kernel given the uniform distribution of pairs for \<open>s\<close>.
\<close>
definition
  "likelihood s os = T' (I s) {\<omega> \<in> space S. \<exists> o\<^sub>0 xs \<omega>'. \<omega> = (s, o\<^sub>0) ## xs @- \<omega>' \<and> map snd xs = os}"

abbreviation (input) "L os \<omega> \<equiv> \<exists> xs \<omega>'. \<omega> = xs @- \<omega>' \<and> map snd xs = os"

lemma likelihood_alt_def: "likelihood s os = T' (I s) {(s, o) ## xs @- \<omega>' |o xs \<omega>'. map snd xs = os}"
  unfolding likelihood_def by (simp add: in_S)

subsection \<open>Iteration Rule For Likelihood\<close>

lemma L_Nil:
  "L [] \<omega> = True"
  by simp

lemma emeasure_T_observation_Nil:
  "T (s, o\<^sub>0) {\<omega> \<in> space S. L [] \<omega>} = 1"
  by simp

lemma L_Cons:
  "L (o # os) \<omega> \<longleftrightarrow> snd (shd \<omega>) = o \<and> L os (stl \<omega>)"
  apply (cases \<omega>; cases "shd \<omega>"; safe; clarsimp)
   apply force
  subgoal for x xs \<omega>'
    by (force intro: exI[where x = "(x, o) # xs"])
  done

lemma L_measurable[measurable]:
  "Measurable.pred S (L os)"
  apply (induction os)
   apply (simp; fail)
  subgoal premises that for o os
    by(subst L_Cons)
      (intro Measurable.pred_intros_logic
        measurable_compose[OF measurable_shd] measurable_compose[OF measurable_stl that];
        measurable)
  done

lemma init_measurable[measurable]:
  "Measurable.pred S (\<lambda>x. \<exists>o\<^sub>0 xs \<omega>'. x = (s, o\<^sub>0) ## xs @- \<omega>' \<and> map snd xs = os)"
  (is "Measurable.pred S ?f")
proof -
  have *: "?f \<omega> \<longleftrightarrow> fst (shd \<omega>) = s \<and> L os (stl \<omega>)" for \<omega>
    by (cases \<omega>) auto
  show ?thesis
    by (subst *)
       (intro Measurable.pred_intros_logic measurable_compose[OF measurable_shd]; measurable)
qed

lemma T_init_observation_eq:
  "T (s, o) {\<omega> \<in> space S. L os \<omega>} = T (s, o') {\<omega> \<in> space S. L os \<omega>}"
  apply (subst emeasure_Collect_T[unfolded space_T], (measurable; fail))
  apply (subst (2) emeasure_Collect_T[unfolded space_T], (measurable; fail))
  apply (simp add: K_def)
  done

text \<open>
  Shows that it is equivalent to define likelihood in terms of the trace space starting at a single
  pair of an internal state \<open>s\<close> and the default observation @{term obs}.
\<close>
lemma (in HMM) likelihood_init:
  "likelihood s os = T (s, obs) {\<omega> \<in> space S. L os \<omega>}"
proof -
  have *: "(\<Sum>o\<in>\<O>\<^sub>s. emeasure (T (s, o)) {\<omega> \<in> space S. L os \<omega>}) =
    of_nat (card \<O>\<^sub>s) * emeasure (T (s, obs)) {\<omega> \<in> space S. L os \<omega>}"
    by (subst sum_constant[symmetric]) (fastforce intro: sum.cong T_init_observation_eq[simplified])
  show ?thesis
    unfolding likelihood_def
    apply (subst emeasure_T')
    subgoal
      by measurable
    using *
    apply (simp add: I_def in_S observations_finite observations_wellformed nn_integral_pmf_of_set)
    apply (subst mult.commute)
    apply (simp add: observations_finite observations_wellformed mult_divide_eq_ennreal)
    done
qed

lemma emeasure_T_observation_Cons:
  "T (s, o\<^sub>0) {\<omega> \<in> space S. L (o\<^sub>1 # os) \<omega>} =
   (\<integral>\<^sup>+ t. ennreal (pmf (\<O> t) o\<^sub>1) * T (t, o\<^sub>1) {\<omega> \<in> space S. L os \<omega>} \<partial>(\<K> s))" (is "?l = ?r")
proof -
  have *:
    "\<integral>\<^sup>+ y. T (s', y) {x \<in> space S. \<exists>xs. (\<exists>\<omega>'. (s', y) ## x = xs @- \<omega>') \<and> map snd xs = o\<^sub>1 # os}
       \<partial>measure_pmf (\<O> s') =
    ennreal (pmf (\<O> s') o\<^sub>1) * T (s', o\<^sub>1) {\<omega> \<in> space S. \<exists>xs. (\<exists>\<omega>'. \<omega> = xs @- \<omega>') \<and> map snd xs = os}"
    (is "?L = ?R") for s'
  proof -
    have "?L = \<integral>\<^sup>+ x. ennreal (pmf (\<O> s') x) *
            T (s', x) {\<omega> \<in> space S. \<exists>xs. (\<exists>\<omega>'. (s', x) ## \<omega> = xs @- \<omega>') \<and> map snd xs = o\<^sub>1 # os}
          \<partial>count_space UNIV"
      by (rule nn_integral_measure_pmf)
    also have "\<dots> =
      \<integral>\<^sup>+ o\<^sub>2. (if o\<^sub>2 = o\<^sub>1
              then ennreal (pmf (\<O> s') o\<^sub>1) * T (s', o\<^sub>1) {\<omega> \<in> space S. L os \<omega>}
              else 0)
       \<partial>count_space UNIV"
      apply (rule nn_integral_cong_AE
          [where v = "\<lambda> o\<^sub>2. if o\<^sub>2 = o\<^sub>1
            then ennreal (pmf (\<O> s') o\<^sub>1) * T (s', o\<^sub>1) {\<omega> \<in> space S. L os \<omega>} else 0"]
          )
       apply (rule AE_I2)
       apply (split if_split, safe)
      subgoal
        by (auto intro!: arg_cong2[where f = times, OF HOL.refl] arg_cong2[where f = emeasure];
            metis list.simps(9) shift.simps(2) snd_conv
           )
      subgoal
        by (subst arg_cong2[where f = emeasure and d = "{}", OF HOL.refl]) auto
      done
    also have "\<dots> = \<integral>\<^sup>+o\<^sub>2\<in>{o\<^sub>1}.
       (ennreal (pmf (\<O> s') o\<^sub>1) * T (s', o\<^sub>1) {\<omega> \<in> space S. L os \<omega>})
      \<partial>count_space UNIV"
      by (rule nn_integral_cong_AE) auto
    also have "\<dots> = ?R"
      by simp
    finally show ?thesis .
  qed
  have "?l = \<integral>\<^sup>+ t. T t {x \<in> space S. \<exists>xs \<omega>'. t ## x = xs @- \<omega>' \<and> map snd xs = o\<^sub>1 # os} \<partial> (K (s, o\<^sub>0))"
    by (subst emeasure_Collect_T[unfolded space_T], measurable)
  also have "\<dots> = ?r"
    using * by (simp add: K_def)
  finally show ?thesis .
qed

subsection \<open>Computation of Likelihood\<close>

fun backward where
  "backward s [] = 1" |
  "backward s (o # os) = (\<integral>\<^sup>+ t. ennreal (pmf (\<O> t) o) * backward t os \<partial>measure_pmf (\<K> s))"

lemma emeasure_T_observation_backward:
  "emeasure (T (s, o)) {\<omega> \<in> space S. L os \<omega>} = backward s os"
  using emeasure_T_observation_Cons by (induction os arbitrary: s o; simp)

lemma (in HMM) likelihood_backward:
  "likelihood s os = backward s os"
  unfolding likelihood_init emeasure_T_observation_backward ..

end (* HMM Defs *)

context HMM2
begin

fun (in HMM2_defs) forward where
  "forward s t_end [] = indicator {t_end} s" |
  "forward s t_end (o # os) =
    (\<Sum>t \<in> \<S>. ennreal (pmf (\<O> t) o) * ennreal (pmf (\<K> s) t) * forward t t_end os)"

lemma forward_split:
  "forward s t (os1 @ os2) = (\<Sum>t' \<in> \<S>. forward s t' os1 * forward t' t os2)"
  if "s \<in> \<S>"
  using that
  apply (induction os1 arbitrary: s)
  subgoal for s
    apply (simp add: sum_indicator_mult[OF states_finite])
    apply (subst sum.cong[where B = "{s}"])
    by auto
  subgoal for a os1 s
    apply simp
    apply (subst sum_distrib_right)
    apply (subst sum.swap)
    apply (simp add: sum_distrib_left algebra_simps)
    done
  done

lemma (in -)
  "(\<Sum>t \<in> S. f t) = f t" if "finite S" "t \<in> S" "\<forall> s \<in> S - {t}. f s = 0"
  thm sum.empty sum.insert sum.mono_neutral_right[of S "{t}"]
  apply (subst sum.mono_neutral_right[of S "{t}"])
  using that
     apply auto
  done
(*
  oops
  by (metis add.right_neutral empty_iff finite.intros(1) insert_iff subsetI sum.empty sum.insert sum.mono_neutral_right that)

  using that
  apply auto
*)

lemma forward_backward:
  "(\<Sum>t \<in> \<S>. forward s t os) = backward s os" if "s \<in> \<S>"
  using \<open>s \<in> \<S>\<close>
  apply (induction os arbitrary: s)
  subgoal for s
    by (subst sum.mono_neutral_right[of \<S> "{s}", OF states_finite])
       (auto split: if_split_asm simp: indicator_def)
  subgoal for a os s
    apply (simp add: sum.swap sum_distrib_left[symmetric])
    apply (subst nn_integral_measure_pmf_support[where A = \<S>])
    using states_finite states_closed by (auto simp: algebra_simps)
  done

theorem likelihood_forward:
  "likelihood s os = (\<Sum>t \<in> \<S>. forward s t os)" if \<open>s \<in> \<S>\<close>
  unfolding likelihood_backward forward_backward[symmetric, OF \<open>s \<in> \<S>\<close>] ..

subsection \<open>Definition of Maximum Probabilities\<close>

abbreviation (input) "V os as \<omega> \<equiv> (\<exists> \<omega>'. \<omega> = zip as os @- \<omega>')"

definition
  "max_prob s os =
  Max {T' (I s) {\<omega> \<in> space S. \<exists>o \<omega>'. \<omega> = (s, o) ## zip as os @- \<omega>'}
       | as. length as = length os \<and> set as \<subseteq> \<S>}"

fun viterbi_prob where
  "viterbi_prob s t_end [] = indicator {t_end} s" |
  "viterbi_prob s t_end (o # os) =
    (MAX t \<in> \<S>. ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * viterbi_prob t t_end os)"

definition
  "is_decoding s os as \<equiv>
    T' (I s) {\<omega> \<in> space S. \<exists>o \<omega>'. \<omega> = (s, o) ## zip as os @- \<omega>'} = max_prob s os \<and>
    length as = length os \<and> set as \<subseteq> \<S>"

subsection \<open>Iteration Rule For Maximum Probabilities\<close>

lemma emeasure_T_state_Nil:
  "T (s, o\<^sub>0) {\<omega> \<in> space S. V [] as \<omega>} = 1"
  by simp

lemma max_prob_T_state_Nil:
  "Max {T (s, o) {\<omega> \<in> space S. V [] as \<omega>} | as. length as = length [] \<and> set as \<subseteq> \<S>} = 1"
  by (simp add: emeasure_T_state_Nil)

lemma V_Cons: "V (o # os) (a # as) \<omega> \<longleftrightarrow> fst (shd \<omega>) = a \<and> snd (shd \<omega>) = o \<and> V os as (stl \<omega>)"
  by (cases \<omega>) auto

lemma measurable_V[measurable]:
  "Measurable.pred S (\<lambda>\<omega>. V os as \<omega>)"
proof (induction os as rule: list_induct2')
  case (4 x xs y ys)
  then show ?case
    by (subst V_Cons)
       (intro Measurable.pred_intros_logic
          measurable_compose[OF measurable_shd] measurable_compose[OF measurable_stl];
        measurable)
qed simp+

lemma init_V_measurable[measurable]:
  "Measurable.pred S (\<lambda>x. \<exists>o \<omega>'. x = (s, o) ## zip as os @- \<omega>')" (is "Measurable.pred S ?f")
proof -
  have *: "?f \<omega> \<longleftrightarrow> fst (shd \<omega>) = s \<and> V os as (stl \<omega>)" for \<omega>
    by (cases \<omega>) auto
  show ?thesis
    by (subst *)
       (intro Measurable.pred_intros_logic measurable_compose[OF measurable_shd]; measurable)
qed

lemma max_prob_Cons':
  "Max {T (s, o\<^sub>1) {\<omega> \<in> space S. V (o # os) as \<omega>} | as. length as = length (o # os) \<and> set as \<subseteq> \<S>} =
  (
    MAX t \<in> \<S>. ennreal (pmf (\<O> t) o * pmf (\<K> s) t) *
      (MAX as \<in> {as. length as = length os \<and> set as \<subseteq> \<S>}. T (t, o) {\<omega> \<in> space S. V os as \<omega>})
  )" (is "?l = ?r")
  and T_V_Cons:
  "T (s, o\<^sub>1) {\<omega> \<in> space S. V (o # os) (t # as) \<omega>}
  = ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * T (t, o) {\<omega> \<in> space S. V os as \<omega>}"
  (is "?l' = ?r'")
  if "length as = length os"
proof -
  let ?S = "\<lambda> os. {as. length as = length os \<and> set as \<subseteq> \<S>}"
  have S_finite: "finite (?S os)" for os :: "'t list"
    using finite_lists_length_eq[OF states_finite] by (rule finite_subset[rotated]) auto
  have S_nonempty: "?S os \<noteq> {}" for os :: "'t list"
  proof -
    let ?a = "SOME a. a \<in> \<S>" let ?as = "replicate (length os) ?a"
    from states_wellformed have "?a \<in> \<S>"
      by (auto intro: someI_ex)
    then have "?as \<in> ?S os"
      by auto
    then show ?thesis
      by force
  qed
  let ?f = "\<lambda>t as os. T t {\<omega> \<in> space S. V os as (t ## \<omega>)}"
  let ?g = "\<lambda>t as os. T t {\<omega> \<in> space S. V os as \<omega>}"
  have *: "?f t as (o # os) = ?g t (tl as) os * indicator {(hd as, o)} t"
    if "length as = Suc n" for t as n
    unfolding indicator_def using that by (cases as) auto
  have **: "K (s, o\<^sub>1) {(t, o)} = pmf (\<O> t) o * pmf (\<K> s) t" for t
    unfolding K_def
    apply (simp add: vimage_def)
    apply (subst arg_cong2[where
          f = nn_integral and d = "\<lambda> x. \<O> x {xa. xa = o \<and> x = t} * indicator {t} x",
          OF HOL.refl])
    subgoal
      by (auto simp: indicator_def)
    by (simp add: emeasure_pmf_single ennreal_mult')
  have "?l = (MAX as \<in> ?S (o # os). \<integral>\<^sup>+ t. ?f t as (o # os) \<partial>K (s, o\<^sub>1))"
    by (subst Max_to_image2; subst emeasure_Collect_T[unfolded space_T]; rule measurable_V HOL.refl)
  also have "\<dots> = (MAX as \<in> ?S (o # os). \<integral>\<^sup>+ t. ?g t (tl as) os * indicator {(hd as,o)} t \<partial>K (s,o\<^sub>1))"
    by (simp cong: Max_image_cong_simp add: *)
  also have "\<dots> = (MAX(t, as)\<in> \<S> \<times> ?S os. ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * ?g (t, o) as os)"
  proof ((rule Max_eq_image_if; clarsimp?), goal_cases)
    case 1
    from S_finite[of "o # os"] show ?case
      by simp
  next
    case 2
    from states_finite show ?case
      by (blast intro: S_finite)
  next
    case (3 as)
    then show ?case
      by - (rule bexI[where x = "hd as"]; cases as; auto simp: algebra_simps **)
  next
    case (4 x as)
    then show ?case
      by - (rule exI[where x = "x # as"], simp add: algebra_simps **)
  qed
  also have "\<dots> = ?r"
    by (subst Max_image_left_mult[symmetric], fact+)
       (rule sym, rule Max_image_pair, rule states_finite, fact+)
  finally show "?l = ?r" .
  have "?l' = \<integral>\<^sup>+ t'. ?f t' (t # as) (o # os) \<partial>K (s, o\<^sub>1)"
    by (rule emeasure_Collect_T[unfolded space_T]; rule measurable_V)
  also from that have "\<dots> = \<integral>\<^sup>+ t'. ?g t' as os * indicator {(t,o)} t' \<partial>K (s,o\<^sub>1)"
    by (subst *[of _ "length as"]; simp)
  also have "\<dots> = ?r'"
    by (simp add: **, simp only: algebra_simps)
  finally show "?l' = ?r'" .
qed

lemmas max_prob_Cons = max_prob_Cons'[OF length_replicate]

subsection \<open>Computation of Maximum Probabilities\<close>

lemma T_init_V_eq:
  "T (s, o) {\<omega> \<in> space S. V os as \<omega>} = T (s, o') {\<omega> \<in> space S. V os as \<omega>}"
  apply (subst emeasure_Collect_T[unfolded space_T], (measurable; fail))
  apply (subst (2) emeasure_Collect_T[unfolded space_T], (measurable; fail))
  apply (simp add: K_def)
  done

lemma T'_I_T:
  "T' (I s) {\<omega> \<in> space S. \<exists>o \<omega>'. \<omega> = (s, o) ## zip as os @- \<omega>'} = T (s,o) {\<omega> \<in> space S. V os as \<omega>}"
proof -
  have "(\<Sum>o\<in>\<O>\<^sub>s. T (s, o) {\<omega> \<in> space S. V os as \<omega>}) =
    of_nat (card \<O>\<^sub>s) * T (s, o) {\<omega> \<in> space S. V os as \<omega>}" for as
    by (subst sum_constant[symmetric]) (fastforce intro: sum.cong T_init_V_eq[simplified])
  then show ?thesis
    unfolding max_prob_def
    apply (subst emeasure_T')
    subgoal
      by measurable
    apply (simp add: I_def in_S observations_finite observations_wellformed nn_integral_pmf_of_set)
    apply (subst mult.commute)
    apply (simp add: observations_finite observations_wellformed mult_divide_eq_ennreal)
    done
qed

lemma max_prob_init:
  "max_prob s os = Max {T (s,o) {\<omega> \<in> space S. V os as \<omega>} | as. length as = length os \<and> set as \<subseteq> \<S>}"
  unfolding max_prob_def by (simp add: T'_I_T[symmetric])

lemma max_prob_Nil[simp]:
  "max_prob s [] = 1"
  unfolding max_prob_init[where o = obs] by auto

lemma Max_start:
  "(MAX t\<in>\<S>. (indicator {t} s :: ennreal)) = 1" if "s \<in> \<S>"
  using states_finite that by (auto simp: indicator_def intro: Max_eqI)

lemma Max_V_viterbi:
  "(MAX t \<in> \<S>. viterbi_prob s t os) =
   Max {T (s, o) {\<omega> \<in> space S. V os as \<omega>} | as. length as = length os \<and> set as \<subseteq> \<S>}" if "s \<in> \<S>"
  using that states_finite states_wellformed
  by (induction os arbitrary: s o; simp
        add: Max_start max_prob_Cons[simplified] Max_image_commute Max_image_left_mult Max_to_image2
        cong: Max_image_cong
      )

lemma max_prob_viterbi:
  "(MAX t \<in> \<S>. viterbi_prob s t os) = max_prob s os" if "s \<in> \<S>"
  using max_prob_init[of s os] Max_V_viterbi[OF \<open>s \<in> \<S>\<close>, symmetric] by simp

end

subsection \<open>Decoding the Most Probable Hidden State Sequence\<close>

context HMM3
begin

fun viterbi where
  "viterbi s t_end [] = ([], indicator {t_end} s)" |
  "viterbi s t_end (o # os) = fst (
    argmax snd (map
      (\<lambda>t. let (xs, v) = viterbi t t_end os in (t # xs, ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * v))
    state_list))"

lemma state_list_nonempty:
  "state_list \<noteq> []"
  using state_list_\<S> states_wellformed by auto

lemma viterbi_viterbi_prob:
  "snd (viterbi s t_end os) = viterbi_prob s t_end os"
proof (induction os arbitrary: s)
  case Nil
  then show ?case
    by simp
next
  case (Cons o os)
  let ?f =
    "\<lambda>t. let (xs, v) = viterbi t t_end os in (t # xs, ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * v)"
  let ?xs = "map ?f state_list"
  from state_list_nonempty have "map ?f state_list \<noteq> []"
    by simp
  from argmax(2,3)[OF this, of snd] have *:
    "snd (fst (argmax snd ?xs)) = snd (argmax snd ?xs)"
    "snd (argmax snd ?xs) = (MAX x \<in> set ?xs. snd x)" .
  then show ?case
    apply (simp add: state_list_\<S>)
    apply (rule Max_eq_image_if)
       apply (intro finite_imageI states_finite; fail)
      apply (intro finite_imageI states_finite; fail)
    subgoal
      apply clarsimp
      subgoal for x
        using Cons.IH[of x] by (auto split: prod.splits)
      done
    apply clarsimp
    subgoal for x
      using Cons.IH[of x] by (force split: prod.splits)
    done
qed

context
begin

private fun val_of where
  "val_of s [] [] = 1" |
  "val_of s (t # xs) (o # os) = ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * val_of t xs os"

lemma val_of_T:
  "val_of s as os = T (s, o\<^sub>1) {\<omega> \<in> space S. V os as \<omega>}" if "length as = length os"
  using that by (induction arbitrary: o\<^sub>1 rule: val_of.induct; (subst T_V_Cons)?; simp)

lemma viterbi_sequence:
  "snd (viterbi s t_end os) = val_of s (fst (viterbi s t_end os)) os"
  if "snd (viterbi s t_end os) > 0"
  using that
proof (induction os arbitrary: s)
  case Nil
  then show ?case
    by (simp add: indicator_def split: if_split_asm)
next
  case (Cons o os s)
  let ?xs = "map
    (\<lambda>t. let (xs, v) = viterbi t t_end os in (t # xs, ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * v))
    state_list"
  from state_list_nonempty have "?xs \<noteq> []"
    by simp
  from argmax(1)[OF this, of snd] obtain t where
    "t \<in> set state_list"
    "fst (argmax snd ?xs) =
    (t # fst (viterbi t t_end os), ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * snd (viterbi t t_end os))"
    by (auto split: prod.splits)
  with Cons show ?case
    by (auto simp: ennreal_zero_less_mult_iff)
qed

lemma viterbi_valid_path:
  "length as = length os \<and> set as \<subseteq> \<S>" if "viterbi s t_end os = (as, v)"
using that proof (induction os arbitrary: s as v)
  case Nil
  then show ?case
    by simp
next
  case (Cons o os s as v)
  let ?xs = "map
    (\<lambda>t. let (xs, v) = viterbi t t_end os in (t # xs, ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * v))
    state_list"
  from state_list_nonempty have "?xs \<noteq> []"
    by simp
  from argmax(1)[OF this, of snd] obtain t where "t \<in> \<S>"
    "fst (argmax snd ?xs) =
    (t # fst (viterbi t t_end os), ennreal (pmf (\<O> t) o * pmf (\<K> s) t) * snd (viterbi t t_end os))"
    by (auto simp: state_list_\<S> split: prod.splits)
  with Cons.prems show ?case
    by (cases "viterbi t t_end os"; simp add: Cons.IH)
qed

definition
  "viterbi_final s os = fst (argmax snd (map (\<lambda> t. viterbi s t os) state_list))"

lemma viterbi_finalE:
  obtains t where
    "t \<in> \<S>" "viterbi_final s os = viterbi s t os"
    "snd (viterbi s t os) = Max ((\<lambda>t. snd (viterbi s t os)) ` \<S>)"
proof -
  from state_list_nonempty have "map (\<lambda> t. viterbi s t os) state_list \<noteq> []"
    by simp
  from argmax[OF this, of snd] show ?thesis
    by (auto simp: state_list_\<S> image_comp comp_def viterbi_final_def intro: that)
qed

theorem viterbi_final_max_prob:
  assumes "viterbi_final s os = (as, v)" "s \<in> \<S>"
  shows "v = max_prob s os"
proof -
  obtain t where "t \<in> \<S>" "viterbi_final s os = viterbi s t os"
    "snd (viterbi s t os) = Max ((\<lambda>t. snd (viterbi s t os)) ` \<S>)"
    by (rule viterbi_finalE)
  with assms show ?thesis
    by (simp add: viterbi_viterbi_prob max_prob_viterbi)
qed

theorem viterbi_final_is_decoding:
  assumes "viterbi_final s os = (as, v)" "v > 0" "s \<in> \<S>"
  shows "is_decoding s os as"
proof -
  from viterbi_valid_path[of s _ os as v] assms have as: "length as = length os" "set as \<subseteq> \<S>"
    by - (rule viterbi_finalE[of s os]; simp)+
  obtain t where "t \<in> \<S>" "viterbi_final s os = viterbi s t os"
    by (rule viterbi_finalE)
  with assms viterbi_sequence[of s t os] have "val_of s as os = v"
    by (cases "viterbi s t os") (auto simp: snd_def split!: prod.splits)
  with val_of_T as have "max_prob s os = T (s, obs) {\<omega> \<in> space S. V os as \<omega>}"
    by (simp add: viterbi_final_max_prob[OF assms(1,3)])
  with as show ?thesis
    unfolding is_decoding_def by (simp only: T'_I_T)
qed

end (* Anonymous context *)

end (* HMM 3 *)

end (* Theory *)


(*  Title:      HOL/Proofs/Lambda/ListOrder.thy
    Author:     Tobias Nipkow
    Copyright   1998 TU Muenchen
*)

section \<open>Lifting an order to lists of elements\<close>

theory ListOrder
imports Main
begin

declare [[syntax_ambiguity_warning = false]]

text \<open>
  Lifting an order to lists of elements, relating exactly one
  element.
\<close>

definition
  step1 :: "('a => 'a => bool) => 'a list => 'a list => bool" where
  "step1 r =
    (\<lambda>ys xs. \<exists>us z z' vs. xs = us @ z # vs \<and> r z' z \<and> ys =
      us @ z' # vs)"

lemma step1_converse [simp]: "step1 (r\<inverse>\<inverse>) = (step1 r)\<inverse>\<inverse>"
  apply (unfold step1_def)
  apply (blast intro!: order_antisym)
  done

lemma in_step1_converse [iff]: "(step1 (r\<inverse>\<inverse>) x y) = ((step1 r)\<inverse>\<inverse> x y)"
  apply auto
  done

lemma not_Nil_step1 [iff]: "\<not> step1 r [] xs"
  apply (unfold step1_def)
  apply blast
  done

lemma not_step1_Nil [iff]: "\<not> step1 r xs []"
  apply (unfold step1_def)
  apply blast
  done

lemma Cons_step1_Cons [iff]:
    "(step1 r (y # ys) (x # xs)) =
      (r y x \<and> xs = ys \<or> x = y \<and> step1 r ys xs)"
  apply (unfold step1_def)
  apply (rule iffI)
   apply (erule exE)
   apply (rename_tac ts)
   apply (case_tac ts)
    apply fastforce
   apply force
  apply (erule disjE)
   apply blast
  apply (blast intro: Cons_eq_appendI)
  done

lemma append_step1I:
  "step1 r ys xs \<and> vs = us \<or> ys = xs \<and> step1 r vs us
    ==> step1 r (ys @ vs) (xs @ us)"
  apply (unfold step1_def)
  apply auto
   apply blast
  apply (blast intro: append_eq_appendI)
  done

lemma Cons_step1E [elim!]:
  assumes "step1 r ys (x # xs)"
    and "!!y. ys = y # xs \<Longrightarrow> r y x \<Longrightarrow> R"
    and "!!zs. ys = x # zs \<Longrightarrow> step1 r zs xs \<Longrightarrow> R"
  shows R
  using assms
  apply (cases ys)
   apply (simp add: step1_def)
  apply blast
  done

lemma Snoc_step1_SnocD:
  "step1 r (ys @ [y]) (xs @ [x])
    ==> (step1 r ys xs \<and> y = x \<or> ys = xs \<and> r y x)"
  apply (unfold step1_def)
  apply (clarify del: disjCI)
  apply (rename_tac vs)
  apply (rule_tac xs = vs in rev_exhaust)
   apply force
  apply simp
  apply blast
  done

lemma Cons_acc_step1I [intro!]:
    "Wellfounded.accp r x ==> Wellfounded.accp (step1 r) xs \<Longrightarrow> Wellfounded.accp (step1 r) (x # xs)"
  apply (induct arbitrary: xs set: Wellfounded.accp)
  apply (erule thin_rl)
  apply (erule accp_induct)
  apply (rule accp.accI)
  apply blast
  done

lemma lists_accD: "listsp (Wellfounded.accp r) xs ==> Wellfounded.accp (step1 r) xs"
  apply (induct set: listsp)
   apply (rule accp.accI)
   apply simp
  apply (rule accp.accI)
  apply (fast dest: accp_downward)
  done

lemma ex_step1I:
  "[| x \<in> set xs; r y x |]
    ==> \<exists>ys. step1 r ys xs \<and> y \<in> set ys"
  apply (unfold step1_def)
  apply (drule in_set_conv_decomp [THEN iffD1])
  apply force
  done

lemma lists_accI: "Wellfounded.accp (step1 r) xs ==> listsp (Wellfounded.accp r) xs"
  apply (induct set: Wellfounded.accp)
  apply clarify
  apply (rule accp.accI)
  apply (drule_tac r=r in ex_step1I, assumption)
  apply blast
  done

end