Eval Dist lift Comp generate Seed bind simulate Q run' take At Index

lemma evalDist_liftComp_generateSeed_bind_simulateQ_run'_takeAtIndex    {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀]    [∀ i, SampleableType (spec₀.Range i)] [unifSpec ⊂ₒ spec₀]    [unifSpec ˡ⊂ₒ spec₀]    [IsUniformSpec spec₀]    (qc : ι₀ → ℕ) (js : List ι₀) (i₀ : ι₀) (k : ℕ)    {α : Type} (oa : OracleComp spec₀ α) :    𝒟[(do      let seed ← liftComp (generateSeed spec₀ qc js) spec₀      (simulateQ seededOracle oa).run' (seed.takeAtIndex i₀ k) : OracleComp spec₀ α)] =    𝒟[oa] := by  classical  revert qc js k  induction oa using OracleComp.inductionOn with  | pure x =>    intro qc js k; apply evalDist_ext; intro a; simp  | query_bind t mx ih =>    intro qc js k    simp only [simulateQ_bind, simulateQ_query, OracleQuery.cont_query,      OracleQuery.input_query, id_map]    apply evalDist_ext; intro x    have hrun' : ∀ s : QuerySeed spec₀,        (do let u ← seededOracle t; simulateQ seededOracle (mx u) :          StateT _ (OracleComp spec₀) α).run' s =        match s.pop t with        | none => liftM (query t) >>= fun u => (simulateQ seededOracle (mx u)).run' s        | some (u, s') => (simulateQ seededOracle (mx u)).run' s' :=      run'_bind_query_eq_pop t mx    simp_rw [hrun']    rw [probOutput_bind_eq_tsum]    simp_rw [probOutput_liftComp]    by_cases hpop_none : qc t * js.count t = 0 ∨ (t = i₀ ∧ k = 0)    · have hpop : ∀ s ∈ support (generateSeed spec₀ qc js),          (s.takeAtIndex i₀ k).pop t = none := by        intro s hs; rw [QuerySeed.pop_eq_none_iff]        rw [support_generateSeed (spec := spec₀)] at hs        rcases hpop_none with hcount | ⟨rfl, rfl⟩        · have hlen : (s t).length = qc t * js.count t := hs t          have hnil := List.eq_nil_of_length_eq_zero (hlen.trans hcount)          by_cases hti : t = i₀          · subst hti; simp [QuerySeed.takeAtIndex, hnil]          · rw [QuerySeed.takeAtIndex_apply_of_ne _ _ _ _ hti]; exact hnil        · simp [QuerySeed.takeAtIndex]      have step1 : ∀ s,          Pr[= s | generateSeed spec₀ qc js] *            Pr[= x | match (s.takeAtIndex i₀ k).pop t with              | none => liftM (query t) >>= fun u =>                  (simulateQ seededOracle (mx u)).run' (s.takeAtIndex i₀ k)              | some (u, s') => (simulateQ seededOracle (mx u)).run' s'] =          Pr[= s | generateSeed spec₀ qc js] *            Pr[= x | liftM (query t) >>= fun u =>                (simulateQ seededOracle (mx u)).run' (s.takeAtIndex i₀ k)] := by        intro s        by_cases hs : s ∈ support (generateSeed spec₀ qc js)        · simp [hpop s hs]        · simp [(probOutput_eq_zero_iff _ _).2 hs]      simp_rw [step1, probOutput_bind_eq_tsum (liftM (query t))]      simp_rw [← ENNReal.tsum_mul_left, mul_left_comm]      rw [ENNReal.tsum_comm]      simp_rw [ENNReal.tsum_mul_left]      congr 1; ext u; congr 1      rw [← probOutput_liftComp_generateSeed_bind_simulateQ_run'_takeAtIndex_eq_tsum]      exact congrFun (congrArg DFunLike.coe (ih u qc js k)) x    · push Not at hpop_none      obtain ⟨hcount_ne, htk⟩ := hpop_none      have hcount : 0 < qc t * js.count t := Nat.pos_of_ne_zero (by lia)      have hpop_take : ∀ s ∈ support (generateSeed spec₀ qc js),          ∃ u s', (s.takeAtIndex i₀ k).pop t = some (u, s') := by        intro s hs        have hne : s t ≠ [] :=          ne_nil_of_mem_support_generateSeed (spec := spec₀) (qc := qc) (js := js) s t hs hcount        by_cases hti : t = i₀        · have hk : 0 < k := Nat.pos_of_ne_zero (htk hti)          have htake_ne : (s.takeAtIndex i₀ k) t ≠ [] := by            rw [hti, QuerySeed.takeAtIndex_apply_self]            obtain ⟨a, l, hl⟩ := List.exists_cons_of_ne_nil (hti ▸ hne)            rw [hl, show k = (k - 1) + 1 from (Nat.succ_pred_eq_of_pos hk).symm,              List.take_succ_cons]            exact List.cons_ne_nil _ _          obtain ⟨u, us, hcons⟩ := List.exists_cons_of_ne_nil htake_ne          exact ⟨u, Function.update (s.takeAtIndex i₀ k) t us,            QuerySeed.pop_eq_some_of_cons _ _ u us hcons⟩        · obtain ⟨u, us, hcons⟩ := List.exists_cons_of_ne_nil hne          exact ⟨u, Function.update (s.takeAtIndex i₀ k) t us,            QuerySeed.pop_eq_some_of_cons _ _ u us              ((QuerySeed.takeAtIndex_apply_of_ne s i₀ k t hti).trans hcons)⟩      have step1 : ∀ s,          Pr[= s | generateSeed spec₀ qc js] *            Pr[= x | match (s.takeAtIndex i₀ k).pop t with              | none => liftM (query t) >>= fun u =>                  (simulateQ seededOracle (mx u)).run' (s.takeAtIndex i₀ k)              | some (u, s') => (simulateQ seededOracle (mx u)).run' s'] =          Pr[= s | generateSeed spec₀ qc js] *            (match (s.takeAtIndex i₀ k).pop t with              | none => 0              | some (u, s') =>                  Pr[= x | (simulateQ seededOracle (mx u)).run' s']) := by        intro s        by_cases hs : s ∈ support (generateSeed spec₀ qc js)        · obtain ⟨u, s', hpop⟩ := hpop_take s hs; simp [hpop]        · simp [(probOutput_eq_zero_iff _ _).2 hs]      simp_rw [step1]      rw [← (QuerySeed.prependValues_singleton_injective t).tsum_eq        (support_generateSeed_mul_subset_range_prependValues_aux qc js t hcount _)]      by_cases hti : t = i₀      · subst hti        have hk : 0 < k := Nat.pos_of_ne_zero (htk rfl)        simp only [QuerySeed.pop_takeAtIndex_prependValues_self _ _ _ hk]        rw [ENNReal.tsum_prod', probOutput_bind_eq_tsum]        congr 1; ext u        exact probOutput_prependValues_takeAtIndex_tsum_eq_query_mul qc js t t (k - 1)          (mx u) u x hcount (ih u _ js.dedup (k - 1))      · simp only [QuerySeed.pop_takeAtIndex_prependValues_of_ne _ _ _ _ hti]        rw [ENNReal.tsum_prod', probOutput_bind_eq_tsum]        congr 1; ext u        exact probOutput_prependValues_takeAtIndex_tsum_eq_query_mul qc js t i₀ k          (mx u) u x hcount (ih u _ js.dedup k)