-- Perfect squares and sums of two squares: squares are closed under
-- products, a square is 0 or 1 modulo 3 and modulo 4, an odd square is 1
-- modulo 8, there is no square strictly between consecutive squares, and the
-- product of two sums of two squares is a sum of two squares
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.NumberTheory.PerfectSquares
import ModularMore

predicate IsSquare (n : Nat) := ∃ k, Nat(k) ∧ n = k * k

theorem IsSquare.intro (k : Nat)
  shows IsSquare(k * k)
proof
  witness k
  rfl

theorem IsSquare.zero
  shows IsSquare(0)
proof
  witness 0
  ring

theorem IsSquare.one
  shows IsSquare(1)
proof
  witness 1
  ring

theorem IsSquare.pow_two (k : Nat)
  shows IsSquare(k ^ 2)
proof
  have square := pow.two at k
  rw square
  exact IsSquare.intro at k

theorem IsSquare.pow_double (a k : Nat)
  shows IsSquare(a ^ (2 * k))
proof
  have double := pow.double_exponent at a, k
  rw double
  exact IsSquare.pow_two at a ^ k

theorem IsSquare.mul (m n : Nat)
  requires first : IsSquare(m)
  requires second : IsSquare(n)
  shows IsSquare(m * n)
proof
  obtain ⟨a, aNat, aShape⟩ := first
  obtain ⟨b, bNat, bShape⟩ := second
  witness a * b
  rw [aShape, bShape]
  ring

theorem IsSquare.mod_four (n : Nat)
  requires square : IsSquare(n)
  shows n % 4 = 0 ∨ n % 4 = 1
proof
  obtain ⟨k, kNat, shape⟩ := square
  rw shape
  have four : 4 ≠ 0 by arith
  have reduce := mod.mul at 4, k, k with four
  rw reduce
  have alternatives := mod.four_cases at k
  rcases alternatives with zero | one | two | three
  . left
    rw zero
    compute
  . right
    rw one
    compute
  . left
    rw two
    compute
  . right
    rw three
    compute

theorem IsSquare.mod_three (n : Nat)
  requires square : IsSquare(n)
  shows n % 3 = 0 ∨ n % 3 = 1
proof
  obtain ⟨k, kNat, shape⟩ := square
  rw shape
  have three : 3 ≠ 0 by arith
  have reduce := mod.mul at 3, k, k with three
  rw reduce
  have alternatives := mod.three_cases at k
  rcases alternatives with zero | one | two
  . left
    rw zero
    compute
  . right
    rw one
    compute
  . right
    rw two
    compute

theorem IsSquare.mod_three_ne_two (n : Nat)
  requires square : IsSquare(n)
  shows n % 3 ≠ 2
proof
  have alternatives := IsSquare.mod_three with square
  intro two
  rw two at alternatives
  cases alternatives with zero | one
  . arith
  . arith

theorem IsSquare.mod_four_ne_two (n : Nat)
  requires square : IsSquare(n)
  shows n % 4 ≠ 2
proof
  have alternatives := IsSquare.mod_four with square
  intro two
  rw two at alternatives
  cases alternatives with zero | one
  . arith
  . arith

theorem IsSquare.mod_four_ne_three (n : Nat)
  requires square : IsSquare(n)
  shows n % 4 ≠ 3
proof
  have alternatives := IsSquare.mod_four with square
  intro three
  rw three at alternatives
  cases alternatives with zero | one
  . arith
  . arith

theorem Odd.square_mod_eight (k : Nat)
  requires odd : Odd(k)
  shows (k * k) % 8 = 1
proof
  obtain ⟨j, jNat, shape⟩ := Odd.iff_two_mul_add_one.forward with odd
  have evenProduct : Even(j * (j + 1))
    have parity := parity.exhaustive at j
    cases parity with even | oddJ
    . exact Even.mul_left at j, j + 1 with even
    . have succEven := Even.succ_of_odd at j with oddJ
      have one := add.right_one at j
      rw one
      exact Even.mul_right at j, succ(j) with succEven
  obtain ⟨m, mNat, mShape⟩ := Even.iff_two_mul.forward with evenProduct
  have expand : k * k = 8 * m + 1
    rw shape
    ring [mShape]
  rw expand
  have eight : 8 ≠ 0 by arith
  have reduce := mod.multiple_add at 8, 1, m with eight
  rw reduce
  compute

theorem IsSquare.odd_mod_eight (n : Nat)
  requires square : IsSquare(n)
  requires odd : Odd(n)
  shows n % 8 = 1
proof
  obtain ⟨k, kNat, shape⟩ := square
  rw shape at odd
  have both := Odd.of_mul at k, k with odd
  rw shape
  exact Odd.square_mod_eight with both.left

theorem IsSquare.even_iff (n : Nat)
  requires square : IsSquare(n)
  shows Even(n) ↔ ∃ k, Nat(k) ∧ Even(k) ∧ n = k * k
proof
  obtain ⟨k, kNat, shape⟩ := square
  constructor
  . intro even
    rw shape at even
    have characterization := Even.mul_self_iff at k
    have kEven := characterization.forward with even
    exact ⟨k, kNat, kEven, shape⟩
  . intro witnessed
    obtain ⟨j, jNat, jEven, jShape⟩ := witnessed
    rw jShape
    exact Even.mul_left at j, j with jEven

-- Strictly increasing squares: a square below another has the smaller root.
theorem IsSquare.root_lt (k m : Nat)
  requires below : k * k < m * m
  shows k < m
proof
  by_contra notLess
  have above := LessEqual.of_not_Less with notLess
  have squares := LessEqual.mul at m, k, m, k with above, above
  have cycle := Less.transitive_LessEqual with below, squares
  exact Less.irreflexive at k * k with cycle

