Zagier's one-sentence proof for Fermat's two-square theorem

Zagier.v

(* Zagier's one-sentence proof for the two-square theorem in Coq (SSReflect) *)
(* - D. Zagier, "A One-Sentence Proof That Every Prime p \equiv 1 (\mod 4) Is a Sum of Two Squares", 
     The American Mathematical Monthly, Vol. 97, No. 2 (Feb., 1990), p. 144 *)

Set Implicit Arguments. Unset Strict Implicit.
From mathcomp Require Import all_ssreflect.

Lemma odd_fixedpoints {X:eqType} (f:X->X) (D:seq X):
  uniq D -> (forall x, x \in D -> (f x \in D) && (f (f x) == x)) ->
  odd (size D) <-> odd (count [pred x | f x == x] D).
Proof.
  move=>Duniq Hf; suff: ~~ odd (count [pred x | f x != x] D)
    by rewrite-[[pred x|f x != x]]/(predC[pred x|f x == x])
              -(count_predC [pred x|f x == x])oddD=>/negbTE->; case:odd.
  move:{2}(size D)(leqnn(size D))Duniq Hf=>n; elim:n D=>[[]//|n IH]D.
  move:leq_eqVlt=>->/orP[/eqP|/IH//].
  move:D cons_uniq=>[//|x D]->[]sz/andP[]/negP nxD Duniq Hf/=; subst n.
  move fx:(f x==x)=>[/eqP|]/=.
  -{ apply:IH=>//z zD; move:in_cons(zD)(Hf z)=>->->/(_ (orbT _))/andP[].
     rewrite in_cons=>/orP[/eqP?/eqP|->//]; congruence. }
  -{ move:mem_head in_cons(fx)(Hf x)=>->->->/(_ isT)/andP[]/=+/eqP/[dup]ffx.
     move:eq_sym fx=>->+/[dup]/perm_to_rem/permP->/=+->=>->fxD; apply/negPn/IH.
     +{ by move:size_rem leq_pred=>->. }
     +{ by apply rem_uniq. }
     +{ move=>z; rewrite(rem_filter (f x) Duniq)!mem_filter=>/andP[]/=/negPf.
        move:(Hf z)=>+fxz zD; rewrite!in_cons zD=>/(_(orbT _))/andP[]/orP[|->].
        *{ by move:eq_sym fxz=>->/negP?/eqP->. }
        *{ move?:(f z==f x)=>[/eqP|//]/eqP; congruence. } } }
Qed.

Require Import ZArith Lia.
Import Z Znumtheory. Open Scope Z. 
Canonical Z_eqType := Eval hnf in EqType Z (EqMixin eqb_spec).

Section seqZ.
  Definition seqZ x := [seq of_nat n | n <- iota 1 (to_nat x)].

  Lemma mem_seqZ x y: 0 < x <= y -> x \in seqZ y.
  Proof.
    case:(lt_ge_cases y 0)=>Hy; first by lia.
    pattern y; apply natlike_ind=>//=; first by lia.
    have seqZs: forall x, 0 <= x -> seqZ (succ x) = seqZ x ++ [:: succ x]
      by move=>??; rewrite-add_1_r/seqZ Z2Nat.inj_add//iotaD map_cat/=; do 2 f_equal; lia.
    move=>z Hz IH[Hx]/Zle_lt_or_eq[/lt_succ_r|->]; rewrite seqZs//mem_cat.
    -{ by move=>/(conj Hx)/IH->. }
    -{ by rewrite mem_head orbT. }
  Qed.

  Lemma seqZ_uniq x: uniq (seqZ x).
  Proof. rewrite/seqZ map_inj_uniq ?iota_uniq=>//>; apply Nat2Z.inj. Qed.
End seqZ.

Section Zagier's_Map.
  Variable k: Z.
  Let p := 4*k + 1.
  Variable Hp: prime p.

  Definition inDom '(x,y,z) := [/\ 0 < x , 0 < y, 0 < z & x^2 + 4*y*z = p].
  Let inDomb '(x,y,z) := [&& 0 <? x , 0 <? y, 0 <? z & x^2 + 4*y*z == p].

  Lemma inDomP xyz: inDom xyz <-> inDomb xyz.
  Proof.
    move:xyz=>[[x y]z]; rewrite/inDomb/inDom; split. 
    -{ by case=>/ltb_lt->/ltb_lt->/ltb_lt->/eqP. }
    -{ by move=>/and4P[/ltb_lt?/ltb_lt?/ltb_lt?/eqP]. }
  Qed.

  Let cubeDom: list (Z*Z*Z) :=
    [seq (xy, z) | xy <- [seq (x, y) | x <- seqZ p, y <- seqZ p], z <- seqZ p].

  Definition seqDom: list (Z*Z*Z) := [seq xyz <- cubeDom | inDomb xyz ].

  Lemma seqDomP xyz: inDom xyz <-> xyz \in seqDom.
  Proof.
    split; last by rewrite mem_filter=>/andP[]/inDomP.
    move:mem_filter=>->/[dup]D/inDomP->/=.
    Ltac ex_mem v e := exists v; first(apply mem_seqZ; split=>//; apply:(le_trans _ e);
                                   first apply le_mul_diag_r; lia).
    move:xyz D=>[[x y]z][???]Ep; apply/flatten_mapP; rewrite-Ep; exists (x, y).
    -{ by apply/flatten_mapP; ex_mem x (x*x); apply/mapP; ex_mem y (y*z). }
    -{ by apply/mapP; ex_mem z (z*y). }
  Qed.
  
  Lemma seqDom_uniq: uniq seqDom.
  Proof.
    by apply filter_uniq; repeat apply allpairs_uniq; apply seqZ_uniq||case=>??[].
  Qed.

  Definition f '(x,y,z) :=
    if x <? y - z then (x + 2 * z, z, y - x - z)
    else if (2 * y <? x) then (x - 2 * y, x - y + z, y)
    else (2 * y - x, y, x - y + z).

  Ltac caseLT_aux H := match goal with |- context[if ?b then _ else _] =>
                                       case_eq b=>[/ltb_lt|/ltb_nlt]H end.
  Tactic Notation "caseLT" := let H := fresh "_LT" in caseLT_aux H.
  Tactic Notation "caseLT" ident(H) := caseLT_aux H.
  Ltac caseDom := caseLT xLyz; last caseLT H2yx.

  Lemma f_closed xyz: inDom xyz -> inDom (f xyz).
  Proof.
    move:xyz=>[[x y]z][???]Ep; rewrite/f/inDom; caseDom; split; try lia.
    -{ move:H2yx Ep=>/le_ngt/Zle_lt_or_eq[|->]; lia.  }
    -{ move:xLyz Ep Hp=>/le_ngt/Zle_lt_or_eq[|<-]; first lia.
       have->:(y-z)^2+4*y*z=(y+z)*(y+z) by ring.
       by move=><-/square_not_prime. }
  Qed.

  Lemma f_involutory xyz: inDom xyz -> f (f xyz) = xyz.
  Proof.
    by move:xyz=>[[x y]z][????]; rewrite/f; repeat caseLT; do 2 f_equal; lia.
  Qed.

  Lemma f_fixedpoint x y z:
    inDom (x,y,z) -> f (x,y,z) == (x,y,z) <-> [&& x == 1, y == 1 & z == k].
  Proof.
    case=>Hx Hy Hz Ep; split.
    -{ rewrite/f; caseDom=>/eqP; rewrite(mul_comm 2)=>[][]; try lia.
       match goal with |- context[?e-x=x] => have<-:(2*y=e) by clear;case:y;lia end.
       move=>{xLyz H2yx}_ Hxyzz; (have ?: x = y by lia); subst y; apply/and3P.
       have ?: x = 1. {
         move:Ep Hp; have->:x^2+4*x*z=x*(x+4*z) by ring.
         move=><-/prime_divisors/(_ x(divide_factor_l _ _))[|[|[]]]; try lia.
         move=>/eq_sym; rewrite mul_id_r; lia. }
       subst; split=>//; apply/eqP; lia. }
    -{ rewrite/f=>/and3P[]/eqP->/eqP->/eqP->; apply/eqP;
       repeat caseLT; do 2 f_equal; lia. }
  Qed.

  Lemma seqDom_odd: ssrnat.odd (size seqDom).
  Proof.
    have Hk:1<=k by move:(prime_ge_2 _ Hp); lia.
    move:seqDom_uniq=>Huniq; rewrite(odd_fixedpoints (f:=f))=>//.
    -{ rewrite(eq_in_count (a1:=fun x => f x == x) (a2:=pred1 (1,1,k))).
       +{ rewrite (count_uniq_mem _ Huniq).
          match goal with |- context[nat_of_bool ?e] => have->//:e end.
          rewrite/seqDom mem_filter/cubeDom; apply/andP; split.
          *{ apply/and4P; split; apply/ltb_lt||(rewrite/p; apply/eqP); lia. }
          *{ apply/allpairsP; exists(1,1,k); split=>//=; last by apply mem_seqZ; lia.
             apply/allpairsP; exists(1,1); split=>//=; apply mem_seqZ; lia. } }
       +{ case=>[[x y]z]/seqDomP/f_fixedpoint H; apply/sym_eq/eqP; move:H.
          case:(f(x,y,z)==(x,y,z))=>[[/(_ isT)]/and3P[]/eqP->/eqP->/eqP->//|[_]H[]].
          by move=>???; apply/notF/H; subst; rewrite!eqtype.eq_refl. } }
    -{ by move=>?/seqDomP/[dup]/f_closed/seqDomP->/f_involutory/eqP->. }
  Qed.

  Definition g '(x,y,z):Z*Z*Z := (x,z,y).

  Lemma g_closed xyz: inDom xyz -> inDom (g xyz).
  Proof. move:xyz=>[[??]?][]; split; lia. Qed.

  Lemma g_involutory xyz: g (g xyz) = xyz.
  Proof. by move:xyz=>[[]]. Qed.

  Lemma g_fixedpoint_exists: exists xyz, inDom xyz /\ g xyz = xyz.
  Proof.
    move E:([seq xyz <- seqDom | g xyz == xyz ])=>G0.
    have: ssrnat.odd (size G0). {
      move:E seqDom_odd seqDom_uniq g_involutory size_filter=><-???->.
      by rewrite-odd_fixedpoints=>//?/seqDomP/g_closed/seqDomP->; apply/eqP. }
    case E0:G0=>//[xyz?] _; exists xyz; have: xyz \in G0 by rewrite E0 mem_head.
    by rewrite-E mem_filter=>/andP[/eqP?/seqDomP].
  Qed.
End Zagier's_Map.

Theorem TwoSquareTheorem k:
  let p := 4*k + 1 in prime p -> exists x y, x^2 + y^2 = p.
Proof.
  move=>p; rewrite/p=>/g_fixedpoint_exists[][][x]y?[][]????[];exists x,(2*y); lia.
Qed.