Set Implicit Arguments.

Require Import Crypto.
Require Import Bernoulli.

(* stuff that needs to go somewhere else *)

Theorem ratSubtract_leRat_same_r :
  forall r1 r2 r3,
    r3 <= r1 ->
    r3 <= r2 ->
    ratSubtract r1 r3 <= ratSubtract r2 r3 ->
    r1 <= r2.

  intuition.
  apply (ratAdd_leRat_compat_l r3) in H1.
  assert (ratSubtract r1 r3 + r3 == r1).
  rewrite ratAdd_comm.
  eapply ratSubtract_ratAdd_inverse_2; intuition.
  rewrite <- H2.
  rewrite H1.
  rewrite ratAdd_comm.
  eapply eqRat_impl_leRat.
  apply ratSubtract_ratAdd_inverse_2; intuition.

Qed.

Theorem sumList_filter_complement :
  forall (A : Set){eqd : EqDec A}(c : Comp A)(a : A),
    well_formed_comp c ->
    sumList (filter (fun x => negb (eqb a x)) (getSupport c)) (evalDist c) ==
    ratSubtract 1 (evalDist c a).

  intuition.
  eapply (eqRat_ratAdd_same_r (evalDist c a)).
  symmetry.
  rewrite ratAdd_comm.
  rewrite ratSubtract_ratAdd_inverse_2.
  rewrite <- (evalDist_lossless H).
  rewrite (sumList_filter_partition (eqb a)).
  rewrite ratAdd_comm.
  eapply ratAdd_eqRat_compat; intuition.
  destruct (eq_Rat_dec (evalDist c a) 0).
  rewrite e.
  eapply sumList_0; intuition.
  apply filter_In in H0.
  intuition.
  rewrite eqb_leibniz in H2; subst.
  trivial.

  rewrite (@sumList_exactly_one _ a); intuition.
  eapply filter_NoDup.
  eapply getSupport_NoDup.
  eapply filter_In; intuition.
  eapply getSupport_In_evalDist; intuition.
  eapply eqb_refl.
  apply filter_In in H0; intuition.
  rewrite eqb_leibniz in H3; subst.
  intuition.

  eapply evalDist_le_1.
Qed.

Theorem evalDist_Repeat_0 :
  forall (A : Set)(c : Comp A)(P : A -> bool) a,
    (forall x, In x (getSupport c) -> P x = true -> x <> a) ->
    evalDist (Repeat c P) a == 0.

  intuition.
  simpl.
  unfold indicator.
  case_eq (P a); intuition.
  rewrite ratMult_1_l.
  eapply ratMult_0.
  right.
  eapply getSupport_not_In_evalDist.
  intuition.
  eapply H;
    eauto.

  rewrite ratMult_0_l.
  eapply ratMult_0_l.

Qed.

Theorem evalDist_event_equiv :
  forall (A : Set){eqd : EqDec A}(c : Comp A) a,
    evalDist c a == Pr[x <-$ c; ret (eqb a x)].

  intuition.

  destruct (in_dec (EqDec_dec _) a (getSupport c)).

  rewrite evalDist_seq_step.
  rewrite (@sumList_exactly_one _ a).
  dist_compute.
  eapply getSupport_NoDup.
  trivial.
  intuition.
  dist_compute.

  rewrite evalDist_seq_step.
  assert (evalDist c a == 0).
  eapply getSupport_not_In_evalDist; trivial.
  rewrite H.
  symmetry.
  eapply sumList_0.
  intuition.
  dist_compute.

Qed.

Theorem filter_ext : forall (A : Set)(ls : list A)(f1 f2 : A -> bool),
  (forall a, f1 a = f2 a) ->
  filter f1 ls = filter f2 ls.

  induction ls; intuition; simpl in *.
  rewrite H.
  destruct (f2 a); eauto.
  f_equal;
  eauto.

Qed.

(* definitions related to the program logic *)

Definition marginal_l(A B : Set){eqd : EqDec A}(c : Comp (A * B))(a : A) :=
  Pr[x <-$ c; ret eqb (fst x) a].

Definition marginal_r(A B : Set){eqd : EqDec B}(c : Comp (A * B))(b : B) :=
  Pr[x <-$ c; ret eqb (snd x) b].

Theorem in_support_marginal_l :
  forall (A B : Set){eqd: EqDec A}(c1 : Comp (A * B))(c2 : Comp A),
    (forall a, evalDist c2 a == marginal_l c1 a) ->
    forall p,
      In p (getSupport c1) ->
      In (fst p) (getSupport c2).

  intuition.
  eapply getSupport_In_evalDist.
  intuition.
  rewrite H in H1.
  clear H.
  unfold marginal_l in *.
  simpl in *.
  eapply sumList_0 in H1; eauto.
  eapply ratMult_0 in H1.
  intuition.
  eapply getSupport_In_evalDist;
    eauto.
  dist_compute.

Qed.

Theorem in_support_marginal_r :
  forall (A B : Set){eqd: EqDec B}(c1 : Comp (A * B))(c2 : Comp B),
    (forall b, evalDist c2 b == marginal_r c1 b) ->
    forall p,
      In p (getSupport c1) ->
      In (snd p) (getSupport c2).

  intuition.
  eapply getSupport_In_evalDist.
  intuition.
  rewrite H in H1.
  clear H.
  unfold marginal_r in *.
  simpl in *.
  eapply sumList_0 in H1; eauto.
  eapply ratMult_0 in H1.
  intuition.
  eapply getSupport_In_evalDist;
    eauto.
  dist_compute.

Qed.

Definition comp_spec (R1 R2 : Set)
  {eqd1 : EqDec R1}{eqd2 : EqDec R2}
  (post : R1 -> R2 -> Prop)(c1 : Comp R1)(c2 : Comp R2) :=
  exists c : Comp (R1 * R2),
    (forall r1, evalDist c1 r1 == marginal_l c r1) /\
    (forall r2, evalDist c2 r2 == marginal_r c r2) /\
    (forall p, In p (getSupport c) -> post (fst p) (snd p)).

Theorem list_choice :
  forall (A B : Type)(eqd : forall (a1 a2 : A), {a1 = a2} + {a1 <> a2}) P (ls : list A) (b : B),
    (forall a, In a ls -> exists (b : B), P a b) ->
    exists (f : A -> B), (forall a, In a ls -> P a (f a)).

  induction ls; intuition.
  exists (fun a => b).
  intuition.
  simpl in *.
  intuition.

  edestruct (H a); intuition.
  edestruct (H0); intuition.
  exists (fun a0 => if (eqd a a0) then x else (x0 a0)).
  intuition.
  inversion H3; clear H3; subst.
  destruct (eqd a0 a0); intuition.
  destruct (eqd a a0); intuition.
  subst.
  trivial.
Qed.

