Require Import UniMath.Foundations.All.
Require Import UniMath.MoreFoundations.All.
Require Import UniMath.CategoryTheory.Core.Prelude.
Require Import UniMath.CategoryTheory.Limits.Terminal.
Require Import UniMath.CategoryTheory.Monoidal.Categories.
Require Import UniMath.CategoryTheory.Monoidal.Structure.Cartesian.
Require Import UniMath.CategoryTheory.Monoidal.Structure.Symmetric.

Require Import UniMath.CategoryTheory.MarkovCategories.MarkovCategory.
Require Import UniMath.CategoryTheory.MarkovCategories.Determinism.
Require Import UniMath.CategoryTheory.MarkovCategories.AlmostSurely.

Import MonoidalNotations.

Local Open Scope cat.
Local Open Scope moncat.
Local Open Scope markov.

Definition is_conditional {C : markov_category} {a x y : C}
                          (f : a --> x ⊗ y) (k : x ⊗ a --> y) : UU
    := f = copy a
           · (f · proj1 · copy x) #⊗ identity a
           · mon_lassociator _ _ _
           · identity x #⊗ k.

Definition has_conditionals (C : markov_category) : UU :=
   ∏ (a x y : C) (f : a --> x ⊗ y),
   ∑ (k : x ⊗ a --> y), is_conditional f k.

Definition markov_category_with_conditionals : UU :=
   ∑ (C : markov_category), has_conditionals C.

Coercion markov_cat_with_conditional_to_markov
         (C : markov_category_with_conditionals)
  : markov_category := pr1 C.

Section Accessors.
  Context {C : markov_category_with_conditionals}.

  Definition conditional {a x y : C} (f : a --> x ⊗ y) : x ⊗ a --> y
    := pr1 (pr2 C a x y f).

  Proposition conditional_eq {a x y : C} (f : a --> x ⊗ y)
    : f = copy a
           · (f · proj1 · copy x) #⊗ identity a
           · mon_lassociator _ _ _
           · identity x #⊗ (conditional f).
  Proof.
    exact (pr2 (pr2 C a x y f)).
  Qed.
End Accessors.

