Set Implicit Arguments.

Require Import Omega.
Require Import Wf_nat.
Require Import Relation_Operators.
Require Import Coq.Wellfounded.Lexicographic_Product.

(* simultaneously produces a proof of the arith. equation eq and uses it to rewrite in the goal *)
Ltac omega_rewrite eq :=
  let Hf := fresh in
  assert eq as Hf; [omega | rewrite Hf; clear Hf].

(* simultaneously produces a proof of the arith. equation eq and uses it to rewrite in H *)
Ltac omega_rewrite_ctxt eq H :=
  let Hf := fresh in
  assert eq as Hf; [omega | rewrite Hf in H; clear Hf].

Definition sp_natpair := symprod _ _ lt lt.

Lemma wf_pair : well_founded sp_natpair.
Proof.
  apply wf_symprod; exact lt_wf.
Qed.

Fixpoint iter{X}(f : X -> X)(n : nat)(x : X) : X :=
  match n with
  | 0   => x
  | S m => f (iter f m x)
  end.

Lemma iter_plus{X} : forall (f : X -> X)(m n : nat)(x : X), iter f (m + n) x = iter f m (iter f n x).
Proof.
  intros f m; induction m; intros.
  auto.
  simpl; congruence.
Qed.

Lemma iter_S{X} : forall (f : X -> X)(n : nat)(x : X), iter f (S n) x = iter f n (f x).
Proof.
  intros.
  omega_rewrite (S n = n + 1).
  rewrite iter_plus; auto.
Qed.

Definition cyclic{X}(f : X -> X) := forall x x', exists n, iter f n x = x'.

Definition inj{X Y}(f : X -> Y) := forall x x', f x = f x' -> x = x'.

Lemma iter_aux_lemma : forall (X : Type)(f : X -> X)(x x' y : X)(p : nat * nat),
  f x = y -> f x' = y -> iter f (fst p) y = x -> iter f (snd p) y = x' -> x = x'.
Proof.
  intros X f x x' y p Hx Hx' Hyx Hyx'.
  induction (wf_pair p) as [p _ IHp].
  destruct p as [m n];  simpl in Hyx,Hyx'.
  destruct (dec_eq_nat m n) as [mn_eq | mn_neq].
  congruence.
  destruct (nat_total_order _ _ mn_neq).
  omega_rewrite_ctxt (n = S (pred (n - m)) + m) Hyx'.
  rewrite iter_plus, Hyx, iter_S, Hx in Hyx'.
  apply (IHp (m, pred (n-m))); [ right; omega | auto | auto ].
  omega_rewrite_ctxt (m = S (pred (m - n)) + n) Hyx.
  rewrite iter_plus, Hyx', iter_S, Hx' in Hyx.
  apply (IHp (pred (m-n), n)); [ left; omega | auto | auto ].
Qed.

Theorem cyclic_inj : forall (X : Type)(f : X -> X), cyclic f -> inj f.
Proof.
  intros X f fcyc x x' Hxx'.
  destruct (fcyc (f x) x) as [m Hm].
  destruct (fcyc (f x) x') as [n Hn].
  apply (@iter_aux_lemma _ f _ _ (f x) (m,n)); auto.
Qed.