From Coq Require Import Arith Lia.
Definition dec X := {X} + {~X}.

(*** Size Recursion Operator *)

Definition size_rec {X: Type} (sigma: X -> nat) {p: X -> Type} :
  (forall x, (forall y, sigma y < sigma x -> p y) -> p x) ->
  (forall x, p x).
Proof.
  intros F.
  enough (forall n x, sigma x < n -> p x) as H.
  { intros x. apply (H (S (sigma x))). lia. }
  induction n as [|n IH]; intros x H.
  - exfalso. lia.
  - apply F. intros y H1. apply IH. lia.
Defined.

(** Alternatively, we can define a size recursion operator 
    with the refine tactic thus making fully explicit the 
    computational structure.  All proofs are inserted by lia. *)                   

Module Size_rec.
Section Size_rec.
  Variables (X: Type) (sigma: X -> nat) (p: X -> Type).
  Variable (F: forall x, (forall y, sigma y < sigma x -> p y) -> p x).

  Definition size_rec': forall n x, sigma x < n -> p x.
  Proof.
    refine (fix f n := match n with 0 => _ | S n' => _ end).
    - refine (fun x h => match _: False with end). lia.
    - refine (fun x h => F x (fun y h' => f n' y _)). lia.
  Defined.
  Definition size_rec : forall x, p x.
  Proof.
    refine (fun x => size_rec' (S (sigma x)) x _). lia.
  Defined.
End Size_rec.
Arguments size_rec {X} sigma {p}.
End Size_rec.

(** We define a size recursion operator for binary predicates *)

Definition size_rec2 {X Y: Type} (sigma: X -> Y -> nat) {p: X -> Y -> Type} :
  (forall x y, (forall x' y', sigma x' y' < sigma x y -> p x' y') -> p x y) ->
  (forall x y, p x y).
Proof.
  intros F.
  enough (forall '(x,y), p x y) as H.
  { intros x y. apply (H (x,y)). } 
  refine (size_rec (fun '(x,y) => sigma x y) (fun '(x,y) IH => _)). cbn in IH.
  apply F. intros x' y' H. apply (IH (x',y')), H.
Defined.

(*** Least Witness operator *)

Definition safe (p: nat -> Prop) n := forall k, p k -> k >= n.
Definition least (p: nat -> Prop) n := p n /\ safe p n.

Fact safe_O p :
  safe p 0.
Proof.
  intros k _. lia.
Qed.

Fact safe_S p n :
  safe p n -> ~p n -> safe p (S n).
Proof.
  intros H1 H2 k H3.
  enough (k <= n -> False) by lia. intros H4.
  specialize (H1 k H3).
  enough (k = n) as -> by auto.
  lia.
Qed.

Fact LW  {p: nat -> Prop} {n} :
  (forall x, dec (p x)) -> p n -> sig (least p).
Proof.
  intros D H.
  enough (forall k, safe p k -> sig (least p)) as H1.
  { apply (H1 0). apply safe_O. }
  refine (size_rec (fun k => n - k) (fun k IH H1 => _)). cbn in IH.
  destruct (D k) as [H2|H2].
  - exists k. hnf. auto.
  - assert (H3: safe p (S k)).
    { apply safe_S; assumption. }
    apply (IH (S k)). 2:exact H3. apply H3 in H. lia.
Defined.

(*** Euclidean Division *)

Definition delta x y z : Prop :=
  z * S y <= x < S z * S y.

Fact delta_fun x y z z' :
  delta x y z -> delta x y z' -> z = z'.
Proof.
  unfold delta. nia.
Qed.

Fact delta_base x y :
  x <= y -> delta x y 0.
Proof.
  unfold delta. lia.
Qed.
 
Fact delta_step x y z :
  x > y -> delta (x - S y) y z -> delta x y (S z).
Proof.
  unfold delta. lia.
Qed.
 
Definition Div :
  forall x y, { z | delta x y z }.
Proof.
  intros x y. revert x.
  refine (size_rec (fun x => x) (fun x IH => _)).
  destruct (le_lt_dec x y) as [H|H].
  - exists 0. apply delta_base, H.
  - specialize (IH (x - S y)) as [z IH]. {lia. }
    exists (S z). apply delta_step; assumption.
Defined.

Definition D x y := proj1_sig (Div x y).

Compute D 100 19.

Fact D_correct x y :
  delta x y (D x y).
Proof.
  apply (proj2_sig (Div x y)).
Qed.

Definition DIV f x y :=
  if le_lt_dec x y then 0 else S (f (x - S y) y).

Fact DIV_agree f :
  (forall x y, delta x y (f x y)) <-> (forall x y, f x y = DIV f x y).
Proof.
  split.
  - intros H x y.
    apply (delta_fun x y).
    { apply H. }
    unfold DIV. destruct le_lt_dec as [H1|H1].
    + apply delta_base, H1.
    + apply delta_step. exact H1. apply H.
  - intros E.
    refine (size_rec (fun x => x) (fun x IH y => _)).
    cbn in IH. rewrite E.
    unfold DIV. destruct le_lt_dec as [H1|H1].
    + apply delta_base, H1.
    + apply delta_step. exact H1. apply IH. lia.
Qed.

Fact DIV_D x y :
  D x y = DIV D x y.
Proof.
  apply DIV_agree, D_correct.
Qed.

(*** Division Theorem *)

Definition M x y := x - D x y * S y.

Fact DM_eq x y :
  x = D x y * S y + M x y.
Proof.
  unfold M.
  enough (x >= D x y * S y) by lia.
  generalize (D_correct x y). unfold delta. lia.
Qed.

Fact M_le x y :
  M x y <= y.
Proof.
  unfold M.
  generalize (D_correct x y). unfold delta. lia.
Qed.

Fact div_unique x y a b :
  x = a * S y + b -> b <= y -> a = D x y /\ b = M x y.
Proof.
  intros H1 H2. unfold M.
  enough (a = D x y) by lia.
  apply (delta_fun x y).
  - unfold delta. lia.
  - apply D_correct.
Qed.

Fact div_unique' a b a' b' y:
  a * S y + b = a' * S y + b' ->
  b <= y -> b' <= y ->
  a=a' /\ b=b'.
Proof.
  set (x:= a * S y + b).
  set (x':= a' * S y + b').
  intros H1 H2 H3.
  enough (a = a') by lia.
  enough (a = D x y /\ a' = D x' y) by nia.
  split.
  - eapply div_unique. 2:exact H2. reflexivity.
  - eapply div_unique. 2:exact H3. reflexivity.
Qed.

(*** Greatest Common Divisors *)

Definition divides n x : Prop
  := exists k, x = k * n.

Fact divides_zero n :
  divides n 0.
Proof.
  exists 0. reflexivity.
Qed.

Fact divides_minus x y n :
  x <= y -> divides n x -> divides n y <->  divides n (y - x).
Proof.
  intros H [k ->]. split.
  - intros [l ->]. exists (l-k). nia.
  - intros [l H1]. exists (k + l). nia.
Qed.

Fact divides_bnd n x :
  x > 0 -> divides n x -> n <= x.
Proof.
  intros H [k ->]. destruct k.
  - exfalso. lia.
  - nia.
Qed.
 
Fact divides_eq' x y :
  x > 0 -> y > 0 -> divides x y -> divides y x -> x = y.
Proof.
  intros H1 H2 H3%divides_bnd H4%divides_bnd; lia.
Qed.

Fact divides_eq x y :
  (forall n, divides n x <-> divides n y) -> x = y.
Proof.
  destruct x, y; intros H.
  - reflexivity.
  - exfalso.
    enough (S(S y) <= S y) by lia.
    apply divides_bnd. lia.
    apply H, divides_zero.
  - exfalso.
    enough (S(S x) <= S x) by lia.
    apply divides_bnd. lia.
    apply H, divides_zero.
  - apply divides_eq'. lia. lia.
    + apply H. exists 1. lia.
    + apply H. exists 1. lia.
Qed.

Definition gamma x y z : Prop :=
  forall n, divides n z <-> divides n x /\ divides n y.

Fact gamma_fun x y z z' :
  gamma x y z -> gamma x y z' -> z = z'.
Proof.
  intros H H'.
  apply divides_eq. intros n. split.
  - intros H1. apply H',H,H1.
  - intros H1. apply H,H',H1.
Qed.

Fact gamma_zero x :
  gamma 0 x x.
Proof.
  intros n. generalize (divides_zero n). tauto.
Qed.

Fact gamma_minus x y z :
  x <= y -> gamma x y z <-> gamma x (y - x) z.
Proof.
  intros H.
  split; intros H1 n;
    specialize (H1 n);
    assert (H2:= divides_minus _ _ n H);
    tauto.
Qed.
  
Fact gamma_sym x y z :
  gamma x y z <-> gamma y x z.
Proof.
  firstorder.
Qed.

(** Gcd function *)

Definition Gcd :
  forall x y, { z | gamma x y z }.
Proof.
  refine (size_rec2 (fun x y => x + y) (fun x y IH => _)). cbn in IH.
  destruct x.
  { exists y. apply gamma_zero. }
  destruct y.
  { exists (S x). apply gamma_sym, gamma_zero. }
  destruct (le_lt_dec x y) as [H|H].
  - specialize (IH (S x) (y - x)) as [z IH]. lia.
    exists z. apply gamma_minus. lia. exact IH.
  - specialize (IH (x - y) (S y)) as [z IH]. lia.
    exists z. apply gamma_sym, gamma_minus. lia. apply gamma_sym, IH.
Defined.

Definition gcd x y := proj1_sig (Gcd x y).

Fact gcd_correct x y :
  gamma x y (gcd x y).
Proof.
  apply (proj2_sig (Gcd x y)).
Qed.

Compute gcd 252 105.

(** Recursive specification *)

Definition GCD f x y :=
  match x, y with
  | 0, y' => y'
  | S x', 0 => S x'
  | S x', S y' => if le_lt_dec x' y' then f (S x') (y' - x') else f (x' - y') (S y')
  end.

Fact GCD_agree f :
  (forall x y, gamma x y (f x y)) <-> (forall x y, f x y = GCD f x y).
Proof.
  split.
  - intros H x y.
    apply (gamma_fun x y).
    { apply H. }
    destruct x.
    { cbn. apply gamma_zero. }
    destruct y; cbn [GCD].
    { apply gamma_sym, gamma_zero. }
    destruct le_lt_dec as [H1|H1].
    + apply gamma_minus. lia. apply H.
    + apply gamma_sym, gamma_minus. lia. apply gamma_sym, H.
  - intros E.
    refine (size_rec2 (fun x y => x + y) (fun x y IH => _)). cbn in IH.
    destruct x.
    { rewrite E. cbn. apply gamma_zero. }
    destruct y; rewrite E; cbn.
    { apply gamma_sym, gamma_zero. }
    destruct le_lt_dec as [H|H].
    + apply gamma_minus. lia. apply IH. lia.
    + apply gamma_sym, gamma_minus. lia. apply gamma_sym, IH. lia.
Qed. 

Fact gcd_eq x y :
  gcd x y = GCD gcd x y.
Proof.
  revert x y. apply GCD_agree, gcd_correct.
Qed.

(*** Discrete Inversion *)

Section Discrete_Inversion.
  Variable f: nat -> nat.
  Variable f_strict: forall n, f n < f (S n).

  Lemma f_strict_mono' n n' :
    n < n' -> f n < f n'.
  Proof.
    induction n' as [|n' IH].
    - lia.
    - intros H.
      assert (n=n' \/ n < n') as [<-|H1%IH] by lia.
      + apply f_strict.
      + generalize (f_strict n'). lia.
  Qed.

  Lemma f_strict_mono n n' :
    n < n' <-> f n < f n'.
  Proof.
    split.
    - apply f_strict_mono'.
    - assert (n = n' \/ n < n' \/ n' < n) as [<-|[H|H%f_strict_mono']]; lia.
  Qed.
  
  Lemma f_mono n n' :
    n <= n' <-> f n <= f n'.
  Proof.
    split; intros H.
    - assert (n = n' \/ n < n') as [<-|?%f_strict_mono']; lia.
    - assert (n <= n' \/ n' < n) as [?|?%f_strict_mono']; lia.
  Qed.

  Lemma f_zero n x :
    f n <= x -> f 0 <= x.
  Proof.
    enough (f 0 <= f n) by lia.
    apply -> f_mono. lia.
  Qed.    
  
  Fact f_inc n :
    n <= f n.
  Proof.
    induction n as [|n IH]. lia.
    specialize (f_strict n). lia.
  Qed.
 
  Lemma f_term k x :
    f (S k) <= x -> x - S k < x - k.
  Proof.
    generalize (f_inc (S k)). lia.
  Qed.
 
  Definition phi x n : Prop :=
    f n <= x < f (S n).

  Fact phi_f n :
    phi (f n) n.
  Proof.
    unfold phi. auto.
  Qed.

  Fact phi_inv k n :
    phi (f n) k -> n = k.
  Proof.
    intros [H1%f_mono H2%f_strict_mono]. lia.
  Qed.

  Fact phi_fun x n n' :
    phi x n -> phi x n' -> n = n'.
  Proof.
    unfold phi. intros H1 H2.
    assert (n = n' \/ S n <= n' \/ S n' <= n)
      as [<-|[H%f_mono|H%f_mono]]; lia.
  Qed.

  Fact phi_zero x n :
    phi x n -> f 0 <= x.
  Proof.
    intros [H _]. eapply f_zero, H.
  Qed.

  Definition psi x k n : Prop :=
    if le_lt_dec (f (S k)) x then  phi x n else k = n.

  Fact psi_fun x k n n' :
    psi x k n -> psi x k n' -> n = n'.
  Proof.
    unfold psi. 
    destruct le_lt_dec as [H|H].
    - apply phi_fun.
    - congruence.
  Qed.

  Fact psi_phi x k n :
    f k <= x -> psi x k n -> phi x n.
  Proof.
    unfold psi.
    destruct le_lt_dec as [H|H].
    - auto.
    - unfold phi. intros H1 <-. auto.
  Qed.

  Fact psi_base x k:
    x < f (S k) -> psi x k k.
  Proof.
    unfold psi.
    destruct le_lt_dec as [H|H].
    - intros H1. exfalso. lia.
    - auto.
  Qed.
  
  Fact psi_step x k n :
    f (S k) <= x  -> psi x (S k) n -> psi x k n.
  Proof.
    intros H1 H2.
    unfold psi. destruct le_lt_dec as [H|H].
    - eapply psi_phi; eassumption.
    - exfalso. lia.
  Qed.

  (** Existence *)

  Definition Inv x :
    forall k, { n | psi x k n }.
  Proof.
    refine (size_rec (fun k => x - k) (fun k IH => _)).
    cbn in IH.
    destruct (le_lt_dec (f (S k)) x) as [H|H].
    - specialize (IH (S k)) as [n H1].
      + apply f_term, H. 
      + exists n. eapply psi_step; eassumption.
    - exists k. apply psi_base, H.
  Defined.
  
  Definition inv x k := proj1_sig (Inv x k).
  
  Fact inv_correct x k :
    psi x k (inv x k).
  Proof.
    apply (proj2_sig (Inv x k)).
  Qed.
    
  Fact inv_correct' x n :
    phi x n <-> inv x 0 = n /\ f 0 <= x.
  Proof.
    generalize (inv_correct x 0). unfold psi.
    destruct le_lt_dec as [H|H]; intros H1; split; intros H2.
    - split.
      + eapply phi_fun; eassumption.
      + eapply f_zero, H.
    - destruct H2 as [<- _]. exact H1.
    - assert (H3: f 0 <= x) by eapply f_zero, H2.
      enough (n=0) as -> by auto.
      eapply phi_fun. exact H2. unfold phi; auto.
    - destruct H2 as [<- H2]. rewrite <- H1. split; auto.
  Qed.
  
  Fact inv_correct_f n :
    inv (f n) 0 = n.
  Proof.
    apply inv_correct', phi_f.
  Qed.

  (** Recursive specification *)

  Definition INV (g: nat -> nat) (x k: nat) : nat
    := if le_lt_dec (f (S k)) x then g (S k) else k.

  Fact INV_agree x g :
    (forall k, psi x k (g k)) <-> (forall k, g k = INV g x k).
  Proof.
    split.
    - intros H k.
      apply (psi_fun x k).
      { apply H. }
      unfold INV. destruct le_lt_dec as [H1|H1]. 
      + apply psi_step. exact H1. apply H.
      + apply psi_base, H1.
  - intros E.
    refine (size_rec (fun k => x - k) (fun k IH => _)).
    cbn in IH. rewrite E.
    unfold INV. destruct le_lt_dec as [H|H].
    + apply psi_step. exact H.
      apply IH. apply f_term, H.
    + apply psi_base, H.
  Qed.

  Fact inv_eq x k :
    inv x k = INV (inv x) x k.
  Proof.
    revert k. apply INV_agree, inv_correct.
  Qed.
  
End Discrete_Inversion.

Definition test_sqrt : nat -> nat.
Proof.
  refine (fun x => inv (fun n => n * n) _ x 0).
  intros *; lia.
Defined.

Compute test_sqrt 80.
Compute test_sqrt 85.
Compute test_sqrt 100.

Definition test_div : nat -> nat -> nat.
Proof.
  refine (fun x y => inv (fun n => n * S y) _ x 0).
  intros *; lia.
Defined.

Compute test_div 35 6.