```(* Author: Johannes Hölzl <hoelzl@in.tum.de> *)

imports "../Discrete_Time_Markov_Chain"
begin

lemma inj_on_upd_PiE:
assumes "i ∉ I" shows "inj_on (λ(x,f). f(i := x)) (M × (Π⇩E i∈I. A i))"
unfolding PiE_def
proof (safe intro!: inj_onI ext)
fix f g :: "'a ⇒ 'b" and x y :: 'b
assume *: "f(i := x) = g(i := y)" "f ∈ extensional I" "g ∈ extensional I"
then show "x = y" by (auto simp: fun_eq_iff split: if_split_asm)
fix i' from * ‹i ∉ I› show "f i' = g i'"
by (cases "i' = i") (auto simp: fun_eq_iff extensional_def split: if_split_asm)
qed

lemma sum_folded_product:
fixes I :: "'i set" and f :: "'s ⇒ 'i ⇒ 'a::{semiring_0, comm_monoid_mult}"
assumes "finite I" "⋀i. i ∈ I ⟹ finite (S i)"
shows "(∑x∈Pi⇩E I S. ∏i∈I. f (x i) i) = (∏i∈I. ∑s∈S i. f s i)"
using assms proof (induct I)
case empty then show ?case by simp
next
case (insert i I)
have *: "Pi⇩E (insert i I) S = (λ(x, f). f(i := x)) ` (S i × Pi⇩E I S)"
by (auto simp: PiE_def intro!: image_eqI ext dest: extensional_arb)
have "(∑x∈Pi⇩E (insert i I) S. ∏i∈insert i I. f (x i) i) =
sum ((λx. ∏i∈insert i I. f (x i) i) ∘ ((λ(x, f). f(i := x)))) (S i × Pi⇩E I S)"
unfolding * using insert by (intro sum.reindex) (auto intro!: inj_on_upd_PiE)
also have "… = (∑(a, x)∈(S i × Pi⇩E I S). f a i * (∏i∈I. f (x i) i))"
using insert by (force intro!: sum.cong prod.cong arg_cong2[where f="(*)"])
also have "… = (∑a∈S i. f a i * (∑x∈Pi⇩E I S. ∏i∈I. f (x i) i))"
finally show ?case
using insert by (simp add: sum_distrib_right)
qed

datatype state = listening | sending | sleeping

type_synonym sys_state = "(nat × nat) ⇒ state"

lemma state_UNIV: "UNIV = {listening, sending, sleeping}"
by (auto intro: state.exhaust)

fixes size :: nat and p :: real
assumes size: "0 < size"
assumes p: "0 < p" "p < 1"
begin

interpretation pmf_as_function .

definition states :: "sys_state set" where
"states = ({..< size} × {..< size}) →⇩E {listening, sending, sleeping}"

definition start :: sys_state where
"start = (λx∈{..< size}×{..< size}. listening)((0, 0) := sending)"

definition neighbour_sending where
"neighbour_sending s = (λ(x,y).
(x > 0 ∧ s (x - 1, y) = sending) ∨
(x < size ∧ s (x + 1, y) = sending) ∨
(y > 0 ∧ s (x, y - 1) = sending) ∨
(y < size ∧ s (x, y + 1) = sending))"

definition node_trans :: "sys_state ⇒ (nat × nat) ⇒ state ⇒ state ⇒ real" where
"node_trans g x s = (case s of
listening ⇒ (if neighbour_sending g x
then (λ_.0) (sending := p, sleeping := 1 - p)
else (λ_.0) (listening := 1))
| sending   ⇒ (λ_.0) (sleeping := 1)
| sleeping  ⇒ (λ_.0) (sleeping := 1))"

lemma node_trans_sum_eq_1[simp]:
"node_trans g x s' listening + (node_trans g x s' sending + node_trans g x s' sleeping) = 1"
by (simp add: node_trans_def split: state.split)

lemma node_trans_nonneg[simp]: "0 ≤ node_trans s x i j"
using p by (auto simp: node_trans_def split: state.split)

lift_definition proto_trans :: "sys_state ⇒ sys_state pmf" is
"λs s'. if s' ∈ states then (∏x∈{..< size}×{..< size}. node_trans s x (s x) (s' x)) else 0"
proof
let ?f = "λs s'. if s' ∈ states then (∏x∈{..< size}×{..< size}. node_trans s x (s x) (s' x)) else 0"
fix s show "∀t. 0 ≤ ?f s t"
using p by (auto intro!: prod_nonneg simp: node_trans_def split: state.split)
show "(∫⇧+t. ?f s t ∂count_space UNIV) = 1"
apply (subst nn_integral_count_space'[of states])
proof -
show "(∑x∈states. ∏xa∈{..<size} × {..<size}. node_trans s xa (s xa) (x xa)) = 1"
unfolding states_def by (subst sum_folded_product) simp_all
show "finite states"
by (auto simp: states_def intro!: finite_PiE)
qed
qed

end

subsection ‹The Gossip-Broadcast forms a DTMC›

sublocale gossip_broadcast ⊆ MC_syntax proto_trans .

end
```