Theorem comp_spec_seq :
  forall {A B : Set} P' {C D : Set} P{eqda : EqDec A}{eqdb : EqDec B}{eqdc : EqDec C}{eqdd : EqDec D}(c1 : Comp A)(c2 : Comp B) (c : C) (d : D)
    (f1 : A -> Comp C)(f2 : B -> Comp D),
    comp_spec P' c1 c2 ->
    (forall a b, In a (getSupport c1) -> In b (getSupport c2) -> P' a b -> comp_spec P (f1 a) (f2 b)) ->
    comp_spec P (Bind c1 f1) (Bind c2 f2).

  intuition.
  destruct H; intuition.

  assert (forall p, In p (getSupport x) -> comp_spec P (f1 (fst p)) (f2 (snd p))).
  intuition.
  eapply H0.
  eapply in_support_marginal_l; eauto.
  eapply in_support_marginal_r; eauto.
  eauto.
  clear H0.

  pose proof H2.
  apply (list_choice) in H2.
  destruct H2.

  exists (Bind x x0).
  intuition.

  rewrite evalDist_seq_step.
  eapply eqRat_trans.
  eapply sumList_body_eq.
  intros.
  eapply ratMult_eqRat_compat.
  eauto.
  eapply eqRat_refl.
  unfold marginal_l.

  eapply eqRat_trans.
  eapply sumList_body_eq.
  intros.
  eapply ratMult_eqRat_compat.
  eapply evalDist_seq_step.
  eapply eqRat_refl.

  eapply eqRat_trans.
  eapply sumList_body_eq.
  intros.
  symmetry.
  eapply sumList_factor_constant_r.
  rewrite sumList_comm.

  dist_inline_first.
  rewrite evalDist_seq_step.
  eapply sumList_body_eq.
  intros.

  rewrite (@sumList_exactly_one _ (fst a)).

  assert (Pr  [ret eqb (fst a) (fst a) ]  == 1).
  dist_compute.
  rewrite H5.
  rewrite ratMult_1_r.
  eapply ratMult_eqRat_compat; intuition.
  specialize (H2 a).
  intuition.
  rewrite H2.

  unfold marginal_l.
  intuition.

  eapply getSupport_NoDup.

  eapply in_support_marginal_l; eauto.

  intuition.
  dist_compute.

  rewrite evalDist_seq_step.
  eapply eqRat_trans.
  eapply sumList_body_eq.
  intros.
  eapply ratMult_eqRat_compat.
  eauto.
  eapply eqRat_refl.
  unfold marginal_r.

  eapply eqRat_trans.
  eapply sumList_body_eq.
  intros.
  eapply ratMult_eqRat_compat.
  eapply evalDist_seq_step.
  eapply eqRat_refl.

  eapply eqRat_trans.
  eapply sumList_body_eq.
  intros.
  symmetry.
  eapply sumList_factor_constant_r.
  rewrite sumList_comm.

  dist_inline_first.
  rewrite evalDist_seq_step.
  eapply sumList_body_eq.
  intros.

  rewrite (@sumList_exactly_one _ (snd a)).
  assert (Pr  [ret eqb (snd a) (snd a) ] == 1).
  dist_compute.
  rewrite H5.
  rewrite ratMult_1_r.
  eapply ratMult_eqRat_compat; intuition.
  specialize (H2 a); intuition.
  rewrite H6.
  unfold marginal_r; intuition.

  eapply getSupport_NoDup.
  eapply in_support_marginal_r; eauto.

  intuition.
  dist_compute.

  simp_in_support.
  specialize (H2 x1); intuition.

  unfold eq_dec; intuition.
  eapply (EqDec_dec _).
  eapply (ret (c, d)).

Qed.

Ltac despec :=
  match goal with
    | [H : comp_spec _ _ _  |- _] => destruct H
  end.

Theorem comp_spec_consequence :
  forall (A B : Set){eqda1 eqda2 : EqDec A}{eqdb1 eqdb2 : EqDec B}(p1 p2 : A -> B -> Prop) c1 c2,
    (@comp_spec _ _ eqda1 eqdb1 p1 c1 c2) ->
    (forall a b, p1 a b -> p2 a b) ->
    (@comp_spec  _ _ eqda2 eqdb2 p2 c1 c2).

  intuition.
  despec; intuition.
  exists x; intuition.

  rewrite H1.
  unfold marginal_l.

  dist_skip.
  eapply evalDist_ret_eq.
  case_eq (eqb (fst x0) r1); intuition.
  rewrite eqb_leibniz in H4.
  subst.
  eapply eqb_refl.

  case_eq ( (@eqb A eqda1 (@fst A B x0) r1)); intuition.
  rewrite eqb_leibniz in H5.
  subst.
  rewrite eqb_refl in H4.
  intuition.

  rewrite H.
  unfold marginal_r.

  dist_skip.
  eapply evalDist_ret_eq.
  case_eq (eqb (snd x0) r2); intuition.
  rewrite eqb_leibniz in H4.
  subst.
  eapply eqb_refl.

  case_eq ( (@eqb B eqdb1 (@snd A B x0) r2)); intuition.
  rewrite eqb_leibniz in H5.
  subst.
  rewrite eqb_refl in H4.
  intuition.

Qed.

Theorem comp_spec_symm :
  forall (A B : Set){eqda : EqDec A}{eqdb : EqDec B}(p : A -> B -> Prop) c1 c2,
    comp_spec p c1 c2 ->
    comp_spec (fun b a => p a b) c2 c1.

  intuition.
  despec; intuition.
  exists (p <-$ x; ret (snd p, fst p)).
  intuition.
  eapply eqRat_trans; eauto.
  unfold marginal_r, marginal_l.
  dist_inline_first.
  dist_skip.
  dist_compute.

  eapply eqRat_trans; eauto.
  unfold marginal_r, marginal_l.
  dist_inline_first.
  dist_skip.
  dist_compute.

  repeat simp_in_support.
  simpl.
  eauto.
Qed.

(* completeness theorems *)
Theorem eq_impl_comp_spec :
  forall (A B : Set){eqda : EqDec A}{eqdb : EqDec B}(c1 : Comp A)(c2 : Comp B) x y,
    well_formed_comp c1 ->
    well_formed_comp c2 ->
    evalDist c1 x == evalDist c2 y ->
    comp_spec (fun a b => a = x <-> b = y) c1 c2.

  intuition.

  (* special case: the events have probability 0 *)
  destruct (eq_Rat_dec (evalDist c1 x) 0).

  exists (a <-$ c1;
    b <-$ c2;
    ret (a, b)).
  intuition.
  unfold marginal_l.
  dist_inline_first.
  rewrite <- evalDist_right_ident.

  dist_skip.
  dist_inline_first.
  dist_irr_r.
  dist_compute.

  unfold marginal_r.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_irr_r.
  dist_inline_first.
  dist_skip.
  dist_compute.

  repeat simp_in_support.
  simpl in *.
  exfalso.
  eapply getSupport_not_In_evalDist.
  eapply e.
  trivial.
  repeat simp_in_support.
  simpl in *.
  exfalso.
  eapply getSupport_not_In_evalDist.
  rewrite <- H1.
  trivial.
  trivial.

  (* special case: the events have probability 1 *)
  destruct (eq_Rat_dec (evalDist c1 x) 1).
  exists (a <-$ c1;
    b <-$ c2;
    ret (a, b)).

  split.
  intuition.
  unfold marginal_l.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_skip.
  dist_inline_first.
  dist_irr_r.
  dist_compute.

  split.
  intuition.
  unfold marginal_r.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_irr_r.
  dist_inline_first.
  dist_skip.
  dist_compute.

  intros.
  repeat simp_in_support; simpl in *.
  erewrite (evalDist_1) in H5; intuition.
  Focus 2.
  rewrite <- H1.
  trivial.
  simpl in *; intuition.

  erewrite (evalDist_1) in H3; eauto.
  simpl in *; intuition.

  (* general case *)

  exists (a <-$ c1;
    b <-$ if (eqb a x) then (ret y) else (Repeat c2 (fun b => negb (eqb b y)));
      ret (a, b)).

  split.
  intuition.
  unfold marginal_l.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_skip.
  dist_inline_first.
  dist_irr_r.

  case_eq (eqb x0 x); intuition; wftac.
  remember (getSupport c2) as s.
  destruct s.
  specialize (getSupport_length_nz H0); intuition.
  rewrite <- Heqs in H4.
  simpl in *.
  omega.
  destruct (EqDec_dec _ b y); subst.
  destruct s.
  exfalso.
  eapply n0.
  rewrite H1.
  rewrite <- (evalDist_lossless H0).
  rewrite <- Heqs.
  rewrite sumList_cons.
  unfold sumList; simpl.
  eapply ratAdd_0_r.
  eapply (@well_formed_Repeat _ _ _ _ b).
  trivial.
  eapply filter_In; intuition.
  rewrite <- Heqs.
  simpl.
  intuition.
  case_eq (eqb b y); intuition.
  rewrite eqb_leibniz in H4.
  subst.
  specialize (getSupport_NoDup c2); intuition.
  rewrite <- Heqs in H4.
  inversion H4; clear H4; subst.
  simpl in *.
  intuition.

  eapply (@well_formed_Repeat _ _ _ _ b).
  trivial.
  eapply filter_In; intuition.
  rewrite <- Heqs; simpl; intuition.
  case_eq (eqb b y); intuition.
  rewrite eqb_leibniz in H4; subst.
  intuition.

  dist_compute.

  split.
  intuition.
  unfold marginal_r.
  dist_inline_first.
  rewrite evalDist_seq_step.
  rewrite (sumList_filter_partition (eqb x)).
  symmetry.
  eapply eqRat_trans.
  eapply ratAdd_eqRat_compat.
  eapply sumList_body_eq.
  intros.
  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  apply filter_In in H2.
  intuition.
  rewrite eqb_leibniz in H4.
  subst.
  rewrite eqb_refl.
  dist_inline_first.
  unfold snd.
  eapply eqRat_refl.

  eapply sumList_body_eq; intuition.
  apply filter_In in H2.
  intuition.
  case_eq (eqb x a); intuition.
  rewrite eqb_leibniz in H2.
  subst.
  rewrite eqb_refl in H4.
  simpl in *.
  discriminate.

  rewrite evalDist_assoc at 1.
  case_eq (eqb a x); intuition.
  rewrite eqb_leibniz in H5.
  subst.
  rewrite eqb_refl in H2.
  discriminate.

  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  eapply evalDist_seq_eq.
  intuition.
  eapply eqRat_refl.
  intuition.
  dist_simp.
  unfold snd.
  eapply eqRat_refl.
  intuition.
  rewrite (@sumList_exactly_one _ x).
  case_eq (eqb y r2); intuition.
  rewrite eqb_leibniz in H2.
  subst.
  assert (Pr[ret true] == 1).
  dist_compute.
  rewrite H2.
  rewrite ratMult_1_r.
  clear H2.
  rewrite H1.
  symmetry.
  rewrite ratAdd_0_r at 1.
  eapply ratAdd_eqRat_compat; intuition.
  symmetry.
  eapply sumList_0.
  intuition.
  apply filter_In in H2.
  intuition.
  assert (Pr  [x0 <-$ Repeat c2 (fun b : B => negb (eqb b r2)); ret eqb x0 r2 ] ==
    Pr  [x0 <-$ Repeat c2 (fun b : B => negb (eqb b r2)); ret eqb r2 x0 ] ).
  dist_skip; dist_compute.
  rewrite H2.
  clear H2.
  rewrite <- evalDist_event_equiv.
  eapply ratMult_0.
  right.
  eapply evalDist_Repeat_0; intuition.
  subst.
  rewrite eqb_refl in H5.
  simpl in *; discriminate.

  assert (Pr [ret false] == 0).
  dist_compute.
  rewrite H3.
  clear H3.
  rewrite ratMult_0_r.
  rewrite <- ratAdd_0_l.
  rewrite sumList_factor_constant_r.
  rewrite sumList_filter_complement.
  rewrite H1.
  rewrite evalDist_seq_step.

  destruct (in_dec (EqDec_dec _) r2 (getSupport c2)).

  rewrite (@sumList_exactly_one _ r2).
  assert (Pr  [ret eqb r2 r2 ] == 1).
  dist_compute.
  rewrite H3.
  clear H3.
  rewrite ratMult_1_r.
  simpl.
  unfold indicator.
  case_eq (eqb r2 y); intuition.
  rewrite eqb_leibniz in H3; subst.
  rewrite eqb_refl in H2.
  discriminate.
  simpl.
  rewrite ratMult_1_l.
  symmetry.
  rewrite <- ratMult_1_l at 1.
  rewrite <- ratMult_assoc.
  eapply ratMult_eqRat_compat; intuition.
  symmetry.
  eapply eqRat_trans.
  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  eapply ratInverse_eqRat_compat.
  intuition.

  assert ( (filter (fun b : B => negb (eqb b y)) (getSupport c2)) =
     (filter (fun b : B => negb (eqb y b)) (getSupport c2))).
  eapply filter_ext; intuition.
  rewrite eqb_symm.
  trivial.
  rewrite H5 in H4.
  clear H5.
  eapply n0.
  eapply leRat_impl_eqRat.
  eapply evalDist_le_1.
  eapply ratSubtract_0_inv.
  rewrite sumList_filter_complement in H4.
  rewrite H1.
  intuition.
  trivial.

  assert ((filter (fun b : B => negb (eqb b y)) (getSupport c2)) =
    (filter (fun b : B => negb (eqb y b)) (getSupport c2))).
  eapply filter_ext.
  intuition.
  rewrite eqb_symm.
  intuition.
  rewrite H4.
  clear H4.

  eapply sumList_filter_complement.
  trivial.

  rewrite ratMult_comm.
  eapply ratInverse_prod_1.
  intuition.
  eapply n0.
  eapply leRat_impl_eqRat.
  eapply evalDist_le_1.
  eapply ratSubtract_0_inv.
  rewrite H1.
  intuition.

  eapply getSupport_NoDup.
  simpl.
  eapply filter_In; intuition.
  rewrite eqb_symm.
  rewrite H2.
  trivial.

  intuition.
  dist_compute.

  eapply eqRat_trans.
  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  eapply sumList_0.
  intuition.
  simpl in H3.
  apply filter_In in H3.
  intuition.
  case_eq (eqb a r2); intuition.
  rewrite eqb_leibniz in H3; subst.
  intuition.
  dist_compute.

  rewrite ratMult_0_r.
  symmetry.
  eapply getSupport_not_In_evalDist.
  trivial.

  trivial.

  eapply filter_NoDup.
  eapply getSupport_NoDup.

  eapply filter_In; intuition.
  eapply getSupport_In_evalDist; intuition.

  eapply eqb_refl.

  intuition.
  apply filter_In in H2.
  intuition.
  rewrite eqb_leibniz in H5.
  intuition.

  intros.
  repeat simp_in_support;
  simpl in *; intuition.
  rewrite eqb_refl in Heqx2.
  discriminate.

  symmetry in Heqx2.
  rewrite eqb_leibniz in Heqx2.
  trivial.

  simpl in *.
  eapply filter_In in H5.
  intuition.
  rewrite eqb_refl in H4.
  simpl in *.
  discriminate.

  Grab Existential Variables.
  unfold eq_dec; intuition.
  eapply (EqDec_dec _).
  unfold eq_dec; intuition.
  eapply (EqDec_dec _).

Qed.

Theorem le_impl_comp_spec :
  forall (A B : Set){eqda : EqDec A}{eqdb : EqDec B}(c1 : Comp A)(c2 : Comp B),
    well_formed_comp c1 ->
    well_formed_comp c2 ->
    forall x y,
      evalDist c1 x <= evalDist c2 y ->
      comp_spec (fun a b => a = x -> b = y) c1 c2.

  intuition.

  (* special case : y is a certain output of c2 *)
  destruct (eq_Rat_dec (evalDist c2 y) 1).
  exists (a <-$ c1; b <-$ c2; ret (a, b)).
  intuition.
  unfold marginal_l; intuition.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_skip.
  dist_inline_first.
  dist_irr_r.
  dist_compute.

  unfold marginal_l; intuition.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_irr_r.
  dist_inline_first.
  dist_skip.
  dist_compute.

  repeat simp_in_support.
  simpl in *.
  erewrite evalDist_1 in H5; eauto.
  simpl in *.
  intuition.

  (* special case : x is never an output of c1 *)
  destruct (eq_Rat_dec (evalDist c1 x) 0).
  exists (a <-$ c1; b <-$ c2; ret (a, b)).
  intuition.
  unfold marginal_l; intuition.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_skip.
  dist_inline_first.
  dist_irr_r.
  dist_compute.

  unfold marginal_l; intuition.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_irr_r.
  dist_inline_first.
  dist_skip.
  dist_compute.

  repeat simp_in_support.
  simpl in *.
  exfalso.
  eapply getSupport_In_evalDist.
  eapply H4.
  trivial.

  (* special case: the probabilities are equal *)
  destruct (eq_Rat_dec (evalDist c1 x) (evalDist c2 y)).
  eapply comp_spec_consequence.
  eapply eq_impl_comp_spec; intuition.
  eauto.
  intuition.
  subst.
  intuition.

  (* general case *)

  Local Open Scope rat_scope.

  exists (a <-$ c1;
    b <-$ if (eqb a x) then (ret y) else
      mDiff <- ratSubtract (evalDist c2 y) (evalDist c1 x);
      c <-$ Bernoulli (mDiff * (ratInverse (ratSubtract 1 (evalDist c1 x))));
      if c then (ret y) else (Repeat c2 (fun z => negb (eqb z y)));
        ret (a, b)).

  intuition.
  unfold marginal_l.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_skip.
  dist_inline_first.
  dist_irr_r.

  remember (getSupport c2) as s.
  destruct s.
  exfalso.
  specialize (getSupport_length_nz H0); intuition.
  rewrite <- Heqs in H3.
  simpl in *.
  omega.

  destruct (EqDec_dec _ b y).
  subst.
  destruct s.
  exfalso.
  eapply n.
  rewrite <- (evalDist_lossless H0).
  rewrite <- Heqs.
  rewrite sumList_cons.
  unfold sumList. simpl.
  eapply ratAdd_0_r.

  wftac.
  eapply Bernoulli_wf.
  eapply (@well_formed_Repeat _ _ _ _ b).
  intuition.
  eapply filter_In; intuition.
  rewrite <- Heqs.
  simpl.
  intuition.
  case_eq (eqb b y); intuition.
  rewrite eqb_leibniz in H4.
  subst.
  specialize (getSupport_NoDup c2); intuition.
  rewrite <- Heqs in H4.
  inversion H4; clear H4; subst.
  simpl in *.
  intuition.

  wftac.
  eapply Bernoulli_wf.
  eapply (@well_formed_Repeat _ _ _ _ b).
  intuition.
  eapply filter_In; intuition.
  rewrite <- Heqs.
  simpl.
  intuition.
  case_eq (eqb b y); intuition.
  rewrite eqb_leibniz in H4.
  subst.
  intuition.

  dist_compute.

(*  destruct (in_dec (EqDec_dec _) r2 (getSupport c2)). *)

  symmetry.
  unfold marginal_r.
  dist_inline_first.
  rewrite evalDist_seq_step.
  rewrite (sumList_filter_partition (eqb x)).
  eapply eqRat_trans.
  eapply ratAdd_eqRat_compat.
  eapply sumList_body_eq.
  intros.
  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  apply filter_In in H2.
  intuition.
  rewrite eqb_leibniz in H4.
  subst.
  rewrite eqb_refl.
  dist_inline_first.
  unfold snd.
  eapply eqRat_refl.

  eapply sumList_body_eq; intuition.
  apply filter_In in H2.
  intuition.
  case_eq (eqb a x); intuition.
  rewrite eqb_leibniz in H2.
  subst.
  rewrite eqb_refl in H4.
  simpl in *.
  discriminate.

  rewrite evalDist_assoc at 1.
  rewrite evalDist_assoc at 1.
  rewrite (evalDist_seq_step) at 1.
  rewrite (sumList_filter_partition (fun z => z)) at 1.
  rewrite (@sumList_exactly_one _ true) at 1.
  rewrite (@sumList_exactly_one _ false) at 1.
  rewrite Bernoulli_correct at 1.
  rewrite Bernoulli_correct_complement at 1.
  dist_simp.

  assert (Pr  [a0 <-$ ret y; x0 <-$ ret (a, a0); ret eqb (snd x0) r2 ] ==
    Pr  [ret eqb y r2 ]).
  dist_simp.
  simpl.
  intuition.
  rewrite H5 at 1.
  clear H5.

  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  eapply ratAdd_eqRat_compat.
  eapply eqRat_refl.
  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  eapply evalDist_seq_eq.
  intuition.
  eapply eqRat_refl.
  intuition.
  dist_simp.
  unfold snd.
  eapply eqRat_refl.

  eapply ratFraction_le_1.
  eapply ratSubtract_leRat.
  eapply evalDist_le_1.
  intuition.

  eapply ratFraction_le_1.
  eapply ratSubtract_leRat.
  eapply evalDist_le_1.
  intuition.

  eapply filter_NoDup.
  eapply getSupport_NoDup.

  eapply filter_In; intuition.
  eapply getSupport_In_evalDist.
  intuition.
  rewrite Bernoulli_correct_complement in H5.
  apply ratSubtract_0_inv in H5.
  apply ratFraction_ge_1_inv in H5.

  eapply ratSubtract_leRat_same_r in H5.
  eapply n.
  eapply leRat_impl_eqRat.
  eapply evalDist_le_1.
  trivial.

  eapply evalDist_le_1.
  intuition.

  eapply ratFraction_le_1.
  eapply ratSubtract_leRat.
  eapply evalDist_le_1.
  intuition.

  intuition.

  destruct b; intuition.
  apply filter_In in H5.
  intuition.

  eapply filter_NoDup.
  eapply getSupport_NoDup.

  eapply filter_In.
  intuition.
  eapply getSupport_In_evalDist.
  intuition.
  rewrite Bernoulli_correct in H5.
  apply ratMult_0 in H5.
  intuition.
  apply ratSubtract_0_inv in H6.
  eapply n1.
  eapply leRat_impl_eqRat; intuition.
  eapply ratInverse_nz;eauto.
  eapply ratFraction_le_1.
  eapply ratSubtract_leRat.
  eapply evalDist_le_1.
  intuition.

  intuition.
  destruct b; intuition.
  apply filter_In in H5; intuition.
  discriminate.

  intuition.
  intuition.

  rewrite (@sumList_exactly_one _ x).
  rewrite sumList_factor_constant_r.

  assert (sumList (filter (fun a : A => negb (eqb x a)) (getSupport c1))
     (evalDist c1) ==
     ratSubtract 1 (evalDist c1 x)).

  eapply sumList_filter_complement.
  trivial.
  rewrite H2.
  clear H2.
  rewrite ratMult_distrib.

  eapply eqRat_trans.
  eapply ratAdd_eqRat_compat.
  eapply eqRat_refl.
  eapply ratAdd_eqRat_compat.
  rewrite <- ratMult_assoc.
  eapply ratMult_eqRat_compat.
  rewrite ratMult_comm.
  rewrite ratMult_assoc.
  rewrite ratInverse_prod_1.
  eapply ratMult_1_r.

  intuition.
  eapply n.
  eapply leRat_impl_eqRat.
  eapply evalDist_le_1.
  intuition.
  eapply leRat_trans; eauto.
  apply ratSubtract_0_inv.
  trivial.

  eapply eqRat_refl.

  rewrite <- ratMult_assoc.
  eapply ratMult_eqRat_compat.
  rewrite ratMult_comm.
  rewrite ratMult_ratSubtract_distrib_r.
  rewrite ratMult_1_l.
  eapply ratSubtract_eqRat_compat.
  eapply eqRat_refl.
  rewrite ratMult_assoc.
  rewrite ratInverse_prod_1.
  eapply ratMult_1_r.

  intuition.
  eapply n.
  eapply leRat_impl_eqRat.
  eapply evalDist_le_1.
  intuition.
  eapply leRat_trans; eauto.
  apply ratSubtract_0_inv.
  trivial.

  eapply eqRat_refl.

  case_eq (eqb y r2); intuition.
  rewrite eqb_leibniz in H2.
  subst.
  assert (Pr [ret true] == 1).
  dist_compute.
  rewrite H2.
  clear H2.
  repeat rewrite ratMult_1_r.

  assert (Pr  [x0 <-$ Repeat c2 (fun z : B => negb (eqb z r2)); ret eqb x0 r2 ] ==
    Pr  [x0 <-$ Repeat c2 (fun z : B => negb (eqb z r2)); ret eqb r2 x0 ]).

  dist_skip; dist_compute.
  rewrite H2; clear H2.
  rewrite <- evalDist_event_equiv.

  rewrite evalDist_Repeat_0.
  rewrite ratMult_0_r.
  rewrite <- ratAdd_0_r.
  eapply ratSubtract_ratAdd_inverse_2.
  trivial.

  intuition.
  subst.
  rewrite eqb_refl in H3.
  simpl in *.
  discriminate.

  assert (Pr [ret false] == 0).
  dist_compute.
  rewrite H3.
  clear H3.
  repeat rewrite ratMult_0_r.
  repeat rewrite <- ratAdd_0_l.

  rewrite <- ratSubtract_ratAdd_distr.
  rewrite ratSubtract_ratAdd_inverse_2; trivial.
  simpl.

  destruct (in_dec (EqDec_dec _) r2 (getSupport c2)).

  rewrite (@sumList_exactly_one _ r2).
  rewrite eqb_refl.
  destruct (EqDec_dec bool_EqDec true true); intuition.
  unfold indicator.
  case_eq (eqb r2 y); intuition.
  rewrite eqb_leibniz in H3.
  subst.
  rewrite eqb_refl in H2.
  discriminate.
  simpl.
  rewrite ratMult_1_l.
  rewrite ratMult_1_r.

  assert ((sumList (filter (fun z : B => negb (eqb y z)) (getSupport c2))
         (evalDist c2) ==
         (ratSubtract 1 (evalDist c2 y)))).
  eapply sumList_filter_complement.
  trivial.

  eapply eqRat_trans.
  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  eapply ratMult_eqRat_compat.
  eapply ratInverse_eqRat_compat.
  intuition.
  assert ((filter (fun z : B => negb (eqb y z)) (getSupport c2)) =
    (filter (fun z : B => negb (eqb z y)) (getSupport c2))).

  eapply filter_ext; intuition.

  rewrite eqb_symm.
  trivial.

  rewrite H6 in H4.
  clear H6.
  rewrite H4 in H5.
  eapply n.
  eapply leRat_impl_eqRat.
  eapply evalDist_le_1.
  eapply ratSubtract_0_inv.
  trivial.

  assert ((filter (fun z : B => negb (eqb z y)) (getSupport c2)) =
    (filter (fun z : B => negb (eqb y z)) (getSupport c2))).

  eapply filter_ext.
  intuition.
  rewrite eqb_symm.
  trivial.
  rewrite H5.
  clear H5.

  eapply H4.
  eapply eqRat_refl.
  clear H4.
  rewrite <- ratMult_assoc.
  symmetry.
  rewrite <- ratMult_1_l at 1.
  eapply ratMult_eqRat_compat; intuition.
  symmetry.
  rewrite ratMult_comm.
  eapply ratInverse_prod_1.

  intuition.
  eapply n.
  eapply leRat_impl_eqRat.
  eapply evalDist_le_1.
  intuition.
  eapply ratSubtract_0_inv.
  trivial.

  eapply filter_NoDup.
  eapply getSupport_NoDup.

  eapply filter_In; intuition.
  case_eq (eqb r2 y); intuition.
  rewrite eqb_leibniz in H3; subst.
  rewrite eqb_refl in H2.
  discriminate.

  intuition.
  apply filter_In in H3.
  intuition.
  case_eq (eqb b r2); intuition.
  rewrite eqb_leibniz in H3.
  subst.
  intuition.
  destruct (EqDec_dec bool_EqDec false true ).
  discriminate.
  rewrite ratMult_0_r.
  intuition.

  eapply eqRat_trans.
  eapply ratMult_eqRat_compat.
  eapply eqRat_refl.
  eapply sumList_0.
  intuition.
  apply filter_In in H3.
  destruct ( EqDec_dec bool_EqDec (eqb a r2) true); intuition.
  rewrite eqb_leibniz in e; subst.
  intuition.
  eapply ratMult_0_r.
  rewrite ratMult_0_r.
  symmetry.
  eapply getSupport_not_In_evalDist.
  intuition.

  eapply filter_NoDup.
  eapply getSupport_NoDup.

  eapply filter_In; intuition.
  eapply getSupport_In_evalDist; intuition.
  apply eqb_refl.

  intuition.
  apply filter_In in H2.
  intuition.
  dist_compute.

  repeat simp_in_support;
  simpl in *; trivial.

  simpl in *.
  rewrite eqb_refl in Heqx2.
  discriminate.

  Grab Existential Variables.
  unfold eq_dec; intuition.
  eapply (EqDec_dec _).
  unfold eq_dec; intuition.
  eapply (EqDec_dec _).

Qed.

(* soundness theorems *)

Theorem comp_spec_impl_le :
  forall (A B : Set){eqda : EqDec A}{eqdb : EqDec B}(c1 : Comp A)(c2 : Comp B),
    forall x y,
      comp_spec (fun a b => a = x -> b = y) c1 c2 ->
      evalDist c1 x <= evalDist c2 y.

  intuition.
  despec.
  intuition.

  eapply leRat_trans.
  eapply eqRat_impl_leRat.
  eauto.
  eapply leRat_trans.
  Focus 2.
  eapply eqRat_impl_leRat.
  symmetry.
  eauto.
  unfold marginal_l, marginal_r.
  dist_skip.
  dist_compute.
  exfalso.
  eapply n.
  specialize (H2 x1).
  intuition.

  eapply rat0_le_all.
Qed.

Theorem comp_spec_eq_symm:
  forall (A : Set){eqda : EqDec A}(c1 c2 : Comp A),
      comp_spec eq c1 c2 ->
      comp_spec eq c2 c1.

  intuition.
  eapply comp_spec_symm.
  eapply comp_spec_consequence; eauto.

Qed.

Theorem comp_spec_impl_eq :
  forall (A B : Set){eqda : EqDec A}{eqdb : EqDec B}(c1 : Comp A)(c2 : Comp B),
    forall x y,
      comp_spec (fun a b => a = x <-> b = y) c1 c2 ->
      evalDist c1 x == evalDist c2 y.

  intuition.
  eapply leRat_impl_eqRat.

  eapply comp_spec_impl_le.
  eapply comp_spec_consequence.
  eauto.
  intuition.
  subst.
  intuition.

  eapply comp_spec_impl_le.
  eapply comp_spec_symm.
  eapply comp_spec_consequence.
  eauto.
  intuition.
  subst.
  intuition.
Qed.

(* facts about basic language constructs *)
Theorem comp_spec_ret :
  forall (A B : Set){eqda1 : EqDec A}{eqdb1 : EqDec B}(eqda2 : eq_dec A)(eqdb2 : eq_dec B)(P : A -> B -> Prop) a b,
    P a b ->
    @comp_spec _ _ eqda1 eqdb1 P (Ret eqda2 a) (Ret eqdb2 b).

  intuition.
  exists (ret (a, b)).
  intuition.
  unfold marginal_l.
  dist_compute.

  unfold marginal_r.
  dist_compute.

  simp_in_support.
  trivial.

Qed.

Theorem comp_spec_rnd :
  forall n x y,
  comp_spec (fun a b : Vector.t bool n => a = x <-> b = y)
     ({ 0 , 1 }^n) ({ 0 , 1 }^n).

  intuition.
  eapply eq_impl_comp_spec; wftac.
  eapply uniformity.

Qed.

(* facts about equality specifications *)
Theorem eq_impl_comp_spec_eq :
  forall (A : Set){eqd1 eqd2 : EqDec A}(c1 c2 : Comp A),
    (forall x, evalDist c1 x == evalDist c2 x) ->
    @comp_spec _ _ eqd1 eqd2 eq c1 c2.

  intuition.

  exists (x <-$ c1; ret (x, x)).
  intuition.

  unfold marginal_l.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_skip.
  dist_compute.

  unfold marginal_r.
  dist_inline_first.
  rewrite <- H.
  rewrite <- evalDist_right_ident.
  dist_skip.
  dist_compute.

  repeat simp_in_support.
  trivial.
Qed.

Theorem comp_spec_eq_refl :
  forall (A : Set){eqd : EqDec A}(c : Comp A),
    comp_spec eq c c.

  intuition.
  eapply eq_impl_comp_spec_eq; intuition.

Qed.

Theorem comp_spec_eq_impl_eq :
  forall (A : Set){eqd1 eqd2 : EqDec A}(c1 c2 : Comp A),
    @comp_spec _ _ eqd1 eqd2 eq c1 c2 ->
    (forall x, evalDist c1 x == evalDist c2 x).

  intuition.
  despec.
  intuition.
  rewrite H0.
  rewrite H.
  unfold marginal_l, marginal_r.
  dist_skip.
  simpl.
  specialize (H2 x1).
  intuition.
  destruct x1.
  simpl in *.
  subst.

  destruct (EqDec_dec bool_EqDec (@eqb _ eqd1 a0 x) true); intuition.
  rewrite eqb_leibniz in e.
  subst.
  destruct (EqDec_dec bool_EqDec (@eqb _ eqd2 x x) true); intuition.
  exfalso.
  eapply n.
  eapply eqb_refl.

  destruct (EqDec_dec bool_EqDec (@eqb _ eqd2 a0 x) true); intuition.
  rewrite eqb_leibniz in e.
  subst.
  exfalso.
  eapply n.
  eapply eqb_leibniz.
  trivial.

Qed.

Theorem comp_spec_eq_trans_l :
  forall (A B : Set){eqd : EqDec A}{eqd : EqDec B}(c1 c2 : Comp A)(c3 : Comp B) P,
    comp_spec eq c1 c2 ->
    comp_spec P c2 c3 ->
    comp_spec P c1 c3.

  intuition.

  despec.
  intuition.
  exists x.
  intuition.

  eapply comp_spec_eq_impl_eq in H; intuition.
  rewrite H.
  trivial.

Qed.

Theorem comp_spec_eq_trans :
  forall (A : Set){eqd : EqDec A}(c1 c2 c3 : Comp A),
    comp_spec eq c1 c2 ->
    comp_spec eq c2 c3 ->
    comp_spec eq c1 c3.

  intuition.
  eapply comp_spec_eq_trans_l;
  eauto.

Qed.

Theorem comp_spec_eq_trans_r :
  forall (A B : Set){eqd : EqDec A}{eqd : EqDec B}(c1 : Comp A)(c2 c3 : Comp B) P,
    comp_spec P c1 c2 ->
    comp_spec eq c2 c3 ->
    comp_spec P c1 c3.

  intuition.

  destruct H.
  intuition.
  exists x.
  intuition.

  eapply comp_spec_eq_impl_eq in H0; intuition.
  rewrite <- H0.
  trivial.

Qed.

Theorem comp_spec_seq_eq :
  forall (A B C : Set){eqda : EqDec A}{eqdb : EqDec B}{eqdc : EqDec C}(c1 c2 : Comp A)(f1 : A -> Comp B)(f2 : A -> Comp C) P (b : B) (c : C),
    comp_spec eq c1 c2 ->
    (forall a, comp_spec P (f1 a) (f2 a)) ->
    comp_spec P (Bind c1 f1) (Bind c2 f2).

  intuition.
  eapply comp_spec_seq; trivial.
  eauto.
  intuition; subst; intuition.

Qed.

Theorem comp_spec_eq_swap :
  forall (A B C : Set){eqdc : EqDec C}(c1 : Comp A)(c2 : Comp B)(f : A -> B -> Comp C),
    comp_spec eq (x <-$ c1; y <-$ c2; f x y) (y <-$ c2; x <-$ c1; f x y).

  intuition.
  eapply eq_impl_comp_spec_eq.
  intuition.
  eapply evalDist_commute_eq.

Qed.

Theorem comp_spec_right_ident :
  forall (A : Set){eqd : EqDec A}(c1 : Comp A),
    comp_spec eq (x <-$ c1; ret x) c1.

  intuition.
  eapply eq_impl_comp_spec_eq.
  eapply evalDist_right_ident.
Qed.

Theorem comp_spec_left_ident :
  forall (A B : Set){eqda : EqDec A}{eqdb : EqDec B}(c1 : A -> Comp B) a,
    comp_spec eq (x <-$ ret a; c1 x) (c1 a).

  intuition.
  eapply eq_impl_comp_spec_eq.
  intuition.

  specialize (@evalDist_left_ident eqRat); intuition.
Qed.

Theorem comp_spec_assoc :
  forall (A B C : Set){eqd : EqDec C}(c1 : Comp A)(c2 : A -> Comp B)(c3 : B -> Comp C),
    comp_spec eq (x <-$ (y <-$ c1; c2 y); c3 x) (y <-$ c1; x <-$ c2 y; c3 x).

  intuition.
  eapply eq_impl_comp_spec_eq.
  intuition.
  eapply evalDist_assoc; intuition.
Qed.

(* other facts *)
Theorem comp_spec_event_equiv :
  forall (A B : Set){eqda : EqDec A}{eqdb : EqDec B}(c1 : Comp A)(c2 : Comp B)(f1 : A -> bool)(f2 : B -> bool)(P : bool -> bool -> Prop),
    comp_spec (fun a b => P (f1 a) (f2 b)) c1 c2 ->
    comp_spec P (x <-$ c1; ret (f1 x)) (x <-$ c2; ret (f2 x)).

  intuition.

  eapply comp_spec_seq.
  apply true.
  apply true.
  eauto.
  intuition.
  eapply comp_spec_ret.
  trivial.

Qed.

Theorem comp_spec_iso :
  forall (A B : Set){eqda : EqDec A}{eqdb : EqDec B}(c1 : Comp A)(c2 : Comp B)(f : A -> B)(f_inv : B -> A),
    (forall x : B, In x (getSupport c2) -> f (f_inv x) = x) ->
    (forall x : A, In x (getSupport c1) -> f_inv (f x) = x) ->
    (forall x : B, In x (getSupport c2) -> In (f_inv x) (getSupport c1)) ->
    (forall x : A, In x (getSupport c1) -> comp_spec (fun a b => a = (f x) <-> b = x) c2 c1) ->
    comp_spec (fun a b => f a = b) c1 c2.

  intuition.

  intuition.
  remember f_inv as f_inv2.
  exists (x <-$ c1; ret (x, f x)).
  intuition; subst.

  unfold marginal_l.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  dist_skip.
  dist_compute.

  unfold marginal_r.
  dist_inline_first.
  rewrite <- evalDist_right_ident.
  eapply (@distro_iso_eq _ _ _ _ f f_inv).
  intuition.
  intuition.
  intuition.
  intuition.

  rewrite evalDist_event_equiv.
  symmetry.
  rewrite evalDist_event_equiv.
  eapply comp_spec_impl_eq.

  eapply (@comp_spec_consequence _ _ _ _ _ _ eq).
  eapply comp_spec_event_equiv.
  eapply comp_spec_symm.
  eapply comp_spec_consequence.
  eapply H2.
  eauto.
  intuition.
  case_eq (eqb x b); intuition.
  rewrite eqb_leibniz in H5.
  subst.
  intuition.
  case_eq (eqb (f b) a); intuition.
  subst.
  rewrite eqb_refl in H6.
  discriminate.
  case_eq (eqb (f x) a); intuition.
  rewrite eqb_leibniz in H4.
  subst.
  intuition.
  subst.
  rewrite eqb_refl in H5.
  discriminate.
  intuition.
  subst.
  trivial.
  subst.
  trivial.

  intuition.
  dist_compute.

  repeat simp_in_support.
  trivial.

Qed.

Theorem comp_spec_irr_r :
  forall (A B C : Set){eqda : EqDec A}{eqdb : EqDec B}{eqdc : EqDec C}(c1 : Comp A)(c2 : Comp B)(f2 : B -> Comp C) P,
    well_formed_comp c2 ->
    (forall b, In b (getSupport c2) -> comp_spec P c1 (f2 b)) ->
    comp_spec P c1 (Bind c2 f2).

  intuition.
  apply list_choice in H0.
  destruct H0.
  exists (y <-$ c2; x y).
  intuition.
  unfold marginal_l.
  dist_inline_first.
  dist_irr_r.
  specialize (H0 x0); intuition.
  eapply H0.

  unfold marginal_r.
  dist_inline_first.
  dist_skip.
  specialize (H0 x0); intuition.
  apply H2.
  repeat simp_in_support.
  specialize (H0 x0); intuition.
  intuition.
  eapply (EqDec_dec _).
  apply (x <-$ c1; y <-$ c2; z <-$ f2 y; ret (x, z)).
Qed.

Theorem comp_spec_irr_l :
  forall (A B C : Set){eqda : EqDec A}{eqdb : EqDec B}{eqdc : EqDec C}(c1 : Comp A)(c2 : Comp B)(f1 : A -> Comp C) P,
    well_formed_comp c1 ->
    (forall a, In a (getSupport c1) -> comp_spec P (f1 a) c2) ->
    comp_spec P (Bind c1 f1) c2.

  intuition.
  apply list_choice in H0.
  destruct H0.
  exists (y <-$ c1; x y).
  intuition.
  unfold marginal_l.
  dist_inline_first.
  dist_skip.
  specialize (H0 x0); intuition.
  eapply H0.

  unfold marginal_r.
  dist_inline_first.
  dist_irr_r.
  specialize (H0 x0); intuition.
  apply H2.
  repeat simp_in_support.
  specialize (H0 x0); intuition.
  intuition.
  eapply (EqDec_dec _).
  apply (x <-$ c1; y <-$ c2; z <-$ f1 x; ret (z, y)).
Qed.

(* equational rules for computations with oracles. *)
Transparent evalDist.
Transparent getSupport.

Theorem oc_comp_spec_eq :
  forall (A B C : Set)(c : OracleComp A B C), forall (*(eqda : EqDec A)*) (eqdb : EqDec B) (eqdc : EqDec C)(S1 S2 : Set)(o1 : S1 -> A -> Comp (B * S1))(o2 : S2 -> A -> Comp (B * S2)) eqds1 eqds2 s1 s2 (P : S1 -> S2 -> Prop),
    P s1 s2 ->
    (forall a x1 x2, P x1 x2 -> comp_spec (fun y1 y2 => fst y1 = fst y2 /\ P (snd y1) (snd y2)) (o1 x1 a) (o2 x2 a)) ->
    comp_spec (fun a b => fst a = fst b /\ P (snd a) (snd b))
    (c _ eqds1 o1 s1)
    (c _ eqds2 o2 s2).

  (* We have to do this by induction on the computation, since it may make an arbitary number of calls to the oracle. *)

  induction c; intuition; simpl in *.

  eapply (@comp_spec_consequence _ _ (@pair_EqDec B S1 eqdb eqds1) _ (@pair_EqDec B S2 eqdb eqds2)).
  eapply H0; intuition.
  intuition.

  (* OC_Run case *)
  assert (EqDec C).
  eapply EqDec_pair_l;
    eauto.

  eapply (@comp_spec_seq _ _ (fun a b => fst a = fst b /\ fst (snd a) = fst (snd b) /\ P (snd (snd a)) (snd (snd b)) ))
  ; eauto.
  intuition.
  eapply oc_base_exists; eauto; intuition.
  assert (B * S).
  eapply oc_base_exists.
  eauto.
  intuition.
  assert (B' * S1).
  eapply comp_base_exists.
  eauto.
  intuition.
  intuition.

  intuition.
  eapply oc_base_exists; eauto; intuition.
  assert (B * S).
  eapply oc_base_exists.
  eauto.
  intuition.
  assert (B' * S1).
  eapply comp_base_exists.
  eauto.
  intuition.
  intuition.

  eapply comp_spec_consequence.
  eapply (IHc _ _ _ _ _ _ _ _ _ _ (fun p1 p2 => fst p1 = fst p2 /\ P (snd p1) (snd p2))).
  simpl; intuition.

  intuition.
  simpl in *.
  subst.

  eapply comp_spec_seq.
  intuition.
  assert (B * S).
  eapply oc_base_exists.
  eauto.
  intuition.
  assert (B' * S1).
  eapply comp_base_exists.
  eapply o1; intuition.
  intuition.
  intuition.

  intuition.
  assert (B * S).
  eapply oc_base_exists.
  eauto.
  intuition.
  assert (B' * S1).
  eapply comp_base_exists.
  eapply o1; intuition.
  intuition.
  intuition.

  eapply H.
  eauto.
  intuition.
  intuition.
  eapply comp_spec_ret.
  simpl in *; intuition.
  destruct b2; simpl in *.
  destruct p; simpl in *; intuition.
  inversion H7; clear H7; subst.
  intuition.
  destruct b2; simpl in *.
  subst.
  simpl.
  intuition.

  intuition.
  intuition.
  simpl in *.
  subst.
  eapply comp_spec_ret.
  simpl in *.
  destruct b; simpl in *.
  intuition.

  eapply comp_spec_seq_eq; intuition.
  eapply comp_base_exists; intuition.
  eapply comp_base_exists; intuition.
  eapply comp_spec_eq_refl.
  eapply comp_spec_ret; simpl; intuition.

  assert (EqDec C).
  eapply oc_EqDec.
  eapply c.
  intuition.
  intuition.
  assert (B * S1).
  eapply comp_base_exists.
  eauto.
  intuition.
  intuition.

  eapply comp_spec_seq.

  intuition.
  eapply oc_base_exists.
  eapply o.
  eapply oc_base_exists.
  eauto.
  intuition.
  assert (B * S1).
  eapply comp_base_exists.
  eapply o1.
  intuition.
  trivial.
  intuition.
  intuition.
  assert (B * S1).
  eapply comp_base_exists.
  eapply o1; intuition.
  intuition.

  intuition.
  eapply oc_base_exists.
  eapply o.
  eapply oc_base_exists.
  eauto.
  intuition.
  assert (B * S1).
  eapply comp_base_exists.
  eapply o1.
  intuition.
  trivial.
  intuition.
  intuition.
  assert (B * S1).
  eapply comp_base_exists.
  eapply o1; intuition.
  intuition.

  eapply IHc.
  eauto.
  intuition.
  intuition.
  simpl in *.
  subst.
  destruct b0.
  simpl.
  eapply H; intuition.

Qed.

Opaque evalDist.

Theorem oc_comp_spec_eq_until_bad :
  forall (A B C : Set)(c : OracleComp A B C),
    well_formed_oc c ->
    forall (eqdb : EqDec B) (eqdc : EqDec C)(S1 S2 : Set)(o1 : S1 -> A -> Comp (B * S1))(o2 : S2 -> A -> Comp (B * S2)) eqds1 eqds2
    (bad1 : S1 -> bool)(bad2 : S2 -> bool)(inv : S1 -> S2 -> Prop),
    (forall a b, bad1 a = true -> well_formed_comp (o1 a b)) ->
    (forall a b, bad2 a = true -> well_formed_comp (o2 a b)) ->
    (forall x1 x2 a,
      inv x1 x2 ->
      bad1 x1 = bad2 x2 ->
      comp_spec
      (fun y1 y2 => (bad1 (snd y1) = bad2 (snd y2)) /\ (bad1 (snd y1) = false -> (inv (snd y1) (snd y2) /\ fst y1 = fst y2)))
      (o1 x1 a) (o2 x2 a)) ->
    (forall a b c d,
      bad1 c = true ->
      In (a, b) (getSupport (o1 c d)) -> bad1 b = true) ->
    (forall a b c d,
      bad2 c = true ->
      In (a, b) (getSupport (o2 c d)) -> bad2 b = true) ->
    ((forall s1 s2,
      inv s1 s2 ->
      bad1 s1 = bad2 s2 ->
      comp_spec
      (fun y1 y2 => (bad1 (snd y1) = bad2 (snd y2)) /\ (bad1 (snd y1) = false -> inv (snd y1) (snd y2) /\ (fst y1 = fst y2)))
      (c _ eqds1 o1 s1)
      (c _ eqds2 o2 s2))).

  induction 1; intuition; simpl in *.

  (* Query case *)
  eapply (@comp_spec_consequence _ _ (@pair_EqDec B S1 eqdb eqds1) _
    (@pair_EqDec B S2 eqdb eqds2)).
  eapply H1; intuition.
  intuition.

  (* Run case *)
  assert (EqDec C).
  eapply EqDec_pair_l; eauto.

  assert C.
  eapply oc_base_exists; eauto.
  eapply (fun a => fst (oc_base_exists (o s a) (fun a' => fst (comp_base_exists (o1 s1 a'))))).

  eapply (@comp_spec_seq _ _
    (fun x y => (bad1 (snd (snd x)) = bad2 (snd (snd y))) /\ (bad1 (snd (snd x)) = false -> (fst (snd x) = fst (snd y) /\ fst x = fst y /\ inv (snd (snd x)) (snd (snd y)))))).
  intuition.
  intuition.

  eapply comp_spec_consequence.
  eapply (@IHwell_formed_oc _ _ (S * S1) (S * S2) _ _ _ _
     (fun a => bad1 (snd a)) (fun a => bad2 (snd a)) (fun a b => bad1 (snd a) = false -> (inv (snd a) (snd b) /\ fst a = fst b)))%type.

  (* ---- *)

  intuition.
  simpl in *.

  Theorem oc_comp_wf_inv
     : forall (A B C : Set) (c : OracleComp A B C),
       well_formed_oc c ->
       forall (S : Set) (P : S -> Prop)(eqds : EqDec S) (o : S -> A -> Comp (B * S)) (s : S),
       (forall (a : S) (b : A), P a -> well_formed_comp (o a b)) ->
       (forall a b s s', P s -> In (b, s') (getSupport (o s a)) -> P s') ->
       P s ->
       well_formed_comp (c S eqds o s).

    induction 1; intuition; simpl in *.

    eapply H.
    trivial.

    wftac.
    eapply (@IHwell_formed_oc _ (fun x => P (snd x))).
    intuition.
    wftac.
    intuition.
    repeat simp_in_support.
    simpl in *.
    destruct x; simpl in *.
    eapply oc_comp_invariant.
    intuition.
    eapply H3; eauto.
    eauto.
    eauto.
    trivial.

    wftac.

    wftac.
    eapply H1; intuition.
    eapply H2; eauto.
    eapply H3; eauto.
    eapply oc_comp_invariant; intuition.
    eapply H3; eauto.
    eauto.
    eauto.

  Qed.

  wftac.
  eapply oc_comp_wf_inv;
  intuition.
  eapply H2.
  eapply H12.
  intuition.
  eapply H5; intuition.
  eauto.
  eauto.
  trivial.
  intuition.
  wftac.
  simpl in *.
  eapply oc_comp_wf_inv;
  intuition.
  eapply H3.
  eapply H12.
  intuition.
  eapply H6; intuition.
  eauto.
  eauto.
  trivial.

  intuition.
  simpl in *.

  case_eq (bad1 b); intuition.
  eapply comp_spec_irr_l; intuition.
  eapply oc_comp_wf_inv; eauto.
  eapply H2.
  intuition; simpl.
  eapply H5.
  eauto.
  eauto.
  trivial.
  eapply comp_spec_irr_r; intuition.
  eapply oc_comp_wf_inv; eauto.
  eapply H3.
  intuition; simpl.
  eapply H6.
  eauto.
  eauto.
  simpl.
  rewrite <- H12.
  trivial.

  eapply comp_spec_ret; simpl in *; intuition.

  destruct a2; simpl in *.
  assert (bad1 s0 = true).
  eapply (oc_comp_invariant_f).
  intuition.
  eapply H5; eauto.
  eauto.
  eauto.
  destruct b1.
  simpl in *.
  rewrite H16.
  symmetry.
  eapply (@oc_comp_invariant_f).
  intuition.
  eapply H6;
  eauto.
  eauto.
  eauto.

  (*contradiction *)
  destruct a2.
  simpl in *.
  assert (bad1 s0 = true).
  eapply (@oc_comp_invariant_f).
  intuition.
  eapply H5; eauto.
  eauto.
  eauto.
  congruence.
  (* contradiction *)
  destruct a2; simpl in *.
  assert (bad1 s0 = true).
  eapply (@oc_comp_invariant_f).
  intuition.
  eapply H5; eauto.
  eauto.
  eauto.
  congruence.
  (* contradiction *)
  destruct a2; simpl in *.
  assert (bad1 s0 = true).
  eapply (@oc_comp_invariant_f).
  intuition.
  eapply H5; eauto.
  eauto.
  eauto.
  congruence.

  subst.
  (* IH *)
  assert B.
  assert (B * S).
  eapply oc_base_exists.
  eapply o; intuition.
  intuition.
  assert (B' * S1).
  eapply comp_base_exists.
  eapply o1; intuition.
  intuition.
  intuition.

  eapply comp_spec_seq.
  intuition.
  intuition.

  eapply H1;
  intuition.
  intuition.
  intuition.
  intuition.
  eapply comp_spec_ret.
  simpl in *; intuition.
  destruct b3. destruct p. simpl in *.
  pairInv.
  trivial.
  destruct b3.
  simpl in *.
  subst.
  trivial.

  intuition.
  repeat simp_in_support.
  simpl in *.
  destruct x; simpl in *.
  eapply oc_comp_invariant_f; intuition.
  eapply H5; eauto.
  eauto.
  eauto.

  intuition.
  repeat simp_in_support.
  simpl in *.
  destruct x; simpl in *.
  eapply oc_comp_invariant_f; intuition.
  eapply H6; eauto.
  eauto.
  eauto.

  intuition.
  intuition.
  intuition.

  intuition.
  eapply comp_spec_ret.
  simpl in *; intuition.
  subst.
  intuition.

  (* Ret case *)
  assert C.
  eapply comp_base_exists; eauto.
  eapply comp_spec_seq; intuition.
  eapply comp_spec_eq_refl.
  subst.
  eapply comp_spec_ret.
  simpl in *; intuition.

  (* Bind case *)
  assert (A -> B).
  intuition.
  eapply (comp_base_exists (o1 s1 H9)).

  assert C.
  eapply oc_base_exists.
  eauto.
  intuition.

  assert C'.
  eapply oc_base_exists.
  eapply f.
  eapply oc_base_exists; eauto.
  intuition.

  assert (EqDec C).
  eapply oc_EqDec.
  eauto.
  intuition.
  intuition.

  eapply comp_spec_seq;
  intuition.
  eapply (@IHwell_formed_oc _ _ _ _ _ _ _ _ bad1 bad2 inv).
  eauto.
  eauto.
  intuition.
  intuition.
  intuition.
  trivial.
  trivial.
  intuition.
  destruct b0.
  simpl in *.
  case_eq (bad1 b); intuition.

  eapply comp_spec_eq_trans_l.
  eapply comp_spec_eq_symm.
  eapply comp_spec_right_ident.
  eapply comp_spec_eq_trans_r.
  Focus 2.
  eapply comp_spec_right_ident.
  eapply comp_spec_irr_l; intuition.
  eapply oc_comp_wf_inv; eauto.
  intuition.
  eapply H2; intuition.
  eapply H15.
  intuition.
  eapply H5; intuition.
  eauto.
  eauto.
  trivial.
  eapply comp_spec_irr_r; intuition.
  eapply oc_comp_wf_inv; eauto.
  intuition.
  eapply H3; intuition.
  eapply H19.
  intuition.
  eapply H6; intuition.
  eauto.
  eauto.
  simpl.
  rewrite <- H17.
  trivial.

  eapply comp_spec_ret.
  simpl in *; intuition.
  destruct a. destruct b0; simpl in *.
  assert (bad1 s0 = true).
  eapply (oc_comp_invariant_f (f a0) bad1); intuition.
  eapply H5.
  eapply H21.
  eapply H20.
  eapply H16.
  eapply H15.
  rewrite H20.
  symmetry.
  assert (bad2 s = true).
  rewrite <- H17.
  trivial.

  eapply (oc_comp_invariant_f (f c0) bad2); intuition.
  eapply H6.
  eapply H23.
  eapply H22.
  eapply H21.
  eapply H19.

  destruct a; simpl in *.
  destruct b0; simpl in *.
  assert (bad1 s0 = true).
  eapply oc_comp_invariant_f.
  intuition.
  eapply H5.
  eapply H22.
  eapply H21.
  eapply H16.
  eapply H15.
  congruence.

  destruct a; simpl in *.
  destruct b0; simpl in *.
  assert (bad1 s0 = true).
  eapply oc_comp_invariant_f.
  intuition.
  eapply H5.
  eapply H22.
  eapply H21.
  eapply H16.
  eapply H15.
  congruence.

  subst.
  eapply H1;
  intuition.

Qed.

Require Import Setoid.

Add Parametric Relation (A : Set){eqd : EqDec A} : (Comp A) (@comp_spec A A eqd eqd eq)
  reflexivity proved by (@comp_spec_eq_refl A eqd)
  symmetry proved by (@comp_spec_eq_symm A eqd)
  transitivity proved by (@comp_spec_eq_trans A eqd)
  as comp_spec_eq_rel.


(**
- Definition of an effective epimorphism.
- Proof that natural transformations that are pointwise effective epis are
 effective epis.
- Proof that if the target category has pushouts, a natural transformation that is
  an epimorphism is pointwise epimorphic
- Faithul functors reflect epimorphisms
- Definition of a split epimorphism
- Split epis are absolute (ie are preserved by any functor)

Ambroise LAFONT January 2017
*)

Require Import UniMath.Foundations.PartD.
Require Import UniMath.Foundations.Propositions.
Require Import UniMath.Foundations.Sets.

Require Import UniMath.MoreFoundations.Tactics.

Require Import UniMath.CategoryTheory.Core.Categories.
Require Import UniMath.CategoryTheory.Core.Isos.
Require Import UniMath.CategoryTheory.Core.Functors.
Require Import UniMath.CategoryTheory.Core.NaturalTransformations.
Require Import UniMath.CategoryTheory.Epis.
Require Import UniMath.CategoryTheory.FunctorCategory.
Require Import UniMath.CategoryTheory.limits.graphs.pullbacks.
Require Import UniMath.CategoryTheory.limits.graphs.limits.
Require Import UniMath.CategoryTheory.limits.graphs.colimits.
Require Import UniMath.CategoryTheory.limits.graphs.coequalizers.
Require Import UniMath.CategoryTheory.limits.pullbacks.
Require Import UniMath.CategoryTheory.limits.graphs.pushouts.
Require Import UniMath.CategoryTheory.limits.graphs.eqdiag.
Require Import UniMath.CategoryTheory.limits.pushouts.
Require Import UniMath.CategoryTheory.limits.coequalizers.

Local Open Scope cat.

(** Definition of an effective epimorphism.
An effective epimorphism p: A -> B is a morphism which has a kernel pair and which
is the coequalizer of its kernel pair.
*)
Section EffectiveEpi.
  Context {C:category} {A B:C}.
  Variable (f: C ⟦A,B⟧).

  Definition kernel_pair := Pullback  f f.

  Definition isEffective :=
    ∑  g:kernel_pair,
         (isCoequalizer (PullbackPr1 g)
                        (PullbackPr2 g) f (PullbackSqrCommutes g)).
End EffectiveEpi.

Definition EpisAreEffective (C:category) :=
  ∏ (A B:C) (f:C⟦A,B⟧), isEpi f -> isEffective f.

(** Let f be a natural transformation. If f is pointwise effective, then f is effective *)
Section IsEffectivePw.

  Context {C : category} {D : category} .

  Local Notation CD := (functor_category C D).

  Lemma eq_pb_pw {X Y Z:functor C D}
        (a: X ⟹ Z) (b: Y ⟹ Z)
        (c:C)
    : eq_diag
        (pullback_diagram D (a c)  (b c))
        (diagram_pointwise (pullback_diagram CD a b) c).
  Proof.
    intros.
    use tpair.
    use StandardFiniteSets.three_rec_dep; apply idpath.
    use StandardFiniteSets.three_rec_dep;  use StandardFiniteSets.three_rec_dep;
      exact (empty_rect _ ) ||  (exact (λ _, idpath _)).
  Defined.

  Lemma eq_coeq_pw {X Y: functor C D} (a b:X ⟹ Y) (c:C) :
    eq_diag
      (Coequalizer_diagram D (a c) (b c))
      (diagram_pointwise (Coequalizer_diagram CD a b) c).
  Proof.
    intros.
    use tpair.
    use StandardFiniteSets.two_rec_dep; reflexivity.
    use StandardFiniteSets.two_rec_dep;  use StandardFiniteSets.two_rec_dep;
       try exact (empty_rect _ ).
    intros g'.
    destruct g'.
    apply idpath.
    apply idpath.
  Defined.

  Context {X Y :functor C D } {a:X ⟹ Y}.

  Lemma isEffectivePw : (∏ (x:C), isEffective (a x)) -> isEffective (C:=CD) a.
  Proof.
    intros h.
    red.
    transparent assert (f:(kernel_pair (C:=CD) a)).
    { apply equiv_Pullback_2.
      apply LimFunctorCone.
      intro c.
      specialize (h c).
      set (f := pr1 h).
      apply equiv_Pullback_1 in f.
      use (eq_diag_liftlimcone _  _  f).
      apply eq_pb_pw.
    }
    exists f.
    apply equiv_isCoequalizer2.
    apply  pointwise_Colim_is_isColimFunctor.
    intro x.
    set (g:= f).
    assert (hf := (pr2 (h x))); simpl in hf.
    apply equiv_isCoequalizer1 in hf.
    set (eqd := eq_coeq_pw (PullbackPr1 g) (PullbackPr2 g) x).
    set (z:= (eq_diag_iscolimcocone _ eqd hf)).
    set (CC := (make_ColimCocone _ _ _ z)).
    apply (is_z_iso_isColim _ CC).
    rewrite <- (colimArrowUnique CC _ _ (identity _)).
    apply identity_is_z_iso.
    use StandardFiniteSets.two_rec_dep;
    cbn beta;
    rewrite id_right;
    apply idpath.
  Qed.

End IsEffectivePw.

(**  if the target category has pushouts, a natural transformation that is
  an epimorphism is pointwise epimorphic *)
Section PointwiseEpi.

  Context {C : category} {D : category}.

  Local Notation CD := (functor_category C D).

  Lemma eq_po_pw {X Y Z :functor C D} {a: X ⟹ Y } {b: X ⟹ Z} x  :
    eq_diag
      (pushout_diagram D (a x) (b x))
      (diagram_pointwise (pushout_diagram CD a b) x).
  Proof.
    use tpair.
    use StandardFiniteSets.three_rec_dep;  apply idpath.
    use StandardFiniteSets.three_rec_dep;  use StandardFiniteSets.three_rec_dep;
      exact (empty_rect _ )||exact (λ _, idpath _).
  Defined.

  Lemma Pushouts_pw_epi (colimD : graphs.pushouts.Pushouts D) (A B : functor C D)
       (a: A ⟹ B)  (epia:isEpi (C:=CD) a) : ∏ (x:C), isEpi (a x).
  Proof.
    intro  x; simpl.
    apply (epi_to_pushout (C:=CD)) in epia.
    apply pushout_to_epi.
    simpl.
    apply equiv_isPushout1 in epia.
    apply equiv_isPushout2.
    red in epia.
    red.
    eapply isColimFunctor_is_pointwise_Colim in epia; cycle 1.
    {
      intro c.
      eapply eq_diag_liftcolimcocone.
      - apply eq_po_pw.
      - apply colimD.
    }
    apply (eq_diag_iscolimcocone _ (sym_eq_diag _ _ (eq_po_pw x)))in epia; cycle 1.
    set (CC := (make_ColimCocone _ _ _ epia)).
    eapply (is_z_iso_isColim  _ CC).
    rewrite <- (colimArrowUnique CC _ _ (identity _)).
    apply identity_is_z_iso.
    use StandardFiniteSets.three_rec_dep;
    cbn beta;
    rewrite id_right;
    apply idpath.
  Qed.

End PointwiseEpi.

(** faithul functors reflect epimorphisms *)
Lemma faithful_reflects_epis {C D:precategory} (U:functor C D) (hU:faithful U)
      {a b:C} (f:C⟦a,b⟧) : isEpi (#U f) -> isEpi f.
Proof.
  intros hf c u v huv.
  eapply invmaponpathsincl.
  apply hU.
  cbn.
  apply hf.
  rewrite <- functor_comp, <- functor_comp.
  now rewrite huv.
Qed.