theorem IsSquare.not_between (k n : Nat)
  requires lower : k * k < n
  requires upper : n < (k + 1) * (k + 1)
  shows ¬IsSquare(n)
proof
  intro square
  obtain ⟨m, mNat, shape⟩ := square
  rw shape at lower
  rw shape at upper
  have first := IsSquare.root_lt at k, m with lower
  have second := IsSquare.root_lt at m, k + 1 with upper
  arith

theorem IsSquare.not_two
  shows ¬IsSquare(2)
proof
  have lower : 1 * 1 < 2 by arith
  have square : (1 + 1) * (1 + 1) = 2 * 2 by compute
  have upper : 2 < (1 + 1) * (1 + 1)
    rw square
    arith
  exact IsSquare.not_between at 1, 2 with lower, upper

theorem IsSquare.not_three
  shows ¬IsSquare(3)
proof
  have lower : 1 * 1 < 3 by arith
  have square : (1 + 1) * (1 + 1) = 2 * 2 by compute
  have upper : 3 < (1 + 1) * (1 + 1)
    rw square
    arith
  exact IsSquare.not_between at 1, 3 with lower, upper

theorem IsSquare.not_five
  shows ¬IsSquare(5)
proof
  have lower : 2 * 2 < 5 by arith
  have square : (2 + 1) * (2 + 1) = 3 * 3 by compute
  have upper : 5 < (2 + 1) * (2 + 1)
    rw square
    arith
  exact IsSquare.not_between at 2, 5 with lower, upper

predicate SumOfTwoSquares (n : Nat) :=
  ∃ a, Nat(a) ∧ ∃ b, Nat(b) ∧ n = a * a + b * b

theorem SumOfTwoSquares.intro (a b : Nat)
  shows SumOfTwoSquares(a * a + b * b)
proof
  refine ⟨a, assumption, b, assumption, _⟩
  rfl

theorem SumOfTwoSquares.of_square (n : Nat)
  requires square : IsSquare(n)
  shows SumOfTwoSquares(n)
proof
  obtain ⟨k, kNat, shape⟩ := square
  have zero : k * k = k * k + 0 * 0 by ring
  rw [shape, zero]
  exact SumOfTwoSquares.intro at k, 0

theorem SumOfTwoSquares.zero
  shows SumOfTwoSquares(0)
proof
  have square := IsSquare.zero
  exact SumOfTwoSquares.of_square with square

theorem SumOfTwoSquares.one
  shows SumOfTwoSquares(1)
proof
  have square := IsSquare.one
  exact SumOfTwoSquares.of_square with square

theorem SumOfTwoSquares.two
  shows SumOfTwoSquares(2)
proof
  have shape : 2 = 1 * 1 + 1 * 1 by compute
  rw shape
  exact SumOfTwoSquares.intro at 1, 1

-- The Brahmagupta–Fibonacci identity, with the signed cross term replaced
-- by the distance between the two cross products.
theorem SumOfTwoSquares.product_identity (a b c d : Nat)
  shows
    (a * a + b * b) * (c * c + d * d) =
      (a * c + b * d) * (a * c + b * d) +
        natDist(a * d, b * c) * natDist(a * d, b * c)
proof
  have total := LessEqual.total at b * c, a * d
  cases total with below | above
  . obtain ⟨e, eNat, split⟩ := LessEqual.exists_add with below
    have characterization := natDist.eq_iff at a * d, b * c, e
    have distance : natDist(a * d, b * c) = e
      apply characterization.backward
      right
      exact split
    rw distance
    ring [split]
  . obtain ⟨e, eNat, split⟩ := LessEqual.exists_add with above
    have characterization := natDist.eq_iff at a * d, b * c, e
    have distance : natDist(a * d, b * c) = e
      apply characterization.backward
      left
      exact split
    rw distance
    ring [split]

theorem SumOfTwoSquares.mul (m n : Nat)
  requires first : SumOfTwoSquares(m)
  requires second : SumOfTwoSquares(n)
  shows SumOfTwoSquares(m * n)
proof
  obtain ⟨a, aNat, b, bNat, mShape⟩ := first
  obtain ⟨c, cNat, d, dNat, nShape⟩ := second
  have identity := SumOfTwoSquares.product_identity at a, b, c, d
  rw [mShape, nShape, identity]
  exact SumOfTwoSquares.intro at a * c + b * d, natDist(a * d, b * c)

theorem SumOfTwoSquares.mod_four (n : Nat)
  requires sum : SumOfTwoSquares(n)
  shows n % 4 ≠ 3
proof
  obtain ⟨a, aNat, b, bNat, shape⟩ := sum
  intro three
  rw shape at three
  have four : 4 ≠ 0 by arith
  have reduce := mod.add at 4, a * a, b * b with four
  rw reduce at three
  have squareA := IsSquare.intro at a
  have squareB := IsSquare.intro at b
  have leftCases := IsSquare.mod_four with squareA
  have rightCases := IsSquare.mod_four with squareB
  cases leftCases with leftZero | leftOne
  . rw leftZero at three
    cases rightCases with rightZero | rightOne
    . rw rightZero at three
      have value : (0 + 0) % 4 = 0 by compute
      rw value at three
      arith
    . rw rightOne at three
      have value : (0 + 1) % 4 = 1 by compute
      rw value at three
      arith
  . rw leftOne at three
    cases rightCases with rightZero | rightOne
    . rw rightZero at three
      have value : (1 + 0) % 4 = 1 by compute
      rw value at three
      arith
    . rw rightOne at three
      have value : (1 + 1) % 4 = 2 by compute
      rw value at three
      arith