Lemma copy_tensor_x_I_mor_comp {C : markov_category} {x y : C} (f : I_{C} --> x ⊗ x) (g : x ⊗ I_{C} --> y) :
     ⟨f, identity _⟩ · mon_lassociator _ _ _ · (identity _ #⊗ g)
   = f · (identity _ #⊗ (mon_rinvunitor _ · g)).
Proof.
  unfold pairing.
  rewrite copy_I_mon_rinvunitor.
  rewrite <- tensor_rinvunitor.
  rewrite <- !assoc.
  apply maponpaths.
  rewrite <- mon_rinvunitor_triangle.
  etrans.
  {
    rewrite !assoc'.
    apply maponpaths.
    rewrite assoc.
    rewrite mon_rassociator_lassociator.
    apply id_left.
  }
  rewrite <- tensor_comp_l_id_r.
  apply idpath.
Qed.

Section ConditionalDistributions.
  Context {C : markov_category_with_conditionals}.

  Definition conditional_distribution_1 {x y : C} (p : I_{C} --> x ⊗ y) : x --> y
    := mon_rinvunitor _ · conditional p.

  Definition conditional_distribution_2 {x y : C} (p : I_{C} --> x ⊗ y) : y --> x
    := conditional_distribution_1 (p · sym_mon_braiding _ _ _).

  Notation "p |1" := (conditional_distribution_1 p) : markov.
  Notation "p |2" := (conditional_distribution_2 p) : markov.

  Proposition conditional_distribution_1_eq {x y : C} (p : I_{C} --> x ⊗ y) :
    p = p · proj1 · ⟨identity x, p|1⟩.
  Proof.
    unfold conditional_distribution_1.
    etrans. { apply (conditional_eq p). }
    rewrite <- pairing_eq.
    rewrite copy_tensor_x_I_mor_comp.
    unfold pairing.
    rewrite assoc.
    reflexivity.
  Qed.

  Proposition conditional_distribution_2_eq {x y : C} (p : I_{C} --> x ⊗ y) :
    p = p · proj2 · ⟨p|2, identity y⟩.
  Proof.
    apply cancel_braiding.
    unfold conditional_distribution_2.
    etrans.
    {
      apply conditional_distribution_1_eq.
    }
    refine (!_).
    do 2 refine (assoc' _ _ _ @ _).
    rewrite pairing_sym_mon_braiding.
    rewrite !assoc'.
    apply maponpaths.
    refine (_ @ assoc' _ _ _).
    apply maponpaths_2.
    rewrite sym_mon_braiding_proj1.
    reflexivity.
  Qed.

  Proposition conditional_distribution_1_ase_unique {x y : C} (p : I_{C} --> x ⊗ y) (f : x --> y) :
    p = p · proj1 · ⟨identity x, f⟩ -> f =_{p · proj1} p|1.
  Proof.
    intros e.
    apply make_equal_almost_surely_r.
    rewrite <- conditional_distribution_1_eq.
    rewrite <- e.
    reflexivity.
  Qed.

  Proposition conditional_distribution_2_ase_unique {x y : C} (p : I_{C} --> x ⊗ y) (f : y --> x) :
    p = p · proj2 · ⟨f, identity y⟩ -> f =_{p · proj2} p|2.
  Proof.
    intros e.
    apply make_equal_almost_surely_r.
    apply cancel_braiding.
    rewrite !assoc'.
    rewrite !pairing_sym_mon_braiding.
    rewrite !assoc.
    rewrite <- conditional_distribution_2_eq.
    rewrite <- e.
    reflexivity.
  Qed.

  Proposition conditional_distribution_1_braiding {x y : C} (p : I_{C} --> x ⊗ y) :
    (p · sym_mon_braiding _ _ _)|1 =_{p · proj2} p|2.
  Proof.
    apply conditional_distribution_2_ase_unique.
    rewrite <- sym_mon_braiding_proj1.
    rewrite assoc.
    apply cancel_braiding.
    rewrite assoc'.
    rewrite pairing_sym_mon_braiding.
    rewrite <- conditional_distribution_1_eq.
    reflexivity.
  Qed.

  Proposition conditional_distribution_2_braiding {x y : C} (p : I_{C} --> x ⊗ y) :
    (p · sym_mon_braiding _ _ _)|2 =_{p · proj1} p|1.
  Proof.
    apply conditional_distribution_1_ase_unique.
    rewrite <- sym_mon_braiding_proj2.
    rewrite assoc.
    apply cancel_braiding.
    rewrite assoc'.
    rewrite pairing_sym_mon_braiding.
    rewrite <- conditional_distribution_2_eq.
    reflexivity.
  Qed.

  Proposition conditional_distribution_1_ase_unique'
      {x y : C} {p : I_{C} --> x ⊗ y} (p1 : I_{C} --> x) (f : x --> y) (e : p · proj1 = p1)
    : p = p1 · ⟨identity x, f⟩ -> f =_{p1} p|1.
  Proof.
    rewrite <- e.
    apply conditional_distribution_1_ase_unique.
  Qed.

  Proposition conditional_distribution_2_ase_unique'
    {x y : C} {p : I_{C} --> x ⊗ y} (p2 : I_{C} --> y) (f : y --> x) (e : p · proj2 = p2)
  : p = p2 · ⟨f, identity y⟩ -> f =_{p2} p|2.
  Proof.
    rewrite <- e.
    apply conditional_distribution_2_ase_unique.
  Qed.

End ConditionalDistributions.

Notation "p |1" := (conditional_distribution_1 p) : markov.
Notation "p |2" := (conditional_distribution_2 p) : markov.

Section DefBayesianInverse.
  Context {C : markov_category}.

  Definition is_bayesian_inverse {a x y : C} (p : a --> x) (f : x --> y) (fi : y --> x) : UU
    := p · ⟨identity x, f⟩ = (p · f) · ⟨fi, identity y⟩.


  Proposition make_bayesian_inverse_l {a x y : C} (p : a --> x) (f : x --> y) (fi : y --> x)
    : p · ⟨identity x, f⟩ = (p · f) · ⟨fi, identity y⟩ -> is_bayesian_inverse p f fi.
  Proof.
    intros bi. exact bi.
  Qed.

  Proposition make_bayesian_inverse_r {a x y : C} (p : a --> x) (f : x --> y) (fi : y --> x)
    : p · ⟨f, identity x⟩ = (p · f) · ⟨identity y, fi⟩ -> is_bayesian_inverse p f fi.
  Proof.
    intros bi.
    apply pairing_flip.
    exact bi.
  Qed.

  Proposition is_bayesian_inverse_l {a x y : C} (p : a --> x) (f : x --> y) (fi : y --> x)
    : is_bayesian_inverse p f fi -> p · ⟨identity x, f⟩ = (p · f) · ⟨fi, identity y⟩.
  Proof.
    intros bi. exact bi.
  Qed.

  Proposition is_bayesian_inverse_r {a x y : C} (p : a --> x) (f : x --> y) (fi : y --> x)
  : is_bayesian_inverse p f fi -> p · ⟨f, identity x⟩ = (p · f) · ⟨identity y, fi⟩.
  Proof.
    intros bi.
    apply pairing_flip.
    exact bi.
  Qed.

  Proposition bayesian_inverse_ase_unique
      {a x y : C} (p : a --> x) (f : x --> y) (g1 g2 : y --> x)
      (b1 : is_bayesian_inverse p f g1) (b2 : is_bayesian_inverse p f g2)
    : g1 =_{p · f} g2.
  Proof.
    unfold is_bayesian_inverse in *.
    apply make_equal_almost_surely_r.
    apply pairing_flip.
    rewrite <- b1, <- b2.
    reflexivity.
  Qed.

  Lemma bayesian_inverse_ase_unique'
    {a x y : C} (p : a --> x) (q : a --> y) (f : x --> y) (g1 g2 : y --> x)
    (e : p · f = q) (b1 : is_bayesian_inverse p f g1) (b2 : is_bayesian_inverse p f g2)
  : g1 =_{q} g2.
  Proof.
    rewrite <- e.
    apply bayesian_inverse_ase_unique; assumption.
  Qed.

  Proposition is_bayesian_inverse_ase_cong
    {a x y : C} (p : a --> x) (f1 f2 : x --> y) (g1 g2 : y --> x)
    (fase : f1 =_{ p } f2) (gase : g1 =_{ p · f1 } g2)
    (b1 : is_bayesian_inverse p f1 g1)
    : is_bayesian_inverse p f2 g2.
  Proof.
    unfold is_bayesian_inverse in *.
    rewrite <- (equal_almost_surely_r _ _ _ fase).
    rewrite (equal_almost_surely_composition _ _ _ fase) in gase.
    rewrite <- (equal_almost_surely_l _ _ _ gase).
    rewrite <- (equal_almost_surely_composition _ _ _ fase).
    rewrite b1.
    reflexivity.
  Qed.

End DefBayesianInverse.

Section BayesianInverseLaws.
  Context {C : markov_category}.

  Proposition is_bayesian_inverse_state_preservation
      {a x y : C} (p : a --> x) (f : x --> y) (g : y --> x)
    : is_bayesian_inverse p f g -> p · f · g = p.
  Proof.
    intros e.
    rewrite <- id_right.
    rewrite <- (pairing_proj1 (identity x) f).
    rewrite <- (pairing_proj1 g (identity y)).
    rewrite !assoc.
    rewrite e.
    reflexivity.
  Qed.

  Proposition bayesian_inverse_identity {a x : C} (p : a --> x) :
    is_bayesian_inverse p (identity x) (identity x).
  Proof.
    unfold is_bayesian_inverse.
    rewrite id_right.
    reflexivity.
  Qed.

  Proposition bayesian_inverse_idempotent
    {a x y : C} (p : a --> x) (f : x --> y) (g : y --> x)
    : is_bayesian_inverse p f g -> is_bayesian_inverse (p · f) g f.
  Proof.
    intros e.
    unfold is_bayesian_inverse in *.
    assert(e2 : p · f · g = p).
    {
      rewrite <- id_right.
      rewrite <- (pairing_proj1 (identity x) f).
      rewrite <- (pairing_proj1 g (identity y)).
      rewrite !assoc.
      rewrite e.
      reflexivity.
    }
    rewrite e2.
    apply pairing_flip.
    rewrite <- e.
    reflexivity.
  Qed.

  Proposition bayesian_inverse_comp
    {a x y z : C} (p : a --> x) (f : x --> y) (g : y --> z) (fi : y --> x) (gi : z --> y)
    (bf : is_bayesian_inverse p f fi) (bg : is_bayesian_inverse (p · f) g gi)
    : is_bayesian_inverse p (f · g) (gi · fi).
  Proof.
    unfold is_bayesian_inverse in *.
    etrans. {
      rewrite <- pairing_tensor_r.
      rewrite assoc.
      rewrite bf.
      rewrite <- assoc.
      rewrite pairing_tensor_r.
      rewrite id_left.
      rewrite pairing_split_l.
      rewrite assoc.
      rewrite bg.
      rewrite <- assoc.
      rewrite pairing_tensor_l.
      reflexivity.
    }
    rewrite assoc.
    reflexivity.
  Qed.

End BayesianInverseLaws.

Section ConstructionBayesianInverse.
  Context {C : markov_category_with_conditionals}.
  Context {x : C}.
  Context (p : I_{C} --> x).

  Definition bayesian_inverse {y : C} (f : x --> y) : y --> x
    := mon_rinvunitor y · conditional (p · ⟨f, identity x⟩).

  Definition bayesian_inverse_eq {y : C} {f : x --> y} :
    is_bayesian_inverse p f (bayesian_inverse f).
  Proof.
    unfold is_bayesian_inverse, bayesian_inverse.
    apply pairing_flip.
    pose(phi := p · ⟨ f, identity x ⟩).
    pose(K := conditional_eq phi).

    assert(Aux1 : phi · proj1 = p · f).
    { unfold phi.
      rewrite <- assoc.
      rewrite pairing_proj1.
      reflexivity. }
    rewrite Aux1 in K.
    clear Aux1.
    rewrite <- pairing_eq in K.
    rewrite copy_tensor_x_I_mor_comp in K.
    unfold phi in K.
    rewrite <- assoc in K.
    etrans.
    { rewrite K. reflexivity. }
    reflexivity.
  Qed.

  Proposition bayesian_inverse_eq_l {y : C} (f : x --> y) :
    p · f · ⟨bayesian_inverse f, identity _⟩ = p · ⟨identity _, f⟩.
  Proof.
    symmetry.
    assert(e : is_bayesian_inverse p f (bayesian_inverse f)).
      { apply bayesian_inverse_eq. }
    unfold is_bayesian_inverse in e.
    rewrite <- e.
    reflexivity.
  Qed.

  Proposition bayesian_inverse_eq_r {y : C} (f : x --> y) :
    p · f · ⟨identity _, bayesian_inverse f⟩ = p · ⟨f, identity _⟩.
  Proof.
    apply cancel_braiding.
    rewrite !assoc', !pairing_sym_mon_braiding, assoc.
    apply bayesian_inverse_eq_l.
  Qed.

  Proposition bayesian_inverse_state_preservation
      {y : C} (f : x --> y)
    : p · f · bayesian_inverse f = p.
  Proof.
    apply is_bayesian_inverse_state_preservation.
    apply bayesian_inverse_eq.
  Qed.

End ConstructionBayesianInverse.

#[global] Opaque bayesian_inverse.