Library prosa.classic.util.counting

Require Import prosa.classic.util.tactics prosa.classic.util.ord_quantifier.
From mathcomp Require Import ssreflect ssrbool eqtype ssrnat seq fintype bigop.

(* Additional lemmas about counting. *)
Section Counting.

  Lemma count_or :
     (T: eqType) (l: seq T) P Q,
      count (fun xP x || Q x) l count P l + count Q l.
  Proof.
    intros T l P Q; rewrite -count_predUI.
    apply leq_trans with (n := count (predU P Q) l);
      last by apply leq_addr.
    by apply sub_count; red; unfold predU; simpl.
  Qed.

  Lemma sub_in_count :
     (T: eqType) (l: seq T) (P1 P2: T bool),
      ( x, x \in l P1 x P2 x)
      count P1 l count P2 l.
  Proof.
    intros T l P1 P2 SUB.
    apply leq_trans with (n := count (fun xP1 x && (x \in l)) l);
      first by apply eq_leq, eq_in_count; red; movex INx; rewrite INx andbT.
    by apply sub_count; red; movex /andP [Px INx]; apply SUB.
  Qed.

  Lemma count_sub_uniqr :
     (T: eqType) (l1 l2: seq T) P,
      uniq l1
      {subset l1 l2}
      count P l1 count P l2.
  Proof.
    intros T l1 l2 P UNIQ SUB.
    rewrite -!size_filter uniq_leq_size ?filter_uniq // ⇒ x.
    by rewrite !mem_filter =>/andP [-> /SUB].
  Qed.

  Lemma count_pred_inj :
     (T: eqType) (l: seq T) (P: T bool),
      uniq l
      ( x1 x2, P x1 P x2 x1 = x2)
      count P l 1.
  Proof.
    intros T l P UNIQ INJ.
    induction l as [| x l']; [by done | simpl in *].
    {
      move: UNIQ ⇒ /andP [NOTIN UNIQ].
      specialize (IHl' UNIQ).
      rewrite leq_eqVlt in IHl'.
      move: IHl' ⇒ /orP [/eqP ONE | ZERO]; last first.
      {
        rewrite ltnS leqn0 in ZERO.
        by move: ZERO ⇒ /eqP ->; rewrite addn0 leq_b1.
      }
      destruct (P x) eqn:Px; last by rewrite add0n ONE.
      {
        move: ONE ⇒ /eqP ONE.
        rewrite eqn_leq in ONE; move: ONE ⇒ /andP [_ ONE].
        rewrite -has_count in ONE.
        move: ONE ⇒ /hasP ONE; destruct ONE as [y IN Py].
        specialize (INJ x y Px Py); subst.
        by rewrite IN in NOTIN.
      }
    }
  Qed.

  Lemma count_exists :
     (T: eqType) (l: seq T) n (P: T 'I_n bool),
      uniq l
      ( y x1 x2, P x1 y P x2 y x1 = x2)
      count (fun (y: T) ⇒ [ x in 'I_n, P y x]) l n.
  Proof.
    intros T l n P UNIQ INJ.
    induction n.
    {
      apply leq_trans with (n := count pred0 l); last by rewrite count_pred0.
      apply sub_count; red; intro x.
      by rewrite exists_ord0 //.
    }
    {
      apply leq_trans with (n := n + 1); last by rewrite addn1.
      apply leq_trans with (n := count (fun y[ x in 'I_n, P y (widen_ord (leqnSn n) x)] || P y ord_max) l).
      {
        apply eq_leq, eq_count; red; intro x.
        by rewrite exists_recr //.
      }
      eapply (leq_trans (count_or _ _ _ _)).
      apply leq_add.
      {
        apply IHn.
        {
          intros y x1 x2 P1 P2.
          by specialize (INJ (widen_ord (leqnSn n) y) x1 x2 P1 P2).
        }
      }
      {
        apply count_pred_inj; first by done.
        by intros x1 x2 P1 P2; apply INJ with (y := ord_max).
      }
    }
  Qed.

End Counting.