From Tealeaves Require Export
  Misc.Prop
  Classes.Functor
  Classes.Monoid.

From Tealeaves Require Import
  Classes.Categorical.ApplicativeCommutativeIdempotent
  Classes.Kleisli.Monad
  Tactics.Debug.

Import Kleisli.Monad.Notations.

#[local] Generalizable Variables A B.

#[local] Notation "( -> B )" := (fun A ⇒ A → B) (at level 50).
#[local] Notation "'subset'" := ((-> Prop)).

Definition subset_one {A}: A → subset A := eq.

Definition subset_empty {A}: subset A :=
  const False.

Definition subset_add {A}: subset A → subset A → subset A :=
  fun x y p ⇒ x p ∨ y p.

Module Notations.
  Notation "∅" := subset_empty: tealeaves_scope.
  Notation "{{ x }}" := (subset_one x): tealeaves_scope.
  Infix "∪" :=
    subset_add (at level 61, left associativity): tealeaves_scope.
  Notation "( -> B )" :=
    (fun A ⇒ A → B) (at level 50): tealeaves_scope.
  Notation "'subset'" := ((-> Prop)): tealeaves_scope.
End Notations.

Import Notations.

Tactic Notation "simpl_subset" :=
  (autorewrite with tea_set).
Tactic Notation "simpl_subset" "in" hyp(H) :=
  (autorewrite with tea_set in H).
Tactic Notation "simpl_subset" "in" "*" :=
  (autorewrite with tea_set in *).

Ltac unfold_subset :=
  unfold subset_empty; unfold subset_add; unfold const.

Ltac solve_basic_subset :=
  unfold transparent tcs; unfold_subset; unfold compose; try setext;
  first [tauto | firstorder (subst; (solve auto + eauto)) ].

