Require Import ZArith Omega MyNat Div2.

Ltac rem t x Hx :=
  pose (x := t); assert (Hx : x = t) by reflexivity;
    try rewrite <- Hx in *; clearbody x.

Ltac destruct_if t :=
  lazymatch t with
  | if ?t then _ else _ ⇒
      let Eqb := fresh "Eqb" in
      let b := fresh "b" in
      rem t b Eqb; destruct b
  | ?u ?v ⇒ destruct_if u || destruct_if v
  end.

Ltac auto_if :=
  match goal with
  | |- ?T ⇒ destruct_if T
  | H : ?T |- _ ⇒ destruct_if T
  end.

Ltac impossible := exfalso; discriminate || omega || tauto.

Ltac dumb_aux :=
  auto || auto with × || discriminate
  || intuition || tauto || omega.

Ltac dumb :=
  try solve
    [ impossible
    | subst; dumb_aux || (simpl in *; dumb_aux)
    | intros; dumb_aux || (simpl in *; dumb_aux) ].

Definition bset := nat → bool.

Definition singleton (x : nat) : bset :=
  fun y ⇒ if eq_nat_dec x y then true else false.

Definition inter (P Q : bset) : bset := fun i ⇒ andb (P i) (Q i).

Definition union (P Q : bset) : bset := fun i ⇒ orb (P i) (Q i).

Definition IBN : bool → nat := fun b ⇒ (if b then 1 else 0)%nat.

Fixpoint count n P :=
  match n with
  | O ⇒ O
  | S n ⇒ (IBN (P n) + count n P)%nat
  end.

Definition same (P Q : bset) : Prop := ∀ i, P i = Q i.

Fixpoint same_bool n P Q :=
  match n with
  | O ⇒ true
  | S n ⇒ if Bool.eqb (P n) (Q n) then same_bool n P Q else false
  end.


Lemma same_true n P Q : same_bool n P Q = true → ∀ i, i < n → P i = Q i.
Proof.
  induction n; intros E i L.
    impossible.

    simpl in E.
    auto_if; dumb.
    destruct (eq_nat_dec i n); dumb.
    apply Bool.eqb_true_iff.
    subst; dumb.
Qed.

Lemma same_false n P Q : same_bool n P Q = false → { i | i < n ∧ P i ≠ Q i }.
Proof.
  induction n; simpl; intros I.
    impossible.

    remember (Bool.eqb (P n) (Q n)) as b.
    destruct b.
      destruct IHn as (i, (Li, Di)); eauto.
      destruct (Bool.eqb_false_iff (P n) (Q n)); eauto.
Qed.


Lemma count_union_inter n P Q : (count n (union P Q) + count n (inter P Q) =
  count n P + count n Q)%nat.
Proof.
  induction n; auto; simpl.
  unfold IBN, union, inter in ×.
  remember (P n) as Pn; remember (Q n) as Qn.
  destruct Pn; destruct Qn; simpl; omega.
Qed.

Lemma not_empty P n : O < count n P → { i | i < n ∧ P i = true }.
Proof.
  intros CP.
  induction n.
    simpl in CP; impossible.

    simpl in CP.
    rem (P n) Pn EPn; destruct Pn.
      ∃ n; auto.

      destruct IHn as (i, Hi); dumb.
      ∃ i; dumb.
Defined.

Lemma count_union n P Q : count n (union P Q) =
  count n P + count n Q - count n (inter P Q).
Proof.
  rewrite <- count_union_inter.
  auto with ×.
Qed.

Lemma count_inter n P Q : count n (inter P Q) =
  count n P + count n Q - count n (union P Q).
Proof.
  rewrite <- count_union_inter.
  auto with ×.
Qed.

Lemma count_bound P n : count n P ≤ n.
Proof.
  induction n.
    dumb.
    simpl; destruct (P n); dumb.
Qed.

Lemma count_drawers : ∀ P Q n,
  n < count n P + count n Q → { i | i < n ∧ P i = true ∧ Q i = true }.
Proof.
  intros P Q n L.
  destruct (not_empty (inter P Q) n) as (i, Hi).
    rewrite count_inter.
    pose proof count_bound (union P Q) n.
    omega.

    ∃ i; intuition; unfold inter in *;
      destruct (P i); destruct (Q i); intuition.
Qed.

Definition injective domain (f : nat → nat) := ∀ i j,
  i < domain → j < domain → f i = f j → i = j.

Definition bounded domain bound (f : nat → nat) := ∀ i, i < domain → f i < bound.

Fixpoint has_antecedent (f : nat → nat) domain y :=
  match domain with
  | O ⇒ false
  | S domain ⇒ if eq_nat_dec (f domain) y then true else has_antecedent f domain y
  end.

Fixpoint number_of_pals_____OUT (f : nat → nat) domain x :=
  match domain with
  | O ⇒ 0
  | S domain ⇒ (if eq_nat_dec (f domain) (f x) then 1 else 0) + number_of_pals_____OUT f domain x
  end.

Definition image (f : nat → nat) domain : bset := has_antecedent f domain.

Lemma image_true f dom y :
  image f dom y = true → { i | i < dom ∧ f i = y }.
Proof.
  induction dom; intros A; dumb.
  destruct (eq_nat_dec (f dom) y) as [E|E].
    ∃ dom; eauto.

    destruct IHdom as (i, Hi).
      simpl in A.
      auto_if; dumb.

      ∃ i; dumb.
Qed.

Lemma image_false : ∀ f domain y,
  image f domain y = false → ∀ x, x < domain → f x = y → False.
Proof.
  intros f dom; induction dom; intros y T x Bx Exy.
    inversion Bx.

    simpl in T.
    auto_if; dumb.
      apply (IHdom y T x); dumb.
      cut (x ≠ dom); dumb.
Qed.

Lemma image_true_pr f dom y :
  ∀ i, i < dom → y = f i → image f dom y = true.
Proof.
  intros i L E.
  remember (image f dom y) as b; symmetry in Heqb.
  destruct b; [ reflexivity | exfalso ].
  exfalso; eapply image_false; eauto.
Qed.

Lemma image_false_pr f dom y :
  (∀ i, i < dom → y ≠ f i) → image f dom y = false.
Proof.
  intros NE.
  remember (image f dom y) as b; symmetry in Heqb.
  destruct b; [ exfalso | reflexivity ].
  destruct (image_true _ _ _ Heqb) as (n, Hn).
  specialize (NE n); intuition.
Qed.

Fixpoint sum n f := match n with O ⇒ O | S n ⇒ f n + sum n f end.

Definition is_duplicate f n := IBN (has_antecedent f n (f n)).

Definition is_duplicate_outside f bound n :=
  if lt_dec bound (f n) then IBN (has_antecedent f n (f n)) else O.

Fixpoint count_outside domain bound f :=
  match domain with
  | O ⇒ O
  | S domain ⇒ (if lt_dec (f domain) bound then 1 else 0) + count_outside domain bound f
  end.


Lemma image_injective_plus_one : ∀ domain f,
  same (image f (S domain)) (union (image f domain) (singleton (f domain))).
Proof.
  intros dom f i.
  unfold union, singleton.
  simpl; auto_if; auto with ×.
Qed.

Lemma count_same S T : same S T → ∀ n, count n S = count n T.
Proof.
  intros E n; induction n; simpl; dumb.
  rewrite E; dumb.
Qed.

Lemma count_same_bounded S T n : (∀ i, i < n → S i = T i) → count n S = count n T.
Proof.
  intros E; induction n; simpl; dumb.
  repeat (f_equal; dumb).
Qed.

Lemma count_union_singleton P a n :
  P a = false → a < n →
    count n (union P (singleton a)) = S (count n P).
Proof.
  intros Ma L.
  induction n; dumb.
  unfold union, singleton; simpl.
  destruct (eq_nat_dec a n).
    subst; rewrite Ma; simpl.
    f_equal; apply count_same_bounded; intros i Li.
    auto_if; dumb.

    rewrite plus_n_Sm, <- IHn; dumb.
    repeat f_equal; auto with ×.
Qed.

Lemma count_image_injective_plus_one : ∀ domain f bound,
  injective (S domain) f → bounded (S domain) bound f →
    count bound (image f (S domain)) = S (count bound (image f domain)).
Proof.
  intros dom f M If Bf.
  pose proof image_injective_plus_one dom f as HS.
  erewrite count_same; [ | eassumption ].
  rewrite count_union_singleton; dumb.
    apply image_false_pr; intros i L E.
    apply If in E; dumb.
Qed.

Lemma injective_S dom f : injective (S dom) f → injective dom f.
Proof.
  intros If i j Bi Bj.
  apply If; dumb.
Qed.

Lemma bounded_S dom M f : bounded (S dom) M f → bounded dom M f.
Proof.
  intros If i Bi.
  apply If; dumb.
Qed.

Lemma count_image_injective : ∀ domain bound f,
  injective domain f → bounded domain bound f →
    count bound (image f domain) = domain.
Proof.
  intros dom M f Injf Bf.
  induction dom.
    simpl; clear.
    induction M; auto.

    rewrite count_image_injective_plus_one; dumb.
    rewrite IHdom; dumb.
      apply injective_S; auto.
      apply bounded_S; auto.
Qed.