From Tealeaves Require Import
  Functors.Option.

#[local] Generalizable Variables A B f g.

Class PartialBijection {A B: Type}
  (f: A → option B)
  (g: B → option A) :=
  { pb_fwd: ∀ (a: A) (b: B),
      f a = Some b → g b = Some a;
    pb_bwd: ∀ (b: B) (a: A),
      g b = Some a → f a = Some b;

  }.

Lemma pb_iff `{PartialBijection A B f g}:
  ∀ (a: A) (b: B),
    f a = Some b ↔ g b = Some a.
Proof.
  intros. split.
  apply pb_fwd.
  apply pb_bwd.
Qed.

Lemma pb_fwd_None `{PartialBijection A B f g}:
  ∀ (a: A),
    f a = None → ∀ b, g b = Some a → False.
Proof.
  introv H1 H2.
  apply pb_bwd in H2.
  congruence.
Qed.

Lemma pb_bwd_None `{PartialBijection A B f g}:
  ∀ (b: B),
    g b = None → ∀ a, f a = Some b → False.
Proof.
  introv H1 H2.
  apply pb_fwd in H2.
  congruence.
Qed.

Lemma pb_iff_None `{PartialBijection A B f g}:
  ∀ (a: A),
    f a = None ↔ ∀ (b: B), ¬ (g b = Some a).
Proof.
  intros. remember (f a) as fa.
  symmetry in Heqfa.
  destruct fa.
  - apply pb_fwd in Heqfa.
    split.
    inversion 1.
    intros contra.
    specialize (contra b Heqfa).
    contradiction.
  - specialize (pb_fwd_None a Heqfa).
    split; auto.
Qed.

Theorem partial_bijection_spec1:
  ∀ (A B: Type) (f: A → option B) (g: B → option A),
    (∀ (a: A), f a = None ∨ map g (f a) = Some (Some a)) →
    (∀ (b: B), g b = None ∨ map f (g b) = Some (Some b)) →
    (∀ (a: A) (b: B), f a = Some b ↔ g b = Some a).
Proof.
  introv HA HB.
  intros a b.
  split.
  - intros Hf.
    specialize (HA a).
    rewrite Hf in HA.
    cbn in HA.
    inversion HA.
    + inversion H.
    + inversion H.
      reflexivity.
  - intros Hg.
    specialize (HB b).
    rewrite Hg in HB.
    cbn in HB.
    inversion HB.
    + inversion H.
    + inversion H.
      reflexivity.
Qed.

Theorem partial_bijection_spec2:
  ∀ (A B: Type) (f: A → option B) (g: B → option A),
    (∀ (a: A) (b: B), f a = Some b ↔ g b = Some a) →
    (∀ (a: A), f a = None ∨ map g (f a) = Some (Some a)).
Proof.
  introv H. intro a.
  specialize (H a).
  destruct (f a) as [b|].
  - right.
    destruct (H b) as [H1 _].
    specialize (H1 ltac:(reflexivity)).
    cbn. fequal.
    assumption.
  - now left.
Qed.

Theorem partial_bijection_spec3:
  ∀ (A B: Type) (f: A → option B) (g: B → option A),
    (∀ (a: A) (b: B), f a = Some b ↔ g b = Some a) →
    (∀ (b: B), g b = None ∨ map f (g b) = Some (Some b)).
Proof.
  introv H. intro b.
  remember (g b) as gb.
  destruct gb as [a|].
  - right.
    cbn.
    fequal.
    apply H.
    symmetry.
    assumption.
  - now left.
Qed.

Theorem partial_bijection_spec:
  ∀ (A B: Type) (f: A → option B) (g: B → option A),
    (∀ (a: A) (b: B), f a = Some b ↔ g b = Some a) ↔
      (∀ (b: B), g b = None ∨ map f (g b) = Some (Some b)) ∧
        (∀ (a: A), f a = None ∨ map g (f a) = Some (Some a)).
Proof.
  intros. split.
  - intros. split.
    now apply partial_bijection_spec2.
    now apply partial_bijection_spec3.
  - intros [H1 H2].
    apply partial_bijection_spec1.
    assumption.
    assumption.
Qed.

Class PartialBijectionSpec
  (A B: Type)
  (P: A → Prop) (Q: B → Prop)
  (f: A → option B)
  (g: B → option A) :=
  { pb_spec_iff: ∀ (a: A) (b: B),
      f a = Some b ↔ g b = Some a;
    pb_P: ∀ (a: A), f a = None ↔ ¬ (P a);
    pb_Q: ∀ (b: B), g b = None ↔ ¬ (Q b);
  }.