Section subset_monoid.

  Context
    {A: Type}.

  Implicit Types (s t: subset A) (a: A).

  Definition subset_add_nil_l: ∀ s, s ∪ ∅ = s :=
    ltac:(solve_basic_subset).
  Definition subset_add_nil_r: ∀ s, ∅ ∪ s = s :=
    ltac:(solve_basic_subset).
  Definition subset_add_assoc: ∀ s t u, s ∪ t ∪ u = s ∪ (t ∪ u) :=
    ltac:(solve_basic_subset).
  Definition subset_add_comm: ∀ s t, s ∪ t = t ∪ s :=
    ltac:(solve_basic_subset).

  Definition subset_in_empty: ∀ a, ∅ a = False
    := ltac:(solve_basic_subset).
  Definition subset_in_add: ∀ s t a, (s ∪ t) a = (s a ∨ t a)
    := ltac:(solve_basic_subset).

  Definition subset_in_one: ∀ a a', {{a}} a' = (a = a')
    := ltac:(solve_basic_subset).

End subset_monoid.

#[export] Hint Rewrite @subset_add_nil_l @subset_add_nil_r
     @subset_add_assoc @subset_in_empty @subset_in_add @subset_in_one: tea_set.

#[export] Instance Monoid_op_subset {A}:
  Monoid_op (subset A) := @subset_add A.

#[export] Instance Monoid_unit_subset {A}:
  Monoid_unit (subset A) := subset_empty.

#[export, program] Instance Monoid_subset {A}:
  @Monoid (subset A) (@Monoid_op_subset A) (@Monoid_unit_subset A).

Solve Obligations with
  (intros; unfold transparent tcs; solve_basic_subset).

#[export] Instance CommutativeMonoidOp_subset: ∀ (A: Type),
    CommutativeMonoidOp (M := subset A) Monoid_op_subset.
Proof.
  intros; constructor; solve_basic_subset.
Qed.

Lemma monoid_subset_unit_rw {A}:
  monoid_unit (subset A) (Monoid_unit := Monoid_unit_subset) = ∅.
Proof.
  reflexivity.
Qed.

Lemma monoid_subset_rw:
  ∀ {A} (l1 l2: subset A),
    monoid_op (Monoid_op := Monoid_op_subset) l1 l2 = l1 ∪ l2.
Proof.
  reflexivity.
Qed.

Ltac simplify_monoid_subset :=
  ltac_trace "simplify_monoid_subset";
  match goal with
  | |- context[monoid_op (Monoid_op := Monoid_op_subset) ?S1 ?S2] ⇒
      rewrite monoid_subset_rw
  | |- context[monoid_unit _ (Monoid_unit := Monoid_unit_subset)] ⇒
      rewrite monoid_subset_unit_rw
  end.

#[export] Instance Monmor_el {A: Type} (a: A):
  @Monoid_Morphism (subset A) Prop
    (@Monoid_op_subset A) (@Monoid_unit_subset A)
    (Monoid_op_or) (Monoid_unit_false)
    (evalAt a).
Proof.
  constructor.
  - typeclasses eauto.
  - typeclasses eauto.
  - reflexivity.
  - reflexivity.
Qed.

#[export] Instance Map_subset: Map subset :=
  fun A B f s b ⇒ ∃ (a: A), s a ∧ f a = b.

Definition map_set_nil `{f: A → B}:
  map f ∅ = ∅ := ltac:(solve_basic_subset).

Lemma map_set_one `{f: A → B} {a: A}:
  map f {{ a }} = {{ f a }}.
Proof.
  ext b. propext.
  - intros [a' [Hin Heq]].
    cbv in Hin; subst.
    solve_basic_subset.
  - cbv. intro; subst.
    eauto.
Qed.

Definition map_set_add `{f: A → B} {x y}:
  map f (x ∪ y) = map f x ∪ map f y
  := ltac:(solve_basic_subset).

#[export] Hint Rewrite
  @map_set_nil @map_set_one @map_set_add
: tea_set.

Lemma map_id_subset: ∀ (A: Type), map id = id (A := subset A).
Proof.
  intros. ext s a.
  cbv. propext.
  - intros [a' [Hin Heq]]. now subst.
  - intros H. eauto.
Qed.

Lemma map_map_subset: ∀ (A B C: Type) (f: A → B) (g: B → C),
    map g ∘ map f = map (F := subset) (g ∘ f).
Proof.
  intros. ext s c.
  unfold compose. cbv.
  propext.
  - intros [b [[a [Hina Heqa]] Heq]].
    ∃ a. rewrite Heqa. eauto.
  - intros [a [Hin Heq]]. eauto.
Qed.

#[export] Instance Functor_subset: Functor subset :=
  {| fun_map_id := map_id_subset;
     fun_map_map := map_map_subset;
  |}.

#[export] Instance Monoid_Morphism_subset_map:
  ∀ (A B: Type) (f: A → B),
    Monoid_Morphism (subset A) (subset B) (map f).
Proof.
  intros.
  constructor.
  - typeclasses eauto.
  - typeclasses eauto.
  - ext b. apply propositional_extensionality.
    firstorder.
  - intros. ext b. apply propositional_extensionality.
    firstorder.
Qed.

#[export] Instance Return_subset: Return subset := fun A a b ⇒ a = b.

#[local] Notation "{{ x }}" := (@ret subset _ _ x).

#[local] Instance Join_subset: Monad.Join subset := fun A P a ⇒ ∃ (Q: subset A), P Q ∧ Q a.

#[export] Instance Natural_Return_subset: Natural (@ret subset _).
Proof.
  constructor.
  - typeclasses eauto.
  - typeclasses eauto.
  - intros. ext a. ext b.
    unfold_ops @Map_I @Map_subset @Return_subset.
    unfold compose.
    propext.
    firstorder (now subst).
    firstorder (now subst).
Qed.

Lemma natural_join_subset:
  ∀ (A B : Type) (f : A → B), map f ∘ Monad.join = Monad.join ∘ map f.
Proof.
  intros.
  unfold compose.
  unfold_ops @Map_compose @Map_subset @Join_subset.
  ext P b. propext.
  + intros [a [[Q [PQ Qa]] Heq]].
    subst.
    ∃ (map (F := subset) f Q).
    unfold_ops @Map_subset.
    split.
    × ∃ Q. split.
      { assumption. }
      { setext; firstorder. }
    × firstorder.
  + intros [Qb [[Qa [PQa rest]] Qbb]].
    subst.
    firstorder.
Qed.

#[export] Instance Natural_Join_subset: Natural (@Monad.join subset _).
Proof.
  constructor.
  - typeclasses eauto.
  - typeclasses eauto.
  - apply natural_join_subset.
Qed.

Lemma join_ret_subset:
  ∀ A: Type, Monad.join (T := subset) ∘ (ret (T := subset)) = (@id (subset A)).
Proof.
  intros. ext P. unfold id.
  unfold compose.
  unfold_ops @Join_subset.
  ext a.
  propext.
  - simpl_subset.
    intros [Q [Qin Qa]].
    simpl_subset in Qin.
    now subst.
  - intros. ∃ P. now simpl_subset.
Qed.

Corollary join_one_subset: ∀ (A: Type) (P: A → Prop),
    Monad.join {{P}} = P.
Proof.
  intros.
  change ((Monad.join ∘ ret (T := subset)) P = P).
  rewrite join_ret_subset.
  reflexivity.
Qed.

Lemma join_map_ret_subset:
  ∀ A : Type, Monad.join (T := subset) (A := A) ∘
                map (ret (A := A) (T := subset)) = id.
Proof.
  intros.
  unfold id, compose.
  ext P a.
  unfold_ops @Join_subset @Map_compose @Map_subset.
  propext.
  - intros [P' [[a' [Pa' rest]] Ha']].
    subst.
    simpl_subset in Ha'. now subst.
  - intros. ∃ {{a}}.
    simpl_subset.
    split; eauto.
Qed.

Lemma join_join_subset:
  ∀ A : Type,
    Monad.join (T := subset) (A := A) ∘
      Monad.join (T := subset) (A := subset A)
    = Monad.join (T := subset) (A := A) ∘ map (Monad.join (T := subset)).
Proof.
  intros.
  unfold id, compose.
  ext P a.
  propext.
  - intros [Q [[R [PR RQ]] Qa]].
    unfold_ops @Map_subset.
    unfold Monad.join at 1.
    unfold Join_subset at 1.
    ∃ (Monad.join (T := subset) R); split; auto.
    ∃ R; auto. ∃ Q; auto.
  - intros [R [[Q [PQ Rspec]] Ra]].
    subst.
    destruct Ra as [X [QX Xa]].
    ∃ X; split; auto.
    ∃ Q; auto.
Qed.

#[local] Instance Categorical_Monad_subset: Categorical.Monad.Monad subset.
Proof.
  constructor.
  - typeclasses eauto.
  - typeclasses eauto.
  - typeclasses eauto.
  - apply join_ret_subset.
  - apply join_map_ret_subset.
  - apply join_join_subset.
Qed.

#[export] Instance Bind_subset: Bind subset subset := fun A B f s_a ⇒
(fun b ⇒ ∃ (a: A), s_a a ∧ f a b).

#[export] Instance Compat_Map_Bind_subset:
  `{Compat_Map_Bind (Map_U := Map_subset) subset subset} :=
  ltac:(reflexivity).

Definition set_in_ret: ∀ A (a b: A),
    (ret a) b = (a = b) := ltac:(reflexivity).

#[export] Hint Rewrite @set_in_ret: tea_set.

Lemma bind_set_nil `{f: A → subset B}:
  bind f ∅ = ∅.
Proof.
  solve_basic_subset.
Qed.

Lemma bind_set_one `{f: A → subset B} {a: A}:
  bind f {{ a }} = f a.
Proof.
  solve_basic_subset.
Qed.

Lemma bind_set_add `{f: A → subset B} {x y}:
  bind f (x ∪ y) = bind f x ∪ bind f y.
Proof.
  solve_basic_subset.
Qed.

#[export] Hint Rewrite
  @bind_set_nil @bind_set_one @bind_set_add
: tea_set.

Lemma set_bind0: ∀ (A B: Type) (f: A → subset B),
    bind f ∘ ret = f.
Proof.
  intros. ext a. unfold compose.
  now autorewrite with tea_set.
Qed.

Lemma set_bind1: ∀ A: Type, bind ret = @id (subset A).
  intros. cbv. ext s a. propext.
  - intros [a' [h1 h2]]. now subst.
  - intro. eexists a. intuition.
Qed.

Lemma set_bind2:
  ∀ (A B C: Type) (g: B → subset C) (f: A → subset B),
    bind g ∘ bind f = bind (g ⋆ f).
Proof.
  intros. ext a. unfold compose.
  cbv. ext c. propext.
  - intros [b [[a' [ha1 ha2]] hb2]].
    ∃ a'; split; [assumption | ∃ b; split; assumption].
  - intros [a' [ha1 [b [hb1 hb2]]]].
    ∃ b; split; [∃ a'; split; assumption | assumption].
Qed.

#[export] Instance RightPreModule_subset: RightPreModule subset subset :=
  {| kmod_bind1 := set_bind1;
     kmod_bind2 := set_bind2;
  |}.

#[export] Instance Monad_subset: Monad subset :=
  {| kmon_bind0 := set_bind0;
  |}.

#[export] Instance RightModule_subset: RightModule subset subset :=
  {| kmod_monad := _;
  |}.

#[export] Instance Monmor_bind {A B f}:
  Monoid_Morphism (subset A) (subset B) (bind f) :=
  {| monmor_unit := @bind_set_nil A B f;
     monmor_op := @bind_set_add A B f;
  |}.

Theorem set_ret_injective: ∀ (A: Type) (a b: A),
    {{ a }} = {{ b }} → a = b.
Proof.
  intros. assert (lemma: ∀ x, {{ a }} x = {{ b }} x).
  intros. now rewrite H. specialize (lemma a).
  cbv in lemma. symmetry. now rewrite <- lemma.
Qed.

Section subset_applicative_instance.

  Import Applicative.Notations.

  #[export] Instance Pure_subset: Pure subset := @eq.

  #[export] Instance Mult_subset: Mult subset :=
    fun A B '(sa, sb) '(a, b) ⇒ sa a ∧ sb b.

  #[export] Instance Applicative_subset: Applicative subset.
  Proof.
    constructor.
    - apply Functor_subset.
    - intros. change (pure ?x) with (ret x).
      now simpl_subset.
    - intros. unfold_ops @Mult_subset.
      ext [c d]. cbv. propext.
      + intros [[a [Ha1 Ha2]] [b [Hb1 Hb2]]].
        ∃ (a, b). subst. intuition.
      + intros [[a b] [Hab1 Hab2]].
        inversion Hab2. subst. intuition eauto.
    - intros. ext [a [b c]]. cbv.
      propext.
      + intros [[[a' b'] c'] [Ht1 Ht2]].
        inversion Ht2; subst. tauto.
      + intros. ∃ (a, b, c). tauto.
    - intros. ext a. cbv.
      propext.
      + intros [[tt' a'] [H1 H2]]. now subst.
      + intros Xa. ∃ (tt, a). tauto.
    - intros. ext a. cbv.
      propext.
      + intros [[a' tt'] [H1 H2]]. now subst.
      + intros Xa. ∃ (a, tt). tauto.
    - intros. ext [a' b']. cbv. propext.
      + intuition. now subst.
      + intros X; inversion X; now subst.
  Qed.

  Lemma subset_ap_spec:
    ∀ (A B: Type) (sf: subset (A → B)) (sa: subset A) (b: B),
      (sf <⋆> sa) b = ∃ (f: A → B) (a: A), sa a ∧ sf f ∧ f a = b.
  Proof.
    intros.
    unfold ap.
    unfold_ops @Map_subset.
    unfold_ops @Mult_subset.
    propext.
    - intros [[f a] [[hyp1 hyp2] hyp3]].
      ∃ f. ∃ a. auto.
    - intros [f [a [hyp1 [hyp2 hyp3]]]].
      subst. ∃ (f, a). tauto.
  Qed.

  Lemma fn_nequal_counterexample:
    ∀ (A B: Type) (f g: A → B),
      (∃ (a: A), (f a ≠ g a)) → f ≠ g.
  Proof.
    intros.
    intro contra.
    destruct H as [a Hneq].
    rewrite contra in Hneq.
    contradiction.
  Qed.

  Corollary subset_nequal_counterexample:
    ∀ (A: Type) (a: A) (s1 s2: subset A),
      s1 a → ¬ s2 a →
      s1 ≠ s2.
  Proof.
    intros.
    apply fn_nequal_counterexample.
    ∃ a. intro contra.
    rewrite contra in H.
    contradiction.
  Qed.

  Example not_subset_idem:
    ¬ (∀ (A: Type) (s: subset A),
          s ⊗ s = map (fun a: A ⇒ (a, a)) s).
  Proof.
    intro contra.
    specialize (contra nat ({{1}} ∪ {{2}})).
    generalize contra.
    pose (ctx := subset_nequal_counterexample (nat × nat) (1, 2)).
    unfold not in ctx.
    unfold map, Map_subset, mult,Mult_subset.
    apply ctx.
    - simpl_subset. tauto.
    - intros [a [_ false]].
      now inversion false.
  Qed.

  Lemma subset_commutative:
    ∀ (A: Type) (sA: subset A), Center subset%tea A sA.
  Proof.
    intros. constructor; intros;
     unfold_ops @Mult_subset;
       unfold_ops @Map_subset.
     - rename x into sB.
       ext [b a].
       propext.
       + intros [Ha Hb].
         ∃ (a, b). tauto.
       + intros [[a' b'] [H1 H2]].
         inversion H2. subst. tauto.
     - rename x into sB.
       ext [a b].
       propext.
       + intros [Ha Hb].
         ∃ (b, a). tauto.
       + intros [[b' a'] [H1 H2]].
         inversion H2. subst. tauto.
  Qed.

  #[export] Instance ApplicativeCommutative_subset: ApplicativeCommutative subset.
  Proof.
    constructor.
    typeclasses eauto.
    apply subset_commutative.
  Qed.

End subset_applicative_instance.