Require Import Utils.

Section FiniteSets.
  Open Scope bool_scope.

  Variable (X: finType).
  Definition finSet := finType_bool ^ X.
  Notation "x 'mem' a" := (applyVect a x = true) (at level 40).

  Definition makeSet (p: X -> Prop ) {decP : forall x, dec (p x)} : finSet := vectorise (fun x => (toBool (p x))).
  Arguments makeSet p {decP}.

  Lemma set_eq (a b : finSet) : (forall x, x mem a <-> x mem b) <-> a = b.
  Proof.
    split.
    - intros E. apply vector_eq. intros y. destruct (applyVect a y) eqn:H.
      + apply E in H. now rewrite H.
      + destruct (applyVect b y) eqn:H'; auto.
        exfalso. apply E in H'. congruence.
    - intros E x. now rewrite E.
  Qed.

  Lemma set_eq' (a b : finSet) : a = b -> (forall x, x mem a <-> x mem b).
  Proof.
    now apply set_eq.
  Qed.

  Lemma makeSet_eq (p q : X -> Prop) {decP : forall x, dec (p x)} {decQ : forall x, dec (q x)} : (forall x, p x <-> q x) -> makeSet p = makeSet q.
  Proof.
    intros E.
    unfold makeSet. apply vector_eq. intros y.
    rewrite 2apply_vectorise_inverse. destruct decP as [P|P], decQ as[Q|Q]; auto.
    - exfalso. apply Q, E, P.
    - exfalso. apply P, E, Q.
  Qed.

  Definition subset (a b:finSet) := forall x, x mem a -> x mem b.

  Lemma subset_eq a b : subset a b -> subset b a -> a = b.
  Proof.
    intros SA SB.
    apply set_eq. intros p.
    split; auto.
  Qed.

  Definition union (a b : finSet) :finSet := vectorise (fun x => applyVect a x || applyVect b x).
  Definition intersect(a b : finSet) : finSet := vectorise(fun x => applyVect a x && applyVect b x).

  Lemma makeSet_correct (p: X -> Prop) {decP:forall x, dec (p x)} : forall y, y mem (makeSet p) <-> p y.
  Proof.
    intros y. unfold makeSet. rewrite apply_vectorise_inverse.
    destruct (decP y) as [D|D].
    - split; auto.
    - split; auto. intros D'. exfalso. discriminate D'.
  Qed.

  Definition finUnion (Y:finType) (sets: finSet ^ Y) := fold_right union (makeSet (fun x => False)) (map (fun y => applyVect sets y) (elem Y)) .

  Lemma union_correct a b x: x mem (union a b) <-> x mem a \/ x mem b.
  Proof.
    unfold union. rewrite apply_vectorise_inverse.
    destruct (applyVect a x), (applyVect b x); simpl; tauto.
  Qed.

  Lemma union_assoc a b c : union a (union b c) = union (union a b) c.
  Proof.
    apply set_eq. intros x.
    rewrite !union_correct. tauto.
  Qed.

  Lemma intersect_correct a b x : x mem (intersect a b) <-> x mem a /\ x mem b.
  Proof.
    unfold intersect. rewrite apply_vectorise_inverse.
    destruct (applyVect a x), (applyVect b x); simpl; tauto.
  Qed.

  Lemma finUnion_correct (Y:finType) (sets: finSet ^ Y) : forall x, x mem (finUnion sets) <-> exists y, x mem (applyVect sets y).
  Proof.
    unfold finUnion.
    enough (forall l x, x mem fold_right union (makeSet (fun x => False)) (map (fun y => applyVect sets y) l) <-> (exists y, y el l /\ x mem (applyVect sets y))) as H.
    - intros x. split.
      + intros U. enough (exists y, y el (elem Y) /\ x mem (applyVect sets y)) as H'.
        * destruct H' as [y H']. now exists y.
        * now apply H.
      + intros [y U]. apply H. exists y. auto.
    - intros l. induction l as [|y l]; intros x.
      + cbn. split.
        * unfold makeSet. rewrite apply_vectorise_inverse.
          intros F. exfalso. destruct False_dec; auto. discriminate F.
        * intros [y [F _]]. now exfalso.
      + simpl. split.
        * intros U. apply union_correct in U. destruct U as [U|U].
          -- exists y. auto.
          -- apply IHl in U.
             destruct U as [y' [U1 U2]]. exists y'. auto.
        * intros [y' [[U|U] M]]; apply union_correct.
          -- subst y'. now left.
          -- right. apply IHl. now exists y'.
  Qed.

  Definition finSet_rem (a: finSet) (x:X):= makeSet (fun y => y mem a /\ y <> x).

  Lemma finSet_rem_correct a x y : y mem (finSet_rem a x) <-> y mem a /\ y <> x.
  Proof.
    unfold finSet_rem. apply makeSet_correct.
  Qed.

  Definition singleton x := makeSet (fun y => x = y).

  Lemma singleton_correct x y : x mem (singleton y) -> y = x.
  Proof.
    unfold singleton. now rewrite makeSet_correct.
  Qed.

  Lemma singleton_correct' x y : y = x -> x mem (singleton y).
  Proof.
    unfold singleton. now rewrite makeSet_correct.
  Qed.

  Definition emptySet := makeSet (fun x => False).

  Lemma emptySet_correct x: ~ x mem emptySet.
  Proof.
    unfold emptySet. now rewrite makeSet_correct.
  Qed.

  Fixpoint card' (a: finSet) (l: list X):= match l with
     | nil => 0
     | cons x l => if (applyVect a x)
               then S(card' a l)
               else card' a l
  end.
  Definition card (a:finSet) := card' a (elem X).



End FiniteSets.

Arguments makeSet {X} p {decP}.
Arguments emptySet {X}.

Notation "x 'mem' a" := (applyVect a x = true) (at level 40).

Definition set_map (X Y : finType) (a: finSet X) (f: X -> Y) := makeSet (fun y => exists x, x mem a /\ y = f x).

Lemma set_map_correct (X Y: finType) (a: finSet X) (f : X -> Y) y : y mem set_map a f <-> exists x, x mem a /\ y = f x.
Proof.
 unfold set_map. now rewrite makeSet_correct.
Qed.