Require Import Omega HFS.HF.

Implicit Types x y z : nat.

Fixpoint mul2 x :=
  match x with
  | O => O
  | S x' => S (S (mul2 x'))
  end.

Fixpoint pow2 x :=
  match x with
    |O => S O
    |S x => mul2 (pow2 x)
  end.

Fixpoint mod2 x : bool :=
  match x with
  | O => false
  | S x => negb (mod2 x)
  end.

Fixpoint div2 x :=
  match x with
  | O => O
  | S x' => if (mod2 x) then div2 x' else S (div2 x')
  end.

Fixpoint mem x y : bool :=
  match x with
  | O => mod2 y
  | S x' => mem x' (div2 y)
  end.

Definition adj x y :=
  if mem x y then y else pow2 x + y.

Definition without z x :=
  if mem z x then x - pow2 z else x.

Lemma mul2div2_evenodd x :
  mod2 x = false /\ mul2 (div2 x) = x \/
  mod2 x = true /\ S (mul2 (div2 x)) = x.
Proof.
  induction x.
  - simpl. left. tauto.
  - simpl. destruct IHx as [[-> IH2]|[-> IH2]]; simpl.
    + right. rewrite IH2. now split.
    + left. rewrite IH2. now split.
Qed.

Lemma mod2_add_mul2 x y :
  mod2 (mul2 x + y) = mod2 y.
Proof.
  induction x as [|x IHx]; [reflexivity | simpl].
  rewrite IHx. now destruct (mod2 y).
Qed.

Lemma div2mul2add x y :
  div2 (mul2 x + y) = x + div2 y.
Proof.
 induction x as [|x IHx]; [reflexivity | simpl].
 rewrite IHx. now destruct (mod2 (mul2 x + y)).
Qed.

Lemma muldiv_smaller x :
  x >= mul2 (div2 x).
Proof.
  destruct (mul2div2_evenodd x) as [[_ IH]|[_ IH]]; omega.
Qed.

Lemma div2SSlt x :
  div2 (S x) < S x /\ div2 (S (S x)) < S (S x).
Proof.
  induction x as [|x IHx]; simpl in *; [omega |].
  destruct (mod2 x); simpl in *; omega.
Qed.

Lemma div2Slt x :
  div2 (S x) < S x.
Proof.
  generalize (div2SSlt x). tauto.
Qed.

Lemma mul2_larger_S x :
  mul2 (S x) > (S x).
Proof.
  induction x; simpl in *; omega.
Qed.

Lemma mul2_larger_pres x y :
  x <= y -> mul2 x <= mul2 y.
Proof.
  induction 1; simpl; omega.
Qed.

Lemma pow2_larger x :
  pow2 x > x.
Proof.
  induction x; simpl; try omega.
  destruct (pow2 x) as [|n]; [omega |].
  pose proof mul2_larger_S n as H'. omega.
Qed.

Lemma emptyEq x :
  mem x O = false.
Proof.
  induction x; auto.
Qed.

Lemma mem_smaller x y :
  mem x y = true -> pow2 x <= y.
Proof.
  revert y. induction x as [|x IHx].
  - intros [|y]; simpl in *. discriminate. omega.
  - intros y A. simpl in *. specialize (IHx _ A).
    apply mul2_larger_pres in IHx. assert (H := muldiv_smaller y). omega.
Qed.

Lemma mem_add x y :
  mem x (pow2 x + y) = negb (mem x y).
Proof.
  revert y.
  induction x as [|x IHx]; intros y; simpl;
    [reflexivity | now rewrite div2mul2add].
Qed.

Lemma complete_induction (p : nat -> Type) :
  (forall x, (forall y, y < x -> p y) -> p x) -> forall x, p x.
Proof.
  intros H x. apply H.
  induction x as [|x IHx]; intros y H1.
  - omega.
  - destruct (lt_eq_lt_dec x y) as [[A| ->]|A].
    omega. now apply H. now apply IHx.
Qed.

Lemma dec_empty y :
  (y = O) + Sigma x, mem x y = true.
Proof.
  pattern y. revert y. apply complete_induction.
  intros [|y].
  - now left.
  - intros IH. right. destruct (IH _ (div2Slt y)) as [IHA|(x&B)].
    + destruct y as [|y]. now exists O.
      simpl in IHA. destruct (mod2 y); discriminate IHA.
    + now exists (S x).
Qed.

Lemma without_inverse z x :
  mem z x = true -> adj z (without z x) = x.
Proof.
  unfold adj, without. intros A. rewrite !A.
  pose proof le_plus_minus _ _ (mem_smaller A) as H.
  rewrite H, mem_add in A. now destruct (mem z (x - pow2 z)).
Qed.

Fact sip (p: nat -> Type) :
  p O -> (forall x y, p x -> p y -> p (adj x y)) -> forall x, p x.
Proof.
  intros base IH x.
  pattern x. revert x. apply complete_induction.
  intros y IH'. destruct (dec_empty y) as [-> | (x & A)].
  - assumption.
  - rewrite <- (without_inverse A). pose (H := pow2_larger x). apply IH.
    + apply IH'. apply mem_smaller in A. omega.
    + apply IH'. destruct y as [|y].
      now rewrite emptyEq in A. unfold without. rewrite A. omega.
Qed.

Lemma adj_mem x y :
  mem x (adj x y) = true.
Proof.
  unfold adj. case_eq (mem x y). trivial.
  rewrite mem_add. now intros ->.
Qed.

Lemma adj_diff z x y :
  z <> x -> mem z (adj x y) = mem z y.
Proof.
  revert x y. unfold adj. induction z as [|z IHz].
  - intros [|x] y A. now contradiction A.
    simpl. destruct (mem x (div2 y)). reflexivity.
    now rewrite mod2_add_mul2.
  - destruct x as [|x]; intros y A; simpl.
    + case_eq (mod2 y); simpl; now try intros ->.
    + assert (z <> x) by congruence.
      specialize (IHz _ (div2 y) H).
      destruct (mem x (div2 y)); [reflexivity |now rewrite div2mul2add].
Qed.

Fact cancel x y :
  adj x (adj x y) = adj x y.
Proof.
  unfold adj at 1. now rewrite (adj_mem x y).
Qed.

Instance nat_eq_dec :
  eq_dec nat.
Proof.
  intros x y. hnf. decide equality.
Qed.

Fact swap x y z :
  adj x (adj y z) = adj y (adj x z).
Proof.
  unfold adj at 1 3.
  decide (x = y) as [-> |e]. reflexivity.
  rewrite !adj_diff; try congruence.
  unfold adj. destruct (mem x z); destruct (mem y z); omega.
Qed.

Fact disjoint x y :
  adj x y <> O.
Proof.
  intros A. assert (H := adj_mem x y).
  now rewrite A, emptyEq in H.
Qed.

Lemma mem_adj_iff x y :
  mem x y = true <-> adj x y = y.
Proof.
  unfold adj. split.
  - now intros ->.
  - destruct (mem x y). reflexivity.
    pose proof pow2_larger x as H. omega.
Qed.

Fact membership z x y :
  adj z (adj x y) = adj x y -> z = x \/ adj z y = y.
Proof.
  rewrite <- !mem_adj_iff. intros A.
  decide (z = x) as [-> |B]. now left.
  right. now rewrite <- (adj_diff y B).
Qed.

Definition Ackermann : HF :=
  @HFS (EqType nat) O adj sip swap cancel disjoint membership.