From Stdlib Require Export Reals Psatz.
From iris.proofmode Require Import base proofmode.
From iris.base_logic.lib Require Export fancy_updates.
From iris.bi Require Export lib.fixpoint_mono big_op.
From iris.prelude Require Import options.

From clutch.bi Require Export weakestpre.
From clutch.prelude Require Import stdpp_ext iris_ext NNRbar.
From clutch.prob Require Export couplings distribution graded_predicate_lifting.
From clutch.delay_prob_lang Require Import lang urn_subst metatheory urn_erasable preserve_atomicity.
From clutch.common Require Export language.

Import uPred.

Local Open Scope NNR_scope.
#[global] Hint Resolve cond_nonneg : core.

Class eltonWpGS (Λ : language) (Σ : gFunctors) := EltonWpGS {
  eltonWpGS_invGS :: invGS_gen HasNoLc Σ;
  state_interp : state Λ → iProp Σ;
  err_interp : nonnegreal → iProp Σ;
}.
Global Opaque eltonWpGS_invGS.
Global Arguments EltonWpGS {Λ Σ}.
Canonical Structure NNRO := leibnizO nonnegreal.

Section rupd.
  Context `{H:!eltonWpGS d_prob_lang Σ}.
  (* Do we need to open invariants? *)
  Definition rupd_def (P: _ -> Prop) Q v : iProp Σ:=
    (∀ σ1, state_interp σ1 -∗
           ⌜∀ f, (urns_f_distr (urns σ1) f > 0)%R -> ∃ v', urn_subst_val f v = Some v' /\ P v'⌝ ∗ (Q ∗ state_interp σ1)).

  Local Definition rupd_aux : seal (@rupd_def). Proof. by eexists. Qed.
  Definition rupd := rupd_aux.(unseal).
  Lemma rupd_unseal : rupd = rupd_def.
  Proof. rewrite -rupd_aux.(seal_eq) //. Qed.

End rupd.

(* Definition total_weight (lis:list(gset nat)) := *)
(*   list_sum (size <*)

(* Definition list_weighted_distr' (lis:list (gset nat)) i:=  *)
(*   if (bool_decide (total_weight lis=0)R *)
(*   else match lis!!i with *)
(*        | Some s => ((size s)/(total_weight lis))*)
(*        | None => 0*)
(*        end. *)

Section modalities.
  Context `{eltonWpGS d_prob_lang Σ}.
  Implicit Types ε : nonnegreal.

  Definition state_step_coupl_pre Z Φ : (expr d_prob_lang * state d_prob_lang * nonnegreal -> iProp Σ) :=
    (λ x,
      let '(e, σ1, ε) := x in
      ⌜(1<=ε)%R⌝ ∨
        Z e σ1 ε ∨
        (∀ (ε':nonnegreal), ⌜(ε<ε')%R⌝ -∗ Φ (e, σ1, ε')) ∨
        (∃ μ (ε2 :_ -> nonnegreal),
            ⌜ exists r, forall ρ, (ε2 ρ <= r)%R ⌝ ∗
            ⌜ (Expval (μ) ε2 <= ε)%R ⌝ ∗
            ⌜urn_erasable μ σ1.(urns)⌝ ∗
            ∀ m', |={∅}=> Φ (e, state_upd_urns (λ _, m') σ1, ε2 m')
        ) ∨
        (∃ (K:ectx (d_prob_ectx_lang)) (v: val d_prob_lang) v',
            ⌜e = fill K (Val v)⌝ ∗
            ⌜∀ f, (urns_f_distr (urns σ1) f > 0)%R -> urn_subst_val f v = Some v' ⌝ ∗
            (|={∅}=> Φ (fill K (Val v'), σ1, ε))
        )
    )%I.


  #[local] Instance state_step_coupl_pre_ne Z Φ :
    NonExpansive (state_step_coupl_pre Z Φ).
  Proof.
    rewrite /state_step_coupl_pre.
    intros ?[[]?][[]?][[[=][=]][=]]. by simplify_eq.
  Qed.

  #[local] Instance state_step_coupl_pre_mono Z : BiMonoPred (state_step_coupl_pre Z).
  Proof.
    split; [|apply _].
    iIntros (Φ Ψ HNEΦ HNEΨ) "#Hwand".
    iIntros ([[]?]) "[H|[?|[H|[(%&%&%&%&%&H)|(%&%&%&%&%&H)]]]]".
    - by iLeft.
    - iRight; by iLeft.
    - iRight; iRight; iLeft.
      iIntros (??).
      iApply "Hwand".
      by iApply "H".
    - do 3 iRight. iLeft.
      repeat iExists _; repeat (iSplit; [done|]).
      iIntros (x).
      iDestruct ("H" $! x) as "H".
      iApply "Hwand".
      by iApply "H".
    - repeat iRight.
      iExists _, _, _. repeat iSplit; try done.
      iFrame.
      by iApply "Hwand".
  Qed.

  Definition state_step_coupl' Z := bi_least_fixpoint (state_step_coupl_pre Z).
  Definition state_step_coupl e σ ε Z:= state_step_coupl' Z (e, σ, ε).

  Lemma state_step_coupl_unfold e σ1 ε Z :
    state_step_coupl e σ1 ε Z ≡
      (⌜(1 <= ε)%R⌝ ∨
         (Z e σ1 ε) ∨
         (∀ ε', ⌜(ε<ε')%R⌝ -∗ state_step_coupl e σ1 ε' Z) ∨
         (∃ μ (ε2 :_ -> nonnegreal),
            ⌜ exists r, forall ρ, (ε2 ρ <= r)%R ⌝ ∗
            ⌜ (Expval (μ) ε2 <= ε)%R ⌝ ∗
            ⌜urn_erasable μ σ1.(urns)⌝ ∗
            ∀ m', |={∅}=> state_step_coupl e (state_upd_urns (λ _, m') σ1) (ε2 m') Z
         ) ∨
         (∃ (K:ectx (d_prob_ectx_lang)) (v: val d_prob_lang) v',
            ⌜e = fill K (Val v)⌝ ∗
            ⌜∀ f, (urns_f_distr (urns σ1) f > 0)%R -> urn_subst_val f v = Some v' ⌝ ∗
            (|={∅}=> state_step_coupl (fill K (Val v')) σ1 ε Z)
        )
      )%I.
  Proof. rewrite /state_step_coupl /state_step_coupl' least_fixpoint_unfold //. Qed.

  Lemma state_step_coupl_ret_err_ge_1 e σ1 Z (ε : nonnegreal) :
    (1 <= ε)%R → ⊢ state_step_coupl e σ1 ε Z.
  Proof. iIntros. rewrite state_step_coupl_unfold. by iLeft. Qed.

  Lemma state_step_coupl_ret e σ1 Z ε :
    Z e σ1 ε -∗ state_step_coupl e σ1 ε Z.
  Proof. iIntros. rewrite state_step_coupl_unfold. by iRight; iLeft. Qed.

  Lemma state_step_coupl_ampl e σ1 Z ε:
    (∀ ε', ⌜(ε<ε')%R⌝ -∗ state_step_coupl e σ1 ε' Z) -∗ state_step_coupl e σ1 ε Z.
  Proof. iIntros. rewrite state_step_coupl_unfold.
         do 2 iRight; by iLeft.
  Qed.

  Lemma state_step_coupl_ampl' e σ1 Z ε:
    (∀ ε', ⌜(ε<ε'<1)%R⌝ -∗ state_step_coupl e σ1 ε' Z) -∗ state_step_coupl e σ1 ε Z.
  Proof.
    iIntros "H".
    iApply state_step_coupl_ampl.
    iIntros (ε' ?).
    destruct (Rlt_or_le ε' 1)%R.
    - iApply "H". iPureIntro. lra.
    - by iApply state_step_coupl_ret_err_ge_1.
  Qed.

  Lemma state_step_coupl_rec e σ1 (ε : nonnegreal) Z :
    (∃ μ (ε2 :_ -> nonnegreal),
            ⌜ exists r, forall ρ, (ε2 ρ <= r)%R ⌝ ∗
            ⌜ (Expval (μ) ε2 <= ε)%R ⌝ ∗
            ⌜urn_erasable μ σ1.(urns)⌝ ∗
            ∀ m', |={∅}=> state_step_coupl e (state_upd_urns (λ _, m') σ1) (ε2 m') Z
        )
    ⊢ state_step_coupl e σ1 ε Z.
  Proof. iIntros "H". rewrite state_step_coupl_unfold. do 3 iRight. iLeft. done. Qed.

  Lemma state_step_coupl_value_promote e σ1 (ε : nonnegreal) Z:
     (∃ (K:ectx (d_prob_ectx_lang)) (v: val d_prob_lang) v',
            ⌜e = fill K (Val v)⌝ ∗
            ⌜∀ f, (urns_f_distr (urns σ1) f > 0)%R -> urn_subst_val f v = Some v' ⌝ ∗
            (|={∅}=> state_step_coupl (fill K (Val v')) σ1 ε Z)
     ) ⊢
     state_step_coupl e σ1 ε Z.
  Proof.
    iIntros "H". rewrite state_step_coupl_unfold. repeat iRight. done.
  Qed.

  Lemma state_step_coupl_rec_complete_split e σ1 (ε : nonnegreal) Z :
    (∃ u s (ε2 :_ -> nonnegreal),
        ⌜s≠∅⌝ ∗
        ⌜σ1.(urns) !! u = Some (urn_unif s) ⌝ ∗
        ⌜ exists r, forall ρ, (ε2 ρ <= r)%R ⌝ ∗
                    ⌜ (SeriesC (λ x, if bool_decide (x ∈ elements s) then ε2 x else 0%NNR)/ size s <= ε)%R ⌝ ∗
                    (∀ x, ⌜x∈s⌝ ={∅}=∗ state_step_coupl e (state_upd_urns <[u:=urn_unif {[x]}]> σ1) (ε2 x) Z )
    )%I
    ⊢ state_step_coupl e σ1 ε Z.
  Proof.
    iIntros "(%u&%s&%ε2&%Hs&%Hlookup&[%r %]&%Hineq&H)".
    iApply state_step_coupl_rec.
    destruct σ1 as [? us]; simpl; simpl in *.
    iExists _, (λ m, match ClassicalEpsilon.excluded_middle_informative
                             (∃ x, m=<[u:=urn_unif {[x]}]> us /\ x ∈ s) with
                     | left P =>ε2 (epsilon P)%NNR
                     | right _ => 1%NNR
                     end
               ).
    repeat iSplit.
    3:{ iPureIntro. by eapply complete_split_urn_erasable. }
    - iPureIntro.
      exists (Rmax 1 r).
      intros. case_match.
      + etrans; last apply Rmax_r. naive_solver.
      + apply Rmax_l.
    - iPureIntro.
      etrans; last exact.
      rewrite /Expval.
      rewrite /dbind/dbind_pmf{1}/pmf.
      setoid_rewrite <-SeriesC_scal_r.
      rewrite fubini_pos_seriesC'; last first.
      + setoid_rewrite Rmult_assoc.
        setoid_rewrite SeriesC_scal_l.
        apply (ex_seriesC_le _ (λ x, unif_set s x *Rmax 1 r))%R; last by apply ex_seriesC_scal_r.
        intros. split.
        * apply Rmult_le_pos; first done.
          apply SeriesC_ge_0'.
          real_solver.
        * apply Rmult_le_compat_l; first done.
          trans (SeriesC (λ x, if bool_decide (x=(<[u:=urn_unif {[n]}]> us) ) then Rmax 1 r else 0%R)); last (by rewrite SeriesC_singleton).
          apply SeriesC_le; last apply ex_seriesC_singleton.
          intros n'.
          split; first real_solver.
          case_bool_decide; subst.
          -- rewrite dret_1_1; last done. rewrite Rmult_1_l.
             rewrite Rmax_Rle.
             case_match; naive_solver.
          -- rewrite dret_0; first lra. done.
      + intros.
        setoid_rewrite Rmult_assoc.
        apply ex_seriesC_scal_l.
        apply (ex_seriesC_le _ (λ x, if bool_decide (x=(<[u:=urn_unif {[a]}]> us)) then Rmax 1 r else 0)); last apply ex_seriesC_singleton.
        intros. split; first real_solver.
        case_bool_decide.
        * rewrite dret_1_1; last done.
          subst.
          rewrite Rmult_1_l.
          case_match;
            rewrite Rmax_Rle; naive_solver.
        * rewrite dret_0; simpl; try lra. done.
      + intros. real_solver.
      + right.
        apply SeriesC_ext.
        intros a.
        setoid_rewrite Rmult_assoc.
        rewrite SeriesC_scal_l.
        case_bool_decide as H'.
        * rewrite unif_set_compute; last set_solver.
          rewrite Rmult_comm. f_equal.
          erewrite (SeriesC_ext _ (λ x, if bool_decide (x=(<[u:=urn_unif {[a]}]> us) ) then nonneg (ε2 a) else 0%R)); first by rewrite SeriesC_singleton.
          intros.
          case_bool_decide; subst.
          -- rewrite dret_1_1; last done.
             rewrite Rmult_1_l.
             case_match.
             ++ do 2 f_equal.
                rename select (∃ _, _) into He.
                pose proof epsilon_correct _ He as [He'].
                apply insert_inv in He'.
                assert (∀ x y, urn_unif x = urn_unif y -> x=y) as Hcut.
                { clear. intros. by simplify_eq. }
                apply Hcut in He'.
                set_solver.
             ++ exfalso.
                apply elem_of_elements in H'. naive_solver.
          -- rewrite dret_0; last done.
             lra.
        * rewrite unif_set_compute'; last set_solver. simpl; lra.
    - iIntros (m').
      case_match.
      + rename select (∃ _, _) into He.
        pose proof epsilon_correct _ He as [H2 ].
        simpl in *.
        iMod ("H" with "[//]").
        by rewrite -H2.
      + by iApply state_step_coupl_ret_err_ge_1.
  Qed.

  (* Lemma state_step_coupl_rec_equiv σ1 (ε : nonnegreal) Z : *)
  (*   (∃ μ (ε2 : state con_prob_lang -> nonnegreal), *)
  (*       ⌜ sch_erasable (λ t _ _ sch, TapeOblivious t sch) μ σ1 ⌝ ∗ *)
  (*       ⌜ exists r, forall ρ, (ε2 ρ <= r)*)
  (*                   ⌜ (Expval μ ε2 <= ε)*)
  (*                   ∀ σ2, |={∅}=> state_step_coupl σ2 (ε2 σ2) Z)*)
  (*   (∃ R μ (ε1 : nonnegreal) (ε2 : state con_prob_lang -> nonnegreal), *)
  (*       ⌜ sch_erasable (λ t _ _ sch, TapeOblivious t sch) μ σ1 ⌝ ∗ *)
  (*       ⌜ exists r, forall ρ, (ε2 ρ <= r)*)
  (*                   ⌜ (ε1 + Expval μ ε2 <= ε)*)
  (*                   ⌜pgl μ R ε1⌝ ∗ *)
  (*                   ∀ σ2, ⌜ R σ2 ⌝ ={∅}=∗ state_step_coupl σ2 (ε2 σ2) Z)*)
  (* Proof. *)
  (*   iSplit. *)
  (*   - iIntros "H". *)
  (*     iDestruct "H" as "(ε2 & r Hineq  &H)". *)
  (*     iExists (λ _, True), μ, 0, (λ σ, if bool_decide (μ σ > 0)*)
  (*     repeat iSplit; first done. *)
  (*     + iPureIntro. exists (Rmax r 1). *)
  (*       intros; case_bool_decide. *)
  (*       * etrans; last apply Rmax_l. done. *)
  (*       * simpl. apply Rmax_r. *)
  (*     + simpl. iPureIntro. rewrite Rplus_0_l. etrans; last exact. *)
  (*       rewrite /Expval. apply Req_le. *)
  (*       apply SeriesC_ext. *)
  (*       intros n. case_bool_decide; first done. *)
  (*       destruct (pmf_pos μ n) as K|K. *)
  (*       * exfalso. apply Rlt_gt in K. naive_solver. *)
  (*       * simpl. rewrite -K. lra. *)
  (*     + iPureIntro. by apply pgl_trivial. *)
  (*     + iIntros. case_bool_decide; first done. *)
  (*       by iApply state_step_coupl_ret_err_ge_1. *)
  (*   - iIntros "H". *)
  (*     iDestruct "H" as "(μ & ε2 & r Hineq & *)
  (*     iExists μ, (λ σ, if bool_decide(R σ) then ε2 σ else 1). *)
  (*     repeat iSplit; try done. *)
  (*     + iPureIntro. exists (Rmax r 1). *)
  (*       intros; case_bool_decide. *)
  (*       * etrans; last apply Rmax_l. done. *)
  (*       * simpl. apply Rmax_r. *)
  (*     + rewrite /Expval in Hineq. rewrite /Expval. iPureIntro. *)
  (*       rewrite /pgl/prob in Hpgl. etrans; last exact. *)
  (*       etrans; last (apply Rplus_le_compat_r; exact). *)
  (*       rewrite -SeriesC_plus. *)
  (*       * apply SeriesC_le. *)
  (*         -- intros. split; first (case_bool_decide; real_solver). *)
  (*            case_bool_decide; simpl; try lra. *)
  (*            rewrite Rmult_1_r. apply Rplus_le_0_compat. *)
  (*            real_solver. *)
  (*         -- apply ex_seriesC_plus. *)
  (*            ++ apply (ex_seriesC_le _ μ); last done. *)
  (*               intros. case_bool_decide; by simpl. *)
  (*            ++ apply pmf_ex_seriesC_mult_fn. naive_solver. *)
  (*       * apply (ex_seriesC_le _ μ); last done. *)
  (*         intros. case_bool_decide; by simpl. *)
  (*       * apply pmf_ex_seriesC_mult_fn. naive_solver. *)
  (*     + iIntros. case_bool_decide. *)
  (*       * iApply ("H" with "//"). *)
  (*       * by iApply state_step_coupl_ret_err_ge_1. *)
  (* Qed. *)

  (* Lemma state_step_coupl_rec' σ1 (ε : nonnegreal) Z : *)
  (*   (∃ R μ (ε1 : nonnegreal) (ε2 : state con_prob_lang -> nonnegreal), *)
  (*         ⌜ sch_erasable (λ t _ _ sch, TapeOblivious t sch) μ σ1 ⌝ ∗ *)
  (*         ⌜ exists r, forall ρ, (ε2 ρ <= r)*)
  (*         ⌜ (ε1 + Expval μ ε2 <= ε)*)
  (*         ⌜pgl μ R ε1⌝ ∗ *)
  (*         ∀ σ2, ⌜ R σ2 ⌝ ={∅}=∗ state_step_coupl σ2 (ε2 σ2) Z)*)
  (*   ⊢ state_step_coupl σ1 ε Z. *)
  (* Proof. *)
  (*   iIntros. iApply state_step_coupl_rec. by rewrite state_step_coupl_rec_equiv. *)
  (* Qed. *)

  Lemma state_step_coupl_ind (Ψ Z : _ -> state d_prob_lang → nonnegreal → iProp Σ) :
    ⊢ (□ (∀ e σ ε,
             state_step_coupl_pre Z (λ '(e, σ, ε'),
                 Ψ e σ ε' ∧ state_step_coupl e σ ε' Z)%I (e, σ, ε) -∗ Ψ e σ ε) →
       ∀ e σ ε, state_step_coupl e σ ε Z -∗ Ψ e σ ε)%I.
  Proof.
    iIntros "#IH" (e σ ε) "H".
    set (Ψ' := (λ '(e, σ, ε), Ψ e σ ε) :
           (prodO (prodO (exprO d_prob_lang )(stateO d_prob_lang)) NNRO) → iProp Σ).
    assert (NonExpansive Ψ').
    { intros n [[e1 σ1] ε1] [[e2 σ2] ε2].
      intros [[[=][=]][=]].
      by simplify_eq/=. }
    iApply (least_fixpoint_ind _ Ψ' with "[] H").
    iIntros "!#" ([[] ?]) "H". by iApply "IH".
  Qed.

  Lemma fupd_state_step_coupl e σ1 Z (ε : nonnegreal) :
    (|={∅}=> state_step_coupl e σ1 ε Z) ⊢ state_step_coupl e σ1 ε Z.
  Proof.
    iIntros "H".
    iApply state_step_coupl_rec.
    iExists (dret σ1.(urns)), (λ x, if bool_decide (x = σ1.(urns)) then ε else 1%NNR).
    repeat iSplit.
    - iExists (Rmax ε 1).
      iPureIntro. intros. case_bool_decide; [apply Rmax_l|apply Rmax_r].
    - iPureIntro. rewrite Expval_dret.
      by case_bool_decide.
    - iPureIntro. apply dret_urn_erasable.
    - iIntros (?).
      iMod "H".
      case_bool_decide.
      + subst; by destruct σ1.
      + by iApply state_step_coupl_ret_err_ge_1.
  Qed.

  Lemma state_step_coupl_mono e σ1 Z1 Z2 ε :
    (∀ e σ2 ε', Z1 e σ2 ε' -∗ Z2 e σ2 ε') -∗
    state_step_coupl e σ1 ε Z1 -∗ state_step_coupl e σ1 ε Z2.
  Proof.
    iIntros "HZ Hs".
    iRevert "HZ".
    iRevert (e σ1 ε) "Hs".
    iApply state_step_coupl_ind.
    iIntros "!#" (e σ ε)
      "[% | [? | [H|[(% & % & % & % & % & H)|H]]]] Hw".
    - iApply state_step_coupl_ret_err_ge_1. lra.
    - iApply state_step_coupl_ret. by iApply "Hw".
    - iApply state_step_coupl_ampl.
      iIntros (??).
      by iApply "H".
    - iApply state_step_coupl_rec.
      repeat iExists _.
      repeat iSplit; try done.
      iIntros (x).
      iDestruct ("H" $! x) as "H".
      iMod "H" as "[IH _]".
      by iApply "IH".
    - iApply state_step_coupl_value_promote.
      iDestruct "H" as "(%&%&%&%&%&H2)".
      iExists _, _, _. repeat iSplit; try done.
      iMod ("H2") as "[H2 _]".
      by iApply "H2".
  Qed.

  (* Prove this when generalize splitting condition *)
  Lemma state_step_coupl_mono_err e1 ε1 ε2 σ1 Z :
    (ε1 <= ε2)%R → state_step_coupl e1 σ1 ε1 Z -∗ state_step_coupl e1 σ1 ε2 Z.
  Proof.
    iIntros (Heps) "Hs".
    iApply state_step_coupl_rec.
    destruct σ1 as [? us].
    iExists (dret us), (λ x, if bool_decide (x=us) then ε1 else 1%NNR).
    repeat iSplit.
    - iPureIntro.
      exists (Rmax ε1 1).
      intros.
      case_bool_decide.
      + apply Rmax_l.
      + apply Rmax_r.
    - rewrite Expval_dret. by rewrite bool_decide_eq_true_2.
    - iPureIntro. apply dret_urn_erasable.
    - iIntros (?).
      iModIntro. case_bool_decide; subst; first done.
      by iApply state_step_coupl_ret_err_ge_1.
  Qed.

  Lemma state_step_coupl_bind e1 σ1 Z1 Z2 ε :
    (∀ e2 σ2 ε', Z1 e2 σ2 ε' -∗ state_step_coupl e2 σ2 ε' Z2) -∗
    state_step_coupl e1 σ1 ε Z1 -∗
    state_step_coupl e1 σ1 ε Z2.
  Proof.
    iIntros "HZ Hs".
    iRevert "HZ".
    iRevert (e1 σ1 ε) "Hs".
    iApply state_step_coupl_ind.
    iIntros "!#" (e σ ε)
      "[% | [H | [H|[(% & % & % & % & % & H)|H]]]] HZ".
    - by iApply state_step_coupl_ret_err_ge_1.
    - iApply ("HZ" with "H").
    - iApply state_step_coupl_ampl.
      iIntros (??).
      iDestruct ("H" with "[//]") as "[H _]".
      by iApply "H".
    - iApply state_step_coupl_rec.
      repeat iExists _.
      repeat iSplit; try done.
      iIntros (x).
      iDestruct ("H" $! x) as "H".
      iMod ("H") as "[H _]".
      by iApply "H".
    - iDestruct "H" as "(%&%&%&%&%&H2)".
      iApply state_step_coupl_value_promote.
      iExists _, _, _.
      repeat iSplit; try done.
      iMod ("H2") as "[H2 _]".
      by iApply "H2".
  Qed.

  Lemma state_step_coupl_rec_two_split e σ1 (ε : nonnegreal) Z :
    (∃ u s (ε2 :_ -> nonnegreal) s1 s2,
        ⌜s1≠∅⌝ ∗
        ⌜s2≠∅⌝ ∗
        ⌜s1 ∪ s2 =s⌝ ∗
        ⌜ s1 ≠ s2 ⌝ ∗
        ⌜ s1 ## s2 ⌝ ∗
        ⌜σ1.(urns) !! u = Some (urn_unif s) ⌝ ∗
        ⌜ exists r, forall ρ, (ε2 ρ <= r)%R ⌝ ∗
                    ⌜ (SeriesC (λ x, if bool_decide (x ∈ [s1; s2]) then ε2 x * size x else 0%NNR)/ size s <= ε)%R ⌝ ∗
                    (∀ x, ⌜x=s1 ∨ x=s2⌝ ={∅}=∗ state_step_coupl e (state_upd_urns <[u:=urn_unif x]> σ1) (ε2 x) Z )
    )%I
    ⊢ state_step_coupl e σ1 ε Z.
  Proof.
    iIntros "(%u&%s&%ε2&%s1&%s2&%Hnonempty1&%Hnonempty2&<-&%Hneq&%Hdisjoint&%Hlookup&[%r %Hbound]&%Hineq&H)".
    iApply state_step_coupl_rec.
    destruct σ1 as [? us].
    iExists _, (λ m, match ClassicalEpsilon.excluded_middle_informative
                             (∃ x, m=<[u:=urn_unif x]> us /\ (x =s1 \/ x=s2)) with
                     | left P =>ε2 (epsilon P)%NNR
                     | right _ => 1%NNR
                     end
               ).
    repeat iSplit.
    3:{ iPureIntro. apply two_split_urn_erasable; [| |done|..]; done. }
    - iPureIntro.
      exists (Rmax 1 r).
      intros. case_match.
      + etrans; last apply Rmax_r. naive_solver.
      + apply Rmax_l.
    - iPureIntro.
      etrans; last exact.
      rewrite /Expval.
      rewrite /dbind/dbind_pmf{1}/pmf.
      setoid_rewrite <-SeriesC_scal_r.
      rewrite fubini_pos_seriesC'; last first.
      + setoid_rewrite Rmult_assoc.
        setoid_rewrite SeriesC_scal_l.
        eapply (ex_seriesC_le _ (λ x, biased_coin _ _ x *Rmax 1 r))%R; last eapply ex_seriesC_scal_r; last apply ex_seriesC_finite.
        intros n. split.
        * apply Rmult_le_pos; first done.
          apply SeriesC_ge_0'.
          real_solver.
        * apply Rmult_le_compat_l; first done.
          simpl.
          trans (SeriesC (λ x,
                            if bool_decide (x= if n then (<[u:=urn_unif s1]> us) else (<[u:=urn_unif s2]> us)) then Rmax 1 r else 0%R)); last (by rewrite SeriesC_singleton).
          apply SeriesC_le; last apply ex_seriesC_singleton.
          intros n'.
          split; first real_solver.
          case_bool_decide; subst.
          -- rewrite dret_1_1; last done. rewrite Rmult_1_l.
             rewrite Rmax_Rle.
             case_match; naive_solver.
          -- rewrite dret_0; first lra. done.
      + intros n.
        setoid_rewrite Rmult_assoc.
        apply ex_seriesC_scal_l.
        apply (ex_seriesC_le _ (λ x, if bool_decide (x=if n then (<[u:=urn_unif s1]> us) else (<[u:=urn_unif s2]> us)) then Rmax 1 r else 0)); last apply ex_seriesC_singleton.
        intros. split; first real_solver.
        case_bool_decide.
        * rewrite dret_1_1; last done.
          subst.
          rewrite Rmult_1_l.
          case_match;
            rewrite Rmax_Rle; naive_solver.
        * rewrite dret_0; simpl; try lra. done.
      + intros. real_solver.
      + right.
        rewrite SeriesC_scal_r.
        rewrite SeriesC_list/=; last first.
        { repeat setoid_rewrite NoDup_cons.
          repeat split; last (by apply NoDup_nil); set_solver.
        }
        rewrite SeriesC_finite_foldr/=.
        do 2 setoid_rewrite Rmult_assoc.
        rewrite !SeriesC_scal_l/=.
        rewrite /biased_coin{2}/pmf/biased_coin_pmf.
        erewrite (SeriesC_ext _ (λ x, if bool_decide (x=(<[u:=urn_unif s1]> us)) then ε2 s1 else 0)); last first.
        { intros. case_bool_decide.
          - subst.
            rewrite dret_1_1; last done.
            case_match.
            + pose proof epsilon_correct _ e0 as [H1].
              simpl in *.
              apply insert_inv in H1.
              simplify_eq. rewrite -H1. lra.
            + exfalso.
              naive_solver.
          - rewrite dret_0; last done. simpl. lra.
        }
        rewrite SeriesC_singleton.
        erewrite (SeriesC_ext _ (λ x, if bool_decide (x=(<[u:=urn_unif s2]> us)) then ε2 s2 else 0)); last first.
        { intros. case_bool_decide.
          - subst.
            rewrite dret_1_1; last done.
            case_match.
            + pose proof epsilon_correct _ e0 as [H1].
              simpl in *.
              apply insert_inv in H1.
              simplify_eq. rewrite -H1. lra.
            + exfalso.
              naive_solver.
          - rewrite dret_0; last done. simpl. lra.
        }
        rewrite SeriesC_singleton.
        rewrite -Rdiv_def.
        rewrite Rdiv_plus_distr.
        rewrite !Rdiv_def.
        f_equal; first lra.
        rewrite !Rplus_0_r.
        rewrite Rmult_comm. rewrite Rmult_assoc. f_equal.
        rewrite size_union; last done.
        rewrite -(Rdiv_diag (size s1 + size s2)); last first.
        * rewrite -plus_INR.
          apply not_0_INR.
          assert (0<size s1); last lia.
          destruct (size s1) eqn:?; last lia.
          exfalso.
          apply size_empty_inv in Heqn.
          rewrite leibniz_equiv_iff in Heqn. naive_solver.
        * rewrite plus_INR. rewrite Rdiv_def.
          rewrite -Rcomplements.Rmult_minus_distr_r.
          by rewrite Rplus_minus_l.
    - iIntros (m').
      case_match.
      + rename select (∃ _, _) into He.
        pose proof epsilon_correct _ He as [H2 ].
        simpl in *.
        iMod ("H" with "[//]").
        by rewrite -H2.
      + by iApply state_step_coupl_ret_err_ge_1.
        Unshelve.
        { split.
         - apply Rcomplements.Rdiv_le_0_compat.
           + apply pos_INR.
           + apply lt_0_INR.
             rewrite size_union; last done.
             assert (0<size s1); last lia.
             destruct (size s1) eqn:?; last lia.
             exfalso.
             apply size_empty_inv in Heqn.
             rewrite leibniz_equiv_iff in Heqn. naive_solver.
         - rewrite -Rcomplements.Rdiv_le_1.
           + apply le_INR. rewrite size_union; first lia. done.
           + apply lt_0_INR.
             rewrite size_union; last done.
             assert (0<size s1); last lia.
             destruct (size s1) eqn:?; last lia.
             exfalso.
             apply size_empty_inv in Heqn.
             rewrite leibniz_equiv_iff in Heqn. naive_solver.
        }
  Qed.

  Lemma state_step_coupl_rec_partial_split e σ1 (ε : nonnegreal) Z :
    (∃ u s (ε2 :_ -> nonnegreal) lis,
        ⌜s≠∅⌝ ∗
        ⌜⋃ lis =s⌝ ∗
        ⌜ NoDup lis ⌝ ∗
        ⌜ ∅ ∉ lis ⌝ ∗
        ⌜∀ x y, x∈ lis -> y ∈ lis -> ∀ z, z ∈ x -> z ∈ y -> x=y⌝ ∗
        ⌜σ1.(urns) !! u = Some (urn_unif s) ⌝ ∗
        ⌜ exists r, forall ρ, (ε2 ρ <= r)%R ⌝ ∗
                    ⌜ (SeriesC (λ x, if bool_decide (x ∈ lis) then ε2 x * size x else 0%NNR)/ size s <= ε)%R ⌝ ∗
                    (∀ x, ⌜x∈lis⌝ ={∅}=∗ state_step_coupl e (state_upd_urns <[u:=urn_unif x]> σ1) (ε2 x) Z )
    )%I
    ⊢ state_step_coupl e σ1 ε Z.
  Proof.
    iIntros "(%l&%s&%ε2&%lis &%Hnonempty&%Hunion&%Hnodup&%Hnonempty'&%Hdisjoint&%Hlookup&%Hbound &%Hineq &H)".
    iInduction (lis) as [|hd tl IHl] forall (σ1 ε Z s ε2 Hnonempty Hunion Hnodup Hnonempty' Hdisjoint Hlookup Hbound Hineq).
    { simpl in *. set_solver. }
    destruct tl as [|hd' tl].
    { simpl.
      simpl in *.
      assert (s = hd) as -> by set_solver.
      replace (σ1) with (state_upd_urns <[l:=urn_unif hd]> σ1) at 1.
      - iApply fupd_state_step_coupl. iApply state_step_coupl_mono_err; last iMod ("H" with "[]") as "$"; try done.
        + etrans; last exact.
          rewrite SeriesC_list; simpl; last apply NoDup_singleton.
          rewrite Rplus_0_r.
          rewrite Rmult_div_l; first done.
          apply not_0_INR.
          apply size_non_empty_iff.
          by rewrite leibniz_equiv_iff.
        + iPureIntro. set_solver.
      - by apply state_upd_urns_no_change.
    }
    iApply state_step_coupl_rec_two_split.
    pose (ε2' x := ((if bool_decide (x=hd) then (ε2 x) else 0%R)+
                     if bool_decide (x=⋃(hd' :: tl)) then
                       SeriesC
                         (λ x , if bool_decide (x ∈ hd' :: tl) then ε2 x * size x else 0)%R / size (⋃ (hd'::tl)) else 0
                   )%R).
    assert (∀ x, 0<=ε2' x)%R as Hε2'.
    { intros.
      rewrite /ε2'.
      apply Rplus_le_le_0_compat; first by case_bool_decide.
      case_bool_decide; last done.
      apply Rcomplements.Rdiv_le_0_compat.
      - apply SeriesC_ge_0'.
        intros. case_bool_decide; real_solver.
      - apply lt_0_INR.
        simpl.
        rewrite size_union_alt.
        assert (0<size hd'); last lia.
        destruct (size hd') eqn:Hcontra; last lia.
        exfalso.
        apply Hnonempty'.
        apply size_empty_inv in Hcontra.
        rewrite leibniz_equiv_iff in Hcontra.
        set_solver.
    }
    assert (hd ≠ ⋃ (hd' :: tl)).
    { intros ->.
      repeat setoid_rewrite NoDup_cons in Hnodup.
      destruct Hnodup as [Hnodup1 ].
      rewrite elem_of_cons in Hnodup1.
      apply Hnodup1.
      left.
      assert (hd'≠∅) as H1.
      { intros ->. set_solver. }
      apply set_choose_L in H1 as [].
      eapply Hdisjoint; last done; simpl; set_solver. }
    assert (hd ## ⋃ (hd' :: tl)).
    { intros ? H' H1.
      repeat setoid_rewrite NoDup_cons in Hnodup.
      destruct Hnodup as [Hnodup1 ].
      rewrite elem_of_cons in Hnodup1.
      apply Hnodup1.
      simpl in H1.
      rewrite elem_of_union in H1.
      destruct!/=.
      - left.
        eapply Hdisjoint; last done; simpl; set_solver.
      - right.
        rewrite elem_of_union_list in H1.
        destruct H1 as (?&H1&H2').
        eapply Hdisjoint in H2'; last apply H'; [|set_solver..].
        by subst. }
    iExists l, s, (λ x, mknonnegreal _ (Hε2' x)), hd, (⋃(hd'::tl)).
    repeat iSplit.
    - iPureIntro. intros ->. set_solver.
    - iPureIntro. rewrite empty_union_list_L.
      rewrite Forall_cons.
      intros [->]. set_solver.
    - iPureIntro. set_solver.
    - done.
    - done.
    - done.
    - rewrite /ε2'.
      simpl.
      iPureIntro.
      eexists ((ε2 hd) + SeriesC (λ x : gset BinNums.Z, if bool_decide (x ∈ hd' :: tl) then ε2 x * size x else 0) /
                           size (hd' ∪ ⋃ tl))%R.
      intros.
      apply Rplus_le_compat.
      + case_bool_decide; by subst.
      + case_bool_decide; first done.
        apply Rcomplements.Rdiv_le_0_compat.
        * apply SeriesC_ge_0'.
          intros. case_bool_decide; last done.
          apply Rmult_le_pos; first done.
          apply pos_INR.
        * apply lt_0_INR.
          rewrite size_union_alt.
          assert (0<size hd')%nat; last lia.
          destruct (size hd') eqn:Hcontra; last lia.
          apply size_non_empty_iff in Hcontra; first done.
          rewrite leibniz_equiv_iff. intros ->. set_solver.
    - iPureIntro.
      assert (0<size s)%R.
      { apply lt_0_INR.
        destruct (size s) eqn:Hcontra; last lia.
        apply size_non_empty_iff in Hcontra; first done.
        rewrite leibniz_equiv_iff. intros ->. set_solver. }
      rewrite Rcomplements.Rle_div_l in Hineq; last lra.
      rewrite Rcomplements.Rle_div_l; last lra.
      etrans; last exact.
      rewrite !SeriesC_list; last first.
      + repeat setoid_rewrite NoDup_cons.
        repeat split; last by apply NoDup_nil.
        * by rewrite list_elem_of_singleton.
        * set_solver.
      + done.
      + simpl.
        rewrite /ε2'.
        rewrite bool_decide_eq_true_2; last done.
        rewrite bool_decide_eq_false_2; last first.
        { apply NoDup_cons in Hnodup. naive_solver. }
        rewrite Rplus_0_r.
        apply Rplus_le_compat_l.
        rewrite bool_decide_eq_false_2; last first.
        { intros <-.
          set_solver.
        }
        rewrite bool_decide_eq_true_2; last done.
        rewrite Rplus_0_l.
        rewrite -Rmult_div_swap.
        rewrite Rmult_div_l; last first.
        { apply not_0_INR.
          intros Hcontra.
          apply size_non_empty_iff in Hcontra; first done.
          rewrite leibniz_equiv_iff.
          apply union_positive_l_alt_L.
          intros ->. set_solver.
        }
        rewrite Rplus_0_r.
        rewrite SeriesC_list; last (apply NoDup_cons in Hnodup; naive_solver).
        done.
    - iIntros (s' [?|?]); subst.
      + rewrite /ε2'.
        iMod ("H" with "[]"); last iApply (state_step_coupl_mono_err with "[$]").
        * iPureIntro. set_solver.
        * simpl.
          rewrite bool_decide_eq_true_2; last done.
          rewrite bool_decide_eq_false_2; last done.
          lra.
      + iDestruct ("IHl" $! (state_upd_urns <[l:=urn_unif (⋃ (hd' :: tl))]> σ1) with "[][][][][][][][][-]") as "H'"; last first.
        * done.
        * iIntros (??).
          iMod ("H" with "[]").
          -- iPureIntro. rewrite elem_of_cons. by right.
          -- by rewrite state_upd_urns_twice.
        * simpl.
          rewrite /ε2'.
          rewrite bool_decide_eq_false_2; last done.
          rewrite bool_decide_eq_true_2; last done.
          by rewrite Rplus_0_l.
        * done.
        * iPureIntro. by rewrite lookup_insert_eq.
        * iPureIntro.
          intros.
          eapply Hdisjoint; set_solver.
        * iPureIntro. set_solver.
        * iPureIntro. apply NoDup_cons in Hnodup. naive_solver.
        * done.
        * iPureIntro.
          apply union_positive_l_alt_L.
          intros ->.
          set_solver.
  Qed.

  Lemma state_step_coupl_preserve e σ ε Z:
    is_well_constructed_expr e = true ->
    expr_support_set e ⊆ urns_support_set (urns σ) ->
    map_Forall (λ (_ : loc) v , is_well_constructed_val v = true) (heap σ) ->
    map_Forall (λ (_ : loc) v , val_support_set v ⊆ urns_support_set (urns σ))
               (heap σ) ->
    state_step_coupl e σ ε Z -∗
    state_step_coupl e σ ε
      (λ e' σ' ε',
         ⌜is_well_constructed_expr e'= true⌝ ∗
         ⌜expr_support_set e' ⊆ urns_support_set (urns σ')⌝ ∗
                ⌜map_Forall (λ (_ : loc) v, is_well_constructed_val v = true) (heap σ')⌝ ∗
                ⌜map_Forall (λ (_ : loc) v, val_support_set v ⊆ urns_support_set (urns σ')) (heap σ')⌝ ∗
                Z e' σ' ε').
  Proof.
    intros H1 H2 H3 H4.
    iIntros "H".
    iRevert (H1 H2 H3 H4).
    iRevert "H".
    iRevert (e σ ε).
    iApply state_step_coupl_ind.
    iModIntro.
    iIntros (e σ ε) "H".
    iIntros (He Hsubset Hforall1 Hforall2).
    iDestruct "H" as "[%|[H|[H|[H|H]]]]".
    - by iApply state_step_coupl_ret_err_ge_1.
    - iApply state_step_coupl_ret. iFrame. iPureIntro. naive_solver.
    - iApply state_step_coupl_ampl.
      iIntros.
      iDestruct ("H" with "[//]") as "[H _]".
      by iApply "H".
    - iDestruct "H" as "(%μ&%ε2&[%r %]&%H2&%&H)".
      iApply state_step_coupl_rec.
      iExists μ, (λ m, if bool_decide (μ m > 0) then ε2 m else 1%NNR)%R.
      repeat iSplit.
      + iPureIntro. exists (Rmax r 1).
        intros. case_bool_decide; [|apply Rmax_r].
        etrans; last apply Rmax_l. naive_solver.
      + iPureIntro. erewrite Expval_support in H2.
        etrans; last exact.
        rewrite /Expval.
        right.
        apply SeriesC_ext.
        intros. repeat f_equal.
        by case_bool_decide.
      + done.
      + iIntros (x).
         iDestruct ("H" $! x) as "H".
         iMod "H" as "[H _]". iModIntro.
         case_bool_decide; last by iApply state_step_coupl_ret_err_ge_1.
         iApply "H"; iPureIntro.
         * done.
         * etrans; first exact.
           by erewrite <-urn_erasable_same_support_set.
         * done.
         * eapply map_Forall_impl; first done.
           simpl.
           intros. etrans; first exact.
           by erewrite <-urn_erasable_same_support_set.
    - iDestruct "H" as "(%&%&%&%&%H1&H2)".
      iApply state_step_coupl_value_promote.
      subst.
      repeat iExists _, _, _.
      iFrame.
      repeat iSplit; try done.
      iMod ("H2") as "[H2 _]".
      iApply ("H2"); iPureIntro; try done.
      + rewrite !is_well_constructed_fill in He *.
        rewrite !andb_true_iff in He *.
        split; last naive_solver.
        apply is_simple_val_well_constructed.
        destruct!/=.
        epose proof urns_f_distr_witness _ as [? H'].
        apply H1 in H'.
        by eapply urn_subst_val_is_simple.
      + subst.
        rewrite support_set_fill.
        rewrite support_set_fill in Hsubset.
        etrans; last exact.
        simpl.
        rewrite is_simple_val_support_set; first set_solver.
        epose proof urns_f_distr_witness _ as [? H'].
        apply H1 in H'.
        by eapply urn_subst_val_is_simple.
  Qed.

  Lemma state_step_coupl_ctx_bind K e1 σ1 Z ε:
    state_step_coupl e1 σ1 ε
      (λ e2 σ2 ε',
         state_step_coupl (fill K e2) σ2 ε' Z) -∗ state_step_coupl (fill K e1) σ1 ε Z.
  Proof.
    iRevert (e1 σ1 ε).
    iApply state_step_coupl_ind.
    iModIntro.
    iIntros (e σ ε) "H".
    iDestruct "H" as "[%|[H|[H|[H|H]]]]".
    - iApply state_step_coupl_ret_err_ge_1. lra.
    - done.
    - iApply state_step_coupl_ampl.
      iIntros.
      by iDestruct ("H" with "[//]") as "[H _]".
    - iDestruct ("H") as "(%&%&%&%&%&H)".
      iApply state_step_coupl_rec.
      repeat iExists _; repeat iSplit; try done.
      iIntros. by iMod ("H" $! _) as "[H _]".
    - iDestruct "H" as "(%&%&%&%&%&H2)".
      subst.
      iApply state_step_coupl_value_promote.
      iExists (comp_ectx K _), _, _.
      repeat iSplit; try done; try iFrame.
      + by rewrite fill_app.
      + rewrite fill_app.
        by iMod ("H2") as "[$ _]".
  Qed.

  Lemma state_step_coupl_preserve_to_val e1 σ1 ε Z:
    state_step_coupl e1 σ1 ε Z ⊣⊢ state_step_coupl e1 σ1 ε (λ e2 σ2 ε2, Z e2 σ2 ε2 ∗ ⌜ssrbool.isSome(to_val e1) = ssrbool.isSome(to_val e2)⌝).
  Proof.
    iSplit.
    - iRevert (e1 σ1 ε).
      iApply state_step_coupl_ind.
      iModIntro.
      iIntros (e σ ε) "H".
      iDestruct "H" as "[%|[H|[H|[H|H]]]]".
      + iApply state_step_coupl_ret_err_ge_1. lra.
      + iApply state_step_coupl_ret. by iFrame.
      + iApply state_step_coupl_ampl.
        iIntros.
        by iDestruct ("H" with "[//]") as "[H _]".
      + iDestruct ("H") as "(%&%&%&%&%&H)".
        iApply state_step_coupl_rec.
        repeat iExists _; repeat iSplit; try done.
        iIntros. by iMod ("H" $! _) as "[H _]".
      + iDestruct "H" as "(%&%v&%v'&%&%&H2)".
        subst.
        iApply state_step_coupl_value_promote.
        iExists _, _, _.
        repeat iSplit; try done; try iFrame.
        iMod ("H2") as "[H2 _]".
        assert (ssrbool.isSome$ to_val (fill K (Val v')) = ssrbool.isSome $ to_val (fill K (Val (v)))) as Hrewrite; last by rewrite Hrewrite.
        simpl.
        destruct K as [|e]; simpl; first done.
        rewrite fill_not_val; first rewrite fill_not_val; try done; by destruct e.
    - iApply state_step_coupl_mono.
      iIntros (???) "[$?]".
  Qed.

  Lemma state_step_coupl_preserve_atomic `{Hatomic: !Atomic StronglyAtomic e1} σ1 ε Z :
    state_step_coupl e1 σ1 ε Z ⊣⊢ state_step_coupl e1 σ1 ε (λ e2 σ2 ε2, Z e2 σ2 ε2 ∗ ⌜Atomic StronglyAtomic e2⌝).
  Proof.
    iSplit; last first.
    { iApply state_step_coupl_mono.
      iIntros (???) "[$?]". }
    iIntros "H".
    iRevert (Hatomic).
    iRevert (e1 σ1 ε) "H".
    iApply state_step_coupl_ind.
    iModIntro.
    iIntros (e σ ε) "H".
    iDestruct "H" as "[%|[H|[H|[H|H]]]]"; iIntros (Hatomic).
    - iApply state_step_coupl_ret_err_ge_1. lra.
    - iApply state_step_coupl_ret. by iFrame.
    - iApply state_step_coupl_ampl.
      iIntros.
      iDestruct ("H" with "[//]") as "[H _]".
      by iApply "H".
    - iDestruct ("H") as "(%&%&%&%&%&H)".
      iApply state_step_coupl_rec.
      repeat iExists _; repeat iSplit; try done.
      iIntros. iMod ("H" $! _) as "[H _]".
      by iApply "H".
    - iDestruct "H" as "(%&%v&%v'&%&%H1&H2)".
      subst.
      iApply state_step_coupl_value_promote.
      iExists _, _, _.
      repeat iSplit; try done; try iFrame.
      iMod ("H2") as "[H _]".
      iApply "H".
      iPureIntro.
      epose proof urns_f_distr_witness _ as [? H'].
      apply H1 in H'.
      by eapply value_promote_preserves_atomicity.
  Qed.

  (* Lemma state_step_coupl_state_step α σ1 Z (ε ε' : nonnegreal) : *)
  (*   α ∈ get_active σ1 → *)
  (*   (∃ R, ⌜pgl (state_step σ1 α) R ε⌝ ∗ *)
  (*         ∀ σ2 , ⌜R σ2 ⌝ ={∅}=∗ state_step_coupl σ2 ε' Z) *)
  (*   ⊢ state_step_coupl σ1 (ε + ε') Z. *)
  (* Proof. *)
  (*   iIntros (Hin) "(&H)". *)
  (*   iApply state_step_coupl_rec'. *)
  (*   iExists R, (state_step σ1 α), ε, (λ _, ε'). *)
  (*   repeat iSplit; try done. *)
  (*   - iPureIntro. *)
  (*     simpl in *. *)
  (*     rewrite elem_of_elements elem_of_dom in Hin. *)
  (*     destruct Hin. *)
  (*     by eapply state_step_sch_erasable. *)
  (*   - iPureIntro; eexists _; done. *)
  (*   - iPureIntro. rewrite Expval_const; last done. rewrite state_step_mass; simpl;lra|done.  *)
  (* Qed. *)

  (* Lemma state_step_coupl_iterM_state_adv_comp N α σ1 Z (ε : nonnegreal) : *)
  (*   (α ∈ get_active σ1 -> *)
  (*    (∃ R ε1 (ε2 : _ -> nonnegreal), *)
  (*       ⌜ exists r, forall ρ, (ε2 ρ <= r)*)
  (*       ⌜ (ε1 + Expval (iterM N (λ σ, state_step σ α) σ1) ε2 <= ε)*)
  (*       ⌜pgl (iterM N (λ σ, state_step σ α) σ1) R ε1⌝ ∗ *)
  (*       ∀ σ2, ⌜ R σ2 ⌝ ={∅}=∗ state_step_coupl σ2 (ε2 σ2) Z) *)
  (*     ⊢ state_step_coupl σ1 ε Z)*)
  (* Proof. *)
  (*   iIntros (Hin) "(ε1 &  &  & H)". *)
  (*   iApply state_step_coupl_rec'. *)
  (*   iExists R, _, ε1, ε2. *)
  (*   repeat iSplit; try done. *)
  (*   - iPureIntro. *)
  (*     simpl in *. *)
  (*     rewrite elem_of_elements elem_of_dom in Hin. *)
  (*     destruct Hin. *)
  (*     by eapply iterM_state_step_sch_erasable. *)
  (* Qed.  *)

  (* Lemma state_step_coupl_state_adv_comp α σ1 Z (ε : nonnegreal) : *)
  (*   (α ∈ get_active σ1 -> *)
  (*    (∃ R ε1 (ε2 : _ -> nonnegreal), *)
  (*       ⌜ exists r, forall ρ, (ε2 ρ <= r)*)
  (*       ⌜ (ε1 + Expval (state_step σ1 α) ε2 <= ε)*)
  (*       ⌜pgl (state_step σ1 α) R ε1⌝ ∗ *)
  (*       ∀ σ2, ⌜ R σ2 ⌝ ={∅}=∗ state_step_coupl σ2 (ε2 σ2) Z) *)
  (*     ⊢ state_step_coupl σ1 ε Z)*)
  (* Proof. *)
  (*   iIntros (Hin) "(ε1 &  &  & H)". *)
  (*   iApply (state_step_coupl_iterM_state_adv_comp 1*)
  (*   repeat iExists _. simpl. rewrite dret_id_right. *)
  (*   by repeat iSplit. *)
  (* Qed. *)

  Definition prog_coupl e1 σ1 ε Z : iProp Σ :=
    (∃ (ε2 : (expr d_prob_lang * state d_prob_lang) -> nonnegreal),
           ⌜reducible (e1, σ1)⌝ ∗
           ⌜∃ r, ∀ ρ, (ε2 ρ <= r)%R ⌝ ∗
           ⌜(Expval (prim_step e1 σ1) ε2 <= ε)%R⌝ ∗
           (∀ e2 σ2, |={∅}=>
                         Z e2 σ2 (ε2 (e2, σ2)))
       )%I.

  Definition prog_coupl_equiv1 e1 σ1 ε Z:
    (∀ e2 σ2, Z e2 σ2 1) -∗
    prog_coupl e1 σ1 ε Z -∗
    (∃ R (ε1 : nonnegreal) (ε2 : (expr d_prob_lang * state d_prob_lang) -> nonnegreal),
           ⌜reducible (e1, σ1)⌝ ∗
           ⌜∃ r, ∀ ρ, (ε2 ρ <= r)%R ⌝ ∗
           ⌜(ε1 + Expval (prim_step e1 σ1) ε2 <= ε)%R⌝ ∗
           ⌜pgl (prim_step e1 σ1) R ε1⌝ ∗
           (∀ e2 σ2, ⌜R (e2, σ2)⌝ ={∅}=∗
                         Z e2 σ2 (ε2 (e2, σ2)))
    )%I.
  Proof.
    rewrite /prog_coupl.
    iIntros "H1 H".
    iDestruct "H" as "(%ε2 & % & [%r %] & %Hineq &H)".
    iExists (λ _, True), 0, (λ x, if bool_decide (prim_step e1 σ1 x > 0)%R then ε2 x else 1).
    repeat iSplit; first done.
    - iPureIntro. exists (Rmax r 1).
      intros; case_bool_decide.
      + etrans; last apply Rmax_l. done.
      + simpl. apply Rmax_r.
    - simpl. iPureIntro. rewrite Rplus_0_l. etrans; last exact.
      rewrite /Expval. apply Req_le.
      apply SeriesC_ext.
      intros n. case_bool_decide; first done.
      destruct (pmf_pos (prim_step e1 σ1) n) as [K|K].
      + exfalso. apply Rlt_gt in K. naive_solver.
      + simpl. rewrite -K. lra.
    - iPureIntro. by apply pgl_trivial.
    - iIntros. case_bool_decide; done.
  Qed.

  Definition prog_coupl_equiv2 e1 σ1 ε Z:
    (∀ e2 σ2, Z e2 σ2 1) -∗
    (∃ R (ε1 : nonnegreal) (ε2 : (expr d_prob_lang * state d_prob_lang) -> nonnegreal),
        ⌜reducible (e1, σ1)⌝ ∗
        ⌜∃ r, ∀ ρ, (ε2 ρ <= r)%R ⌝ ∗
                   ⌜(ε1 + Expval (prim_step e1 σ1) ε2 <= ε)%R⌝ ∗
                   ⌜pgl (prim_step e1 σ1) R ε1⌝ ∗
                   (∀ e2 σ2, ⌜R (e2, σ2)⌝ ={∅}=∗
                                 Z e2 σ2 (ε2 (e2, σ2)))
    )%I-∗
    prog_coupl e1 σ1 ε Z.
  Proof.
    iIntros "H1 H".
    iDestruct "H" as "(%R & %ε1 & %ε2 & % & [%r %] & %Hineq & %Hpgl &H)".
    iExists (λ σ, if bool_decide(R σ) then ε2 σ else 1).
    repeat iSplit; try done.
    - iPureIntro. exists (Rmax r 1).
      intros; case_bool_decide.
      + etrans; last apply Rmax_l. done.
      + simpl. apply Rmax_r.
    - rewrite /Expval in Hineq. rewrite /Expval. iPureIntro.
      rewrite /pgl/prob in Hpgl. etrans; last exact.
      etrans; last (apply Rplus_le_compat_r; exact).
      rewrite -SeriesC_plus.
      + apply SeriesC_le.
        * intros. split; first (case_bool_decide; real_solver).
          case_bool_decide; simpl; try lra.
          rewrite Rmult_1_r. apply Rplus_le_0_compat.
          real_solver.
        * apply ex_seriesC_plus.
          -- apply (ex_seriesC_le _ (prim_step e1 σ1)); last done.
             intros. case_bool_decide; by simpl.
          -- apply pmf_ex_seriesC_mult_fn. naive_solver.
      + apply (ex_seriesC_le _ (prim_step e1 σ1)); last done.
        intros. case_bool_decide; by simpl.
      + apply pmf_ex_seriesC_mult_fn. naive_solver.
    - iIntros. case_bool_decide; last done.
      iApply ("H" with "[//]").
  Qed.

  Lemma prog_coupl_mono_err e σ Z ε ε':
    (ε<=ε')%R -> prog_coupl e σ ε Z -∗ prog_coupl e σ ε' Z.
  Proof.
    iIntros (?) "(%&%&%&%&H)".
    repeat iExists _.
    repeat iSplit; try done.
    iPureIntro.
    etrans; exact.
  Qed.

  Lemma prog_coupl_strong_mono e1 σ1 Z1 Z2 ε :
    □(∀ e2 σ2, Z2 e2 σ2 1) -∗
    (∀ e2 σ2 ε', ⌜∃ σ, (prim_step e1 σ (e2, σ2) > 0)%R⌝ ∗ Z1 e2 σ2 ε' -∗ Z2 e2 σ2 ε') -∗
    prog_coupl e1 σ1 ε Z1 -∗ prog_coupl e1 σ1 ε Z2.
  Proof.
    iIntros "#H1 Hm (%&%&[%r %]&%Hineq & Hcnt) /=".
    rewrite /prog_coupl.
    iExists (λ x, if bool_decide (∃ σ, prim_step e1 σ x >0)%R then ε2 x else 1).
    repeat iSplit.
    - done.
    - iPureIntro. exists (Rmax r 1).
      intros; case_bool_decide.
      + etrans; last apply Rmax_l. done.
      + simpl. apply Rmax_r.
    - rewrite /Expval in Hineq. rewrite /Expval. iPureIntro.
      rewrite /Expval. etrans; last exact. apply Req_le.
      apply SeriesC_ext.
      intros n. case_bool_decide; first done.
      destruct (pmf_pos (prim_step e1 σ1) n) as [K|K].
      + exfalso. apply Rlt_gt in K. naive_solver.
      + simpl. rewrite -K. lra.
    - simpl. iIntros (??).
      case_bool_decide; last done.
      iApply "Hm". iMod ("Hcnt" $! _ _).
      by iFrame.
  Qed.

  Lemma prog_coupl_mono e1 σ1 Z1 Z2 ε :
    (∀ e2 σ2 ε', Z1 e2 σ2 ε' -∗ Z2 e2 σ2 ε') -∗
    prog_coupl e1 σ1 ε Z1 -∗ prog_coupl e1 σ1 ε Z2.
  Proof.
    iIntros "Hm".
    rewrite /prog_coupl.
    iIntros "(%&%&%&%&?)".
    repeat iExists _; repeat iSplit; try done.
    iIntros. by iApply "Hm".
  Qed.
    Lemma prog_coupl_strengthen e1 σ1 Z ε :
    □(∀ e2 σ2, Z e2 σ2 1) -∗
    prog_coupl e1 σ1 ε Z -∗
    prog_coupl e1 σ1 ε (λ e2 σ2 ε', ⌜(∃ σ, (prim_step e1 σ (e2, σ2) > 0)%R)\/ (1<=ε')%R⌝ ∧ Z e2 σ2 ε').
  Proof.
    iIntros "#Hmono (%&%&[%r %]&%Hineq & Hcnt) /=".
    rewrite /prog_coupl.
    iExists (λ x, if bool_decide (∃ σ, prim_step e1 σ x >0)%R then ε2 x else 1).
    repeat iSplit.
    - done.
    - iPureIntro. exists (Rmax r 1).
      intros; case_bool_decide.
      + etrans; last apply Rmax_l. done.
      + simpl. apply Rmax_r.
    - rewrite /Expval in Hineq. rewrite /Expval. iPureIntro.
      rewrite /Expval. etrans; last exact. apply Req_le.
      apply SeriesC_ext.
      intros n. case_bool_decide; first done.
      destruct (pmf_pos (prim_step e1 σ1) n) as [K|K].
      + exfalso. apply Rlt_gt in K. naive_solver.
      + simpl. rewrite -K. lra.
    - simpl. iIntros (??).
      case_bool_decide.
      + iMod ("Hcnt" $! _ _).
        iFrame. iPureIntro. naive_solver.
      + iMod ("Hcnt" $! e2 σ2 ).
        iModIntro. iSplit; last iApply "Hmono".
        iPureIntro. naive_solver.
  Qed.

  Lemma prog_coupl_ctx_bind K e1 σ1 Z ε:
    to_val e1 = None ->
    □(∀ e2 σ2, Z e2 σ2 1) -∗
    prog_coupl e1 σ1 ε (λ e2 σ2 ε', Z (fill K e2) σ2 ε') -∗ prog_coupl (fill K e1) σ1 ε Z.
  Proof.
    iIntros (Hv) "#H' H".
    (* iDestruct (prog_coupl_strengthen with "$") as "H". *)
    (* { iModIntro. by iIntros. } *)
    iDestruct "H" as "(%ε2&%&[%r %]&%&H)".