ROMUpToBad

  funext t'
  by_cases h : t' = t
  · subst h
    simp only [cacheOr_apply, QueryCache.cacheQuery_self, Option.some_or]
  · simp only [cacheOr_apply, QueryCache.cacheQuery_of_ne _ _ h]

  funext t'
  by_cases h : t' = t
  · subst h
    simp only [cacheOr_apply, QueryCache.cacheQuery_self, Option.some_or, hc, hp,
      Option.none_or]
  · simp only [cacheOr_apply, QueryCache.cacheQuery_of_ne _ _ h]

  rw [unifFwdFlagImpl_run, unifFwdImpl_run]
  refine relTriple_map (relTriple_post_mono (relTriple_refl _) ?_)
  intro x y hxy
  exact ⟨hxy, hs⟩

  classical
  cases hcache : cache t with
  | some v =>
    simp only [QueryImpl.withProgramming_apply, StateT.run_mk, hcache, probOutput_pure,
      Prod.mk.injEq, and_true, QueryImpl.withCachingTrackingPolicy_apply, Bool.if_true_left,
      Bool.decide_eq_true, Functor.map_map]
  | none =>
    cases hpol : policy t with
    | none =>
      simp only [QueryImpl.withProgramming_apply, hpol, Functor.map_map, StateT.run_mk, hcache,
        QueryImpl.withCachingTrackingPolicy_apply, Option.isSome_none, Bool.false_eq_true,
        ↓reduceIte]
    | some v =>
      have hne : ∀ (w : hashSpec.Range t) (c : hashSpec.QueryCache),
          ((u, (cache', false)) : hashSpec.Range t × hashSpec.QueryCache × Bool)
            ≠ (w, (c, true)) := by
        intro w c hcontr
        injection hcontr with _ h2
        injection h2 with _ h3
        cases h3
      have hL : (QueryImpl.withProgramming uniformSampleImpl policy t).run (cache, false) =
          (pure (v, (cache.cacheQuery t v, true)) :
            ProbComp (hashSpec.Range t × hashSpec.QueryCache × Bool)) := by
        simp only [QueryImpl.withProgramming_apply, hpol, map_pure, StateT.run_mk, hcache]
      have hR : (QueryImpl.withCachingTrackingPolicy uniformSampleImpl policy t).run
            (cache, false) =
          (uniformSampleImpl t >>= fun u' => pure (u', (cache.cacheQuery t u', true)) :
            ProbComp (hashSpec.Range t × hashSpec.QueryCache × Bool)) := by
        simp only [QueryImpl.withCachingTrackingPolicy_apply, hpol, Option.isSome_some, ↓reduceIte,
          Functor.map_map, StateT.run_mk, hcache, bind_pure_comp]
      rw [hL, hR, probOutput_pure, if_neg (hne v _), probOutput_bind_eq_tsum]
      symm
      refine ENNReal.tsum_eq_zero.mpr (fun u' => ?_)
      rw [probOutput_pure, if_neg (hne u' _), mul_zero]

  refine OracleComp.ProgramLogic.Relational.tvDist_simulateQ_run_le_probEvent_output_bad
    (programmedROImpl policy) (trackingROImpl policy) oa cache₀ ?_ ?_ ?_
  · rintro (n | t) s u s'
    · rfl
    · exact probOutput_withProgramming_eq_withCachingTrackingPolicy_of_not_bad_output
        policy t s u s'
  · rintro (n | t) ⟨c, b⟩ hp z hz
    · cases (show b = true from hp)
      simp only [add_apply_inl, programmedROImpl, QueryImpl.add_apply_inl, unifFwdFlagImpl_run,
        support_map, ProbComp.support_uniformFin, Set.image_univ] at hz
      obtain ⟨x, rfl⟩ := hz
      rfl
    · cases (show b = true from hp)
      exact QueryImpl.withProgramming_bad_monotone (so := uniformSampleImpl) (policy := policy)
        t c z hz
  · rintro (n | t) ⟨c, b⟩ hp z hz
    · cases (show b = true from hp)
      simp only [add_apply_inl, trackingROImpl, QueryImpl.add_apply_inl, unifFwdFlagImpl_run,
        support_map, ProbComp.support_uniformFin, Set.image_univ] at hz
      obtain ⟨x, rfl⟩ := hz
      rfl
    · cases (show b = true from hp)
      exact QueryImpl.withCachingTrackingPolicy_bad_monotone (so := uniformSampleImpl)
        (policy := policy) t c z hz

  refine relTriple_simulateQ_run' (programmedROImpl policy) (roImpl hashSpec)
    (fun p c => c = cacheOr p.1 policy) oa ?_ (cache₀, false) _ rfl
  rintro (n | t) ⟨c, bad⟩ c' hc'
  · exact relTriple_unifFwdFlagImpl_unifFwdImpl (R := fun p c => c = cacheOr p.1 policy) n
      (c, bad) c' hc'
  · subst hc'
    simp only [programmedROImpl, roImpl, QueryImpl.add_apply_inr, OracleSpec.randomOracle]
    cases hcache : c t with
    | some v =>
      simp only [add_apply_inr, QueryImpl.withProgramming_apply, StateT.run_mk, hcache,
        QueryImpl.withCaching_apply, StateT.run_bind, StateT.run_get, pure_bind, cacheOr_apply,
        Option.some_or, StateT.run_pure, relTriple_iff_relWP, MAlgRelOrdered.relWP_pure, and_self]
    | none =>
      cases hpol : policy t with
      | some p =>
        simp only [add_apply_inr, QueryImpl.withProgramming_apply, hpol, map_pure, StateT.run_mk,
          hcache, QueryImpl.withCaching_apply, StateT.run_bind, StateT.run_get, pure_bind,
          cacheOr_apply, Option.or_some, Option.getD_none, StateT.run_pure, relTriple_iff_relWP,
          MAlgRelOrdered.relWP_pure, true_and]
        exact (cacheOr_cacheQuery_of_policy policy c t p hcache hpol).symm
      | none =>
        simp only [add_apply_inr, QueryImpl.withProgramming_apply, hpol, Functor.map_map,
          StateT.run_mk, hcache, QueryImpl.withCaching_apply, StateT.run_bind, StateT.run_get,
          pure_bind, cacheOr_apply, Option.or_self, StateT.run_monadLift, monadLift_self,
          bind_pure_comp, StateT.run_modifyGet]
        refine relTriple_map (relTriple_post_mono (relTriple_refl _) ?_)
        intro x y hxy
        cases hxy
        exact ⟨rfl, (cacheOr_cacheQuery policy c t x).symm⟩

  refine relTriple_simulateQ_run' (trackingROImpl policy) (roImpl hashSpec)
    (fun p c => c = p.1) oa ?_ (cache₀, false) _ rfl
  rintro (n | t) ⟨c, bad⟩ c' hc'
  · exact relTriple_unifFwdFlagImpl_unifFwdImpl (R := fun p c => c = p.1) n (c, bad) c' hc'
  · subst hc'
    simp only [trackingROImpl, roImpl, QueryImpl.add_apply_inr, OracleSpec.randomOracle]
    cases hcache : c' t with
    | some v =>
      simp only [add_apply_inr, QueryImpl.withCachingTrackingPolicy_apply, Bool.if_true_left,
        Bool.decide_eq_true, Functor.map_map, StateT.run_mk, hcache, QueryImpl.withCaching_apply,
        StateT.run_bind, StateT.run_get, pure_bind, StateT.run_pure, relTriple_iff_relWP,
        MAlgRelOrdered.relWP_pure, and_self]
    | none =>
      simp only [add_apply_inr, QueryImpl.withCachingTrackingPolicy_apply, Bool.if_true_left,
        Bool.decide_eq_true, Functor.map_map, StateT.run_mk, hcache, QueryImpl.withCaching_apply,
        StateT.run_bind, StateT.run_get, pure_bind, StateT.run_monadLift, monadLift_self,
        bind_pure_comp, StateT.run_modifyGet]
      refine relTriple_map (relTriple_post_mono (relTriple_refl _) ?_)
      intro x y hxy
      cases hxy
      exact ⟨rfl, rfl⟩

  have hB := evalDist_eq_of_relTriple_eqRel (relTriple_run'_programmedROImpl_roImpl policy oa cache₀)
  have hC := evalDist_eq_of_relTriple_eqRel (relTriple_run'_trackingROImpl_roImpl policy oa cache₀)
  calc tvDist ((simulateQ (roImpl hashSpec) oa).run' (cacheOr cache₀ policy))
          ((simulateQ (roImpl hashSpec) oa).run' cache₀)
      = tvDist ((simulateQ (programmedROImpl policy) oa).run' (cache₀, false))
          ((simulateQ (trackingROImpl policy) oa).run' (cache₀, false)) := by
        rw [tvDist, tvDist, hB, hC]
    _ ≤ tvDist ((simulateQ (programmedROImpl policy) oa).run (cache₀, false))
          ((simulateQ (trackingROImpl policy) oa).run (cache₀, false)) := by
        rw [StateT.run', StateT.run']
        exact tvDist_map_le Prod.fst _ _
    _ ≤ _ := tvDist_programmedROImpl_trackingROImpl_le policy oa cache₀

  refine OracleComp.ProgramLogic.Relational.probEvent_output_bad_eq'
    (programmedROImpl policy) (trackingROImpl policy) ?_ ?_ ?_ oa cache₀
  · rintro (n | t) s u s'
    · rfl
    · exact probOutput_withProgramming_eq_withCachingTrackingPolicy_of_not_bad_output
        policy t s u s'
  · rintro (n | t) ⟨c, b⟩ hp z hz
    · cases (show b = true from hp)
      simp only [add_apply_inl, programmedROImpl, QueryImpl.add_apply_inl, unifFwdFlagImpl_run,
        support_map, ProbComp.support_uniformFin, Set.image_univ] at hz
      obtain ⟨x, rfl⟩ := hz
      rfl
    · cases (show b = true from hp)
      exact QueryImpl.withProgramming_bad_monotone (so := uniformSampleImpl) (policy := policy)
        t c z hz
  · rintro (n | t) ⟨c, b⟩ hp z hz
    · cases (show b = true from hp)
      simp only [add_apply_inl, trackingROImpl, QueryImpl.add_apply_inl, unifFwdFlagImpl_run,
        support_map, ProbComp.support_uniformFin, Set.image_univ] at hz
      obtain ⟨x, rfl⟩ := hz
      rfl
    · cases (show b = true from hp)
      exact QueryImpl.withCachingTrackingPolicy_bad_monotone (so := uniformSampleImpl)
        (policy := policy) t c z hz

  funext t'
  by_cases h : t' = t
  · subst h
    simp only [cacheOr_apply, QueryCache.empty_apply, Option.none_or, pointPolicy, dite_true,
      QueryCache.cacheQuery_self]
  · simp only [cacheOr_apply, QueryCache.empty_apply, Option.none_or, pointPolicy, h, dite_false,
      QueryCache.cacheQuery_of_ne _ _ h]

  classical
  set f : γ → ProbComp Bool := fun x =>
    (simulateQ (roImpl hashSpec) (oa x)).run' ((∅ : hashSpec.QueryCache).cacheQuery (pt x) (val x))
  set g : γ → ProbComp Bool := fun x => (simulateQ (roImpl hashSpec) (oa x)).run' ∅
  set bad : γ → ℝ≥0∞ := fun x =>
    Pr[fun z : Bool × hashSpec.QueryCache × Bool => z.2.2 = true |
      (simulateQ (programmedROImpl (pointPolicy (pt x) (val x))) (oa x)).run (∅, false)]
  have hterm : ∀ x, tvDist (f x) (g x) ≤ (bad x).toReal := by
    intro x
    have := tvDist_run'_roImpl_cacheOr_le_probEvent_bad (pointPolicy (pt x) (val x)) (oa x) ∅
    rwa [cacheOr_empty_pointPolicy] at this
  have hbadGame : Pr[= true | badQueryGame pre oa pt val] = ∑' x, Pr[= x | pre] * bad x := by
    rw [badQueryGame, probOutput_bind_eq_tsum]
    refine tsum_congr fun x => ?_
    rw [probOutput_map]
  have hne_top : ∀ x, Pr[= x | pre] * bad x ≠ ⊤ := fun x =>
    ENNReal.mul_ne_top (ne_top_of_le_ne_top one_ne_top probOutput_le_one)
      (ne_top_of_le_ne_top one_ne_top probEvent_le_one)
  have hsum_ne_top : (∑' x, Pr[= x | pre] * bad x) ≠ ⊤ := by
    rw [← hbadGame]
    exact ne_top_of_le_ne_top one_ne_top probOutput_le_one
  have hsummable : Summable fun x => (Pr[= x | pre] * bad x).toReal :=
    ENNReal.summable_toReal hsum_ne_top
  calc ProbComp.boolDistAdvantage (pre >>= f) (pre >>= g)
      ≤ tvDist (pre >>= f) (pre >>= g) := abs_probOutput_toReal_sub_le_tvDist _ _
    _ ≤ ∑' x, Pr[= x | pre].toReal * tvDist (f x) (g x) := tvDist_bind_left_le pre f g
    _ ≤ ∑' x, (Pr[= x | pre] * bad x).toReal := by
        refine Summable.tsum_le_tsum (fun x => ?_) ?_ hsummable
        · rw [ENNReal.toReal_mul]
          exact mul_le_mul_of_nonneg_left (hterm x) ENNReal.toReal_nonneg
        · exact hsummable.of_nonneg_of_le
            (fun x => mul_nonneg ENNReal.toReal_nonneg (tvDist_nonneg _ _))
            (fun x => by
              rw [ENNReal.toReal_mul]
              exact mul_le_mul_of_nonneg_left (hterm x) ENNReal.toReal_nonneg)
    _ = (Pr[= true | badQueryGame pre oa pt val]).toReal := by
        rw [hbadGame, ENNReal.tsum_toReal_eq hne_top]

  rw [probEvent_bind_eq_tsum]
  calc ∑' x, Pr[= x | mx] * Pr[q | my x]
      ≤ ∑' x, Pr[= x | mx] * c := by
        refine ENNReal.tsum_le_tsum fun x => ?_
        by_cases hx : x ∈ support mx
        · exact mul_le_mul' le_rfl (h x hx)
        · simp [probOutput_eq_zero_of_not_mem_support hx]
    _ = (∑' x, Pr[= x | mx]) * c := ENNReal.tsum_mul_right
    _ ≤ 1 * c := mul_le_mul' tsum_probOutput_le_one le_rfl
    _ = c := one_mul c

  have hrel := relTriple_simulateQ_run (trackingROImpl (pointPolicy pt v)) (roImpl hashSpec)
    (fun p c => c = p.1 ∧ p.2 = (p.1 pt).isSome) oa ?_ (cache₀, false) cache₀ ⟨rfl, by simp [h₀]⟩
  · refine probEvent_le_of_relTriple hrel ?_
    rintro a b ⟨-, hc, hflag⟩ hbad
    show (b.2 pt).isSome = true
    rw [hc, ← hflag]
    exact hbad
  · rintro (n | t) ⟨c, bad⟩ c' ⟨hc', hbad⟩
    · exact relTriple_unifFwdFlagImpl_unifFwdImpl
        (R := fun p c => c = p.1 ∧ p.2 = (p.1 pt).isSome) n (c, bad) c' ⟨hc', hbad⟩
    · subst hc'
      simp only [trackingROImpl, roImpl, QueryImpl.add_apply_inr, OracleSpec.randomOracle]
      cases hcache : c' t with
      | some w =>
        simp only [add_apply_inr, QueryImpl.withCachingTrackingPolicy_apply, Functor.map_map,
          StateT.run_mk, hcache, QueryImpl.withCaching_apply, StateT.run_bind, StateT.run_get,
          pure_bind, StateT.run_pure, relTriple_iff_relWP, MAlgRelOrdered.relWP_pure]
        refine ⟨trivial, trivial, ?_⟩
        simpa using hbad
      | none =>
        simp only [add_apply_inr, QueryImpl.withCachingTrackingPolicy_apply, Functor.map_map,
          StateT.run_mk, hcache, QueryImpl.withCaching_apply, StateT.run_bind, StateT.run_get,
          pure_bind, StateT.run_monadLift, monadLift_self, bind_pure_comp, StateT.run_modifyGet]
        refine relTriple_map (relTriple_post_mono (relTriple_refl _) ?_)
        intro x y hxy
        cases hxy
        dsimp only at hbad ⊢
        refine ⟨rfl, rfl, ?_⟩
        by_cases ht : t = pt
        · subst ht
          simp [pointPolicy]
        · simp only [pointPolicy, ht, dite_false, Option.isSome_none, Bool.false_eq_true,
            if_false, QueryCache.cacheQuery_of_ne _ _ (Ne.symm ht)]
          exact hbad

  induction oa using OracleComp.inductionOn generalizing n cache₀ with
  | pure x =>
    intro z hz
    simp only [simulateQ_pure, StateT.run_pure, support_pure, Set.mem_singleton_iff] at hz
    subst hz
    exact ⟨∅, by simp, fun t h => Or.inl h⟩
  | query_bind t mx ih =>
    intro z hz
    rw [isQueryBoundP_query_bind_iff] at hq
    obtain ⟨hvalid, hrest⟩ := hq
    simp only [simulateQ_bind, simulateQ_query, OracleQuery.input_query, OracleQuery.cont_query,
      id_map, StateT.run_bind, support_bind, Set.mem_iUnion] at hz
    obtain ⟨⟨u, c'⟩, hstep, hz⟩ := hz
    rcases t with n' | t
    · have hrun : (roImpl hashSpec (Sum.inl n')).run cache₀ =
          ((fun x => (x, cache₀)) <$> ($[0..n']) :
            ProbComp ((unifSpec + hashSpec).Range (Sum.inl n') × hashSpec.QueryCache)) := rfl
      have hb := hrest u
      simp only [Sum.isRight_inl, Bool.false_eq_true, if_false] at hb
      rw [hrun] at hstep
      simp only [support_map, Set.mem_image, Prod.mk.injEq] at hstep
      obtain ⟨x, -, rfl, rfl⟩ := hstep
      exact ih _ n hb _ z hz
    · have hb := hrest u
      simp only [Sum.isRight_inr, if_true] at hb
      simp only [Sum.isRight_inr, not_true_eq_false, false_or] at hvalid
      cases hcache : cache₀ t with
      | some w =>
        have hrun : (roImpl hashSpec (Sum.inr t)).run cache₀ =
            (pure (w, cache₀) :
              ProbComp ((unifSpec + hashSpec).Range (Sum.inr t) × hashSpec.QueryCache)) := by
          show (hashSpec.randomOracle t).run cache₀ = _
          simp [StateT.run_bind, StateT.run_get, hcache]
        rw [hrun] at hstep
        simp only [support_pure, Set.mem_singleton_iff, Prod.mk.injEq] at hstep
        obtain ⟨rfl, rfl⟩ := hstep
        obtain ⟨S, hS, hkeys⟩ := ih _ (n - 1) hb _ z hz
        exact ⟨S, hS.trans (Nat.sub_le n 1), hkeys⟩
      | none =>
        have hrun : (roImpl hashSpec (Sum.inr t)).run cache₀ =
            (($ᵗ hashSpec.Range t) >>= fun x => pure (x, cache₀.cacheQuery t x) :
              ProbComp ((unifSpec + hashSpec).Range (Sum.inr t) × hashSpec.QueryCache)) := by
          show (hashSpec.randomOracle t).run cache₀ = _
          simp [StateT.run_bind, StateT.run_get, hcache, StateT.run_monadLift, monadLift_self,
            StateT.run_modifyGet, uniformSampleImpl]
        rw [hrun] at hstep
        simp only [support_bind, support_pure, Set.mem_iUnion, Set.mem_singleton_iff,
          Prod.mk.injEq] at hstep
        obtain ⟨x, -, rfl, rfl⟩ := hstep
        obtain ⟨S, hS, hkeys⟩ := ih _ (n - 1) hb _ z hz
        refine ⟨insert t S, ?_, ?_⟩
        · calc (insert t S).card ≤ S.card + 1 := Finset.card_insert_le _ _
            _ ≤ n := by omega
        · intro t' h
          rcases hkeys t' h with h' | h'
          · by_cases htt : t' = t
            · subst htt
              exact Or.inr (Finset.mem_insert_self _ _)
            · rw [QueryCache.isCached_cacheQuery_of_ne _ _ htt] at h'
              exact Or.inl h'
          · exact Or.inr (Finset.mem_insert_of_mem h')
 
omit [DecidableEq ι] [∀ t : hashSpec.Domain, SampleableType (hashSpec.Range t)] in

  calc Pr[fun y => QueryCache.isCached c (f y) = true | samp]
      ≤ Pr[fun y => ∃ t ∈ S, f y = t | samp] :=
        probEvent_mono'' fun y hy => ⟨f y, hS _ hy, rfl⟩
    _ ≤ ∑ t ∈ S, Pr[fun y => f y = t | samp] :=
        probEvent_exists_finset_le_sum S samp (fun t y => f y = t)
    _ ≤ ∑ _t ∈ S, beta := Finset.sum_le_sum fun t _ => hB t
    _ = S.card * beta := by rw [Finset.sum_const, nsmul_eq_mul]