-- Pythagorean triples: the parametrization (m² − n², 2mn, m² + n²) gives
-- every primitive triple with an even second leg, and every triple is a
-- multiple of a primitive one
-- 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.PythagoreanTriples
import CoprimeSquares

-- Positive naturals a, b, c with a² + b² = c².
predicate Pythagorean (a b c : Nat) := 0 < a ∧ 0 < b ∧ a ^ 2 + b ^ 2 = c ^ 2

-- A Pythagorean triple whose legs are coprime.
predicate PrimitivePythagorean (a b c : Nat) :=
  Pythagorean(a, b, c) ∧ gcd(a, b) = 1

-- Squaring is injective on the naturals.
theorem IsSquare.root_unique (x y : Nat)
  requires squares : x * x = y * y
  shows x = y
proof
  have order := Less.trichotomy at x, y with *
  rcases order with below | equal | above
  . exfalso
    have strict := Less.mul at x, y, x, y with below, below
    rw squares at strict
    exact Less.irreflexive at y * y with strict
  . exact equal
  . exfalso
    have strict := Less.mul at y, x, y, x with above, above
    rw squares at strict
    exact Less.irreflexive at y * y with strict

theorem Odd.square_mod_four (k : Nat)
  requires odd : Odd(k)
  shows (k * k) % 4 = 1
proof
  obtain ⟨j, jNat, shape⟩ := Odd.iff_two_mul_add_one.forward with odd
  have expand : k * k = 4 * (j * j + j) + 1
    rw shape
    ring
  rw expand
  have four : 4 ≠ 0 by arith
  have reduce := mod.multiple_add at 4, 1, j * j + j with four
  rw reduce
  compute

-- An odd divisor of 2t divides t.
theorem divides.of_double_odd (d t : Nat)
  requires odd : Odd(d)
  requires double : d ∣ 2 * t
  shows d ∣ t
proof
  obtain ⟨k, kNat, product⟩ := double
  have evenK : Even(k)
    have evenProduct : Even(d * k)
      rw product
      exact Even.two_mul at t
    have split := Even.of_mul at d, k with evenProduct
    cases split with evenD | evenK
    . exfalso
      exact Odd.not_even with odd, evenD
    . exact evenK
  obtain ⟨j, jNat, kShape⟩ := Even.iff_two_mul.forward with evenK
  witness j
  have scaled : 2 * (d * j) = 2 * t
    rw [← product, kShape]
    ring
  have two : 2 ≠ 0 by arith
  exact mul.left_cancel_positive at d * j, t, 2 with two, scaled

-- With m² = n² + d, the difference d has the parity of m + n.
theorem Even.root_sum_iff (m n d : Nat)
  requires shape : m * m = n * n + d
  shows Even(m + n) ↔ Even(d)
proof
  have squareM := Even.mul_self_iff at m
  have squareN := Even.mul_self_iff at n
  have sumMN := Even.add_iff at m, n
  have sumND := Even.add_iff at n * n, d
  rw shape at squareM
  tauto

-- The parametrization: for m > n > 0 the triple
-- (m² − n², 2mn, m² + n²) is Pythagorean.
theorem pythagorean.of_parameters (m n : Nat)
  requires positive : 0 < n
  requires below : n < m
  shows Pythagorean(m ^ 2 - n ^ 2, 2 * m * n, m ^ 2 + n ^ 2)
proof
  have squares := Less.mul at n, m, n, m with below, below
  have squaresLe : n * n ≤ m * m by arith
  obtain ⟨d, dNat, shape⟩ := LessEqual.exists_add at n * n, m * m
    with squaresLe
  have difference : m ^ 2 - n ^ 2 = d
    rw [pow.two, pow.two, shape]
    exact sub.add_left at d, n * n
  have dPositive : 0 < d by arith
  have productPositive : 0 < 2 * m * n
    have mNonzero : m ≠ 0 by arith
    have nNonzero : n ≠ 0 by arith
    have two : 2 ≠ 0 by arith
    have twoM := mul.nonzero at 2, m with two, mNonzero
    have all := mul.nonzero at 2 * m, n with twoM, nNonzero
    exact Less.positive_of_ne_zero at 2 * m * n with all
  unfold Pythagorean
  rw difference
  refine ⟨dPositive, productPositive, _⟩
  simp [pow.two]
  ring [shape]

-- With m, n coprime and of opposite parity, the parametrized triple is
-- primitive: a common divisor e of m² − n² and 2mn is odd, hence divides
-- mn, and is coprime to m and to n, hence to mn.
theorem pythagorean.primitive_of_parameters (m n : Nat)
  requires positive : 0 < n
  requires below : n < m
  requires coprime : gcd(m, n) = 1
  requires parity : Odd(m + n)
  shows PrimitivePythagorean(m ^ 2 - n ^ 2, 2 * m * n, m ^ 2 + n ^ 2)
proof
  constructor
  . exact pythagorean.of_parameters at m, n with positive, below
  . have squares := Less.mul at n, m, n, m with below, below
    have squaresLe : n * n ≤ m * m by arith
    obtain ⟨d, dNat, shape⟩ := LessEqual.exists_add at n * n, m * m
      with squaresLe
    have difference : m ^ 2 - n ^ 2 = d
      rw [pow.two, pow.two, shape]
      exact sub.add_left at d, n * n
    rw difference
    have oddD : Odd(d)
      have parities := Even.root_sum_iff at m, n, d with shape
      have notEvenSum := Odd.not_even at m + n with parity
      have oddIff := Odd.iff_not_even at d
      apply oddIff.backward
      tauto
    have every := gcd.eq_one_iff at d, 2 * m * n
    apply every.backward
    intros e eNat ed eProduct
    have oddE : Odd(e)
      obtain ⟨k, kNat, product⟩ := ed
      rw ← product at oddD
      have both := Odd.of_mul at e, k with oddD
      exact both.left
    have regroup : 2 * m * n = 2 * (m * n) by ring
    rw regroup at eProduct
    have eDividesMN := divides.of_double_odd at e, m * n with oddE, eProduct
    have coprimeEM : gcd(e, m) = 1
      have everyM := gcd.eq_one_iff at e, m
      apply everyM.backward
      intros f fNat fe fm
      have fd := divides.transitive at f, e, d with fe, ed
      have fmm := divides.mul_self at f, m with fm
      have fnn : f ∣ n * n
        have cancel := sub.add_right at n * n, d
        rw [← cancel, ← shape]
        exact divides.sub at f, m * m, d with fmm, fd
      have coprimeFN := coprime.of_divides_left at m, n, f with coprime, fm
      have coprimeFNN := coprime.pow_right at f, n, 2 with coprimeFN
      rw pow.two at coprimeFNN
      have fSelf := divides.reflexive at f
      exact coprime.divisor_one at f, n * n, f with coprimeFNN, fSelf, fnn
    have coprimeEN : gcd(e, n) = 1
      have everyN := gcd.eq_one_iff at e, n
      apply everyN.backward
      intros f fNat fe fn
      have fd := divides.transitive at f, e, d with fe, ed
      have fnn := divides.mul_self at f, n with fn
      have fmm : f ∣ m * m
        rw shape
        exact divides.add at f, n * n, d with fnn, fd
      have coprimeNM := coprime.symmetric at m, n with coprime
      have coprimeFM := coprime.of_divides_left at n, m, f with coprimeNM, fn
      have coprimeFMM := coprime.pow_right at f, m, 2 with coprimeFM
      rw pow.two at coprimeFMM
      have fSelf := divides.reflexive at f
      exact coprime.divisor_one at f, m * m, f with coprimeFMM, fSelf, fmm
    have coprimeME := coprime.symmetric at e, m with coprimeEM
    have coprimeNE := coprime.symmetric at e, n with coprimeEN
    have coprimeMNE := coprime.mul_of_coprime at m, n, e
      with coprimeME, coprimeNE
    have coprimeEMN := coprime.symmetric at m * n, e with coprimeMNE
    have eSelf := divides.reflexive at e
    exact coprime.divisor_one at e, m * n, e with coprimeEMN, eSelf, eDividesMN

-- In a primitive triple exactly one leg is even: two even legs share the
-- factor 2, and two odd legs give c² ≡ 2 modulo 4.
theorem pythagorean.one_even (a b c : Nat)
  requires primitive : PrimitivePythagorean(a, b, c)
  shows Odd(a + b)
proof
  obtain ⟨triple, coprime⟩ := primitive
  obtain ⟨positiveA, positiveB, equation⟩ := triple
  simp [pow.two] at equation
  have notBothEven : ¬(Even(a) ∧ Even(b))
    intro both
    have twoIffA := divides.two_iff_even at a
    have twoIffB := divides.two_iff_even at b
    have twoA := twoIffA.backward with both.left
    have twoB := twoIffB.backward with both.right
    have two := coprime.divisor_one at a, b, 2 with coprime, twoA, twoB
    arith
  have notBothOdd : ¬(Odd(a) ∧ Odd(b))
    intro both
    have modA := Odd.square_mod_four at a with both.left
    have modB := Odd.square_mod_four at b with both.right
    have four : 4 ≠ 0 by arith
    have reduce := mod.add at 4, a * a, b * b with four
    rw [modA, modB, equation] at reduce
    have value : (1 + 1) % 4 = 2 by compute
    rw value at reduce
    have square := IsSquare.intro at c
    have notTwo := IsSquare.mod_four_ne_two at c * c with square
    exact notTwo with reduce
  have evenIffA := Even.iff_not_odd at a
  have evenIffB := Even.iff_not_odd at b
  have sum := Even.add_iff at a, b
  have oddIff := Odd.iff_not_even at a + b
  apply oddIff.backward
  tauto

-- Every primitive triple with b even comes from the parametrization: with
-- c = a + 2u and b = 2w, the equation gives w² = u(a + u) with u and a + u
-- coprime, so u = n² and a + u = m².
theorem pythagorean.primitive_form (a b c : Nat)
  requires primitive : PrimitivePythagorean(a, b, c)
  requires even : Even(b)
  shows ∃ m, Nat(m) ∧ ∃ n, Nat(n) ∧
    0 < n ∧ n < m ∧ gcd(m, n) = 1 ∧ Odd(m + n) ∧
    a = m ^ 2 - n ^ 2 ∧ b = 2 * m * n ∧ c = m ^ 2 + n ^ 2
proof
  have parity := pythagorean.one_even at a, b, c with primitive
  obtain ⟨triple, coprime⟩ := primitive
  obtain ⟨positiveA, positiveB, equation⟩ := triple
  simp [pow.two] at equation
  have oddA : Odd(a)
    have swapped : Odd(b + a)
      have commute : b + a = a + b by ring
      rw commute
      exact parity
    exact Odd.of_add_odd_even at b, a with swapped, even
  have oddC : Odd(c)
    have squareA := Odd.mul_self_iff at a
    have squareB := Even.mul_self_iff at b
    have oddSquareA := squareA.backward with oddA
    have evenSquareB := squareB.backward with even
    have oddSum := Odd.add_even at a * a, b * b with oddSquareA, evenSquareB
    rw equation at oddSum
    have squareC := Odd.mul_self_iff at c
    exact squareC.forward with oddSum
  have aBelowC : a < c
    have bNonzero : b ≠ 0 by arith
    have bbNonzero := mul.nonzero at b, b with bNonzero, bNonzero
    have bbPositive := Less.positive_of_ne_zero at b * b with bbNonzero
    have squares : a * a < c * c by arith
    exact IsSquare.root_lt at a, c with squares
  have aLeC : a ≤ c by arith
  obtain ⟨t, tNat, cShape⟩ := LessEqual.exists_add at a, c with aLeC
  have evenT : Even(t)
    rw cShape at oddC
    exact Even.of_add_odd_odd at a, t with oddC, oddA
  obtain ⟨u, uNat, tShape⟩ := Even.iff_two_mul.forward with evenT
  obtain ⟨w, wNat, bShape⟩ := Even.iff_two_mul.forward with even
  have key : u * (a + u) = w * w
    have expanded : 4 * (u * (a + u)) = 4 * (w * w)
      rw [bShape, cShape, tShape] at equation
      have leftSide : a * a + 2 * w * (2 * w) = a * a + 4 * (w * w) by ring
      have rightSide : (a + 2 * u) * (a + 2 * u) = a * a + 4 * (u * (a + u))
        ring
      rw [leftSide, rightSide] at equation
      have cancel := add.left_cancel at 4 * (w * w), 4 * (u * (a + u)), a * a
      have swapped := cancel with equation
      symm
      exact swapped
    have four : 4 ≠ 0 by arith
    exact mul.left_cancel_positive at u * (a + u), w * w, 4 with four, expanded
  have uPositive : 0 < u by arith
  have coprimeUV : gcd(u, a + u) = 1
    have every := gcd.eq_one_iff at u, a + u
    apply every.backward
    intros d dNat du dv
    have da := divides.sub at d, a + u, u with dv, du
    have cancel := sub.add_right at a, u
    rw cancel at da
    have dSquares := divides.mul_mono at d, u, d, a + u with du, dv
    rw key at dSquares
    have dw := divides.of_mul_self at d, w with dSquares
    have db : d ∣ b
      rw bShape
      exact divides.multiply_left at d, w, 2 with *
    exact coprime.divisor_one at a, b, d with coprime, da, db
  obtain ⟨squareU, squareV⟩ := IsSquare.of_coprime_mul at u, a + u, w
    with coprimeUV, key
  obtain ⟨n, nNat, uShape⟩ := squareU
  obtain ⟨m, mNat, vShape⟩ := squareV
  have nPositive : 0 < n
    have nNonzero : n ≠ 0
      intro nZero
      rw nZero at uShape
      have zero : 0 * 0 = 0 by compute
      rw zero at uShape
      arith
    exact Less.positive_of_ne_zero at n with nNonzero
  have nBelowM : n < m
    have squares : n * n < m * m
      rw [← uShape, ← vShape]
      arith
    exact IsSquare.root_lt at n, m with squares
  have coprimeMN : gcd(m, n) = 1
    have coprimeVU := coprime.symmetric at u, a + u with coprimeUV
    rw [vShape, uShape] at coprimeVU
    have coprimeMNN := coprime.of_mul_left at m, m, n * n with coprimeVU
    have coprimeNNM := coprime.symmetric at m, n * n with coprimeMNN
    have coprimeNM := coprime.of_mul_left at n, n, m with coprimeNNM
    exact coprime.symmetric at n, m with coprimeNM
  have parityMN : Odd(m + n)
    have shape : m * m = n * n + a
      rw [← vShape, ← uShape]
      ring
    have parities := Even.root_sum_iff at m, n, a with shape
    have notEvenA := Odd.not_even at a with oddA
    have oddIff := Odd.iff_not_even at m + n
    apply oddIff.backward
    tauto
  have aForm : a = m ^ 2 - n ^ 2
    rw [pow.two, pow.two, ← vShape, ← uShape]
    symm
    exact sub.add_right at a, u
  have wForm : w = m * n
    have squares : w * w = m * n * (m * n)
      rw [← key, vShape, uShape]
      ring
    exact IsSquare.root_unique at w, m * n with squares
  have bForm : b = 2 * m * n
    rw [bShape, wForm]
    ring
  have cForm : c = m ^ 2 + n ^ 2
    rw [pow.two, pow.two, ← vShape, ← uShape, cShape, tShape]
    ring
  exact ⟨m, mNat, n, nNat, nPositive, nBelowM, coprimeMN, parityMN, aForm,
    bForm, cForm⟩

-- Every Pythagorean triple is gcd(a, b) times a primitive one.
theorem pythagorean.form (a b c : Nat)
  requires triple : Pythagorean(a, b, c)
  shows ∃ x, Nat(x) ∧ ∃ y, Nat(y) ∧ ∃ z, Nat(z) ∧
    PrimitivePythagorean(x, y, z) ∧
    a = gcd(a, b) * x ∧ b = gcd(a, b) * y ∧ c = gcd(a, b) * z
proof
  obtain ⟨positiveA, positiveB, equation⟩ := triple
  simp [pow.two] at equation
  have aNonzero : a ≠ 0 by arith
  have gcdNonzero := gcd.nonzero at a, b with aNonzero
  obtain ⟨x, xNat, aShape⟩ := gcd.divides_left at a, b
  obtain ⟨y, yNat, bShape⟩ := gcd.divides_right at a, b
  have cofactors := gcd.cofactor_coprime at a, b, x, y
    with gcdNonzero, aShape, bShape
  have squareDividesC : gcd(a, b) * gcd(a, b) ∣ c * c
    have dividesA := gcd.divides_left at a, b
    have dividesB := gcd.divides_right at a, b
    have left := divides.mul_mono at gcd(a, b), a, gcd(a, b), a
      with dividesA, dividesA
    have right := divides.mul_mono at gcd(a, b), b, gcd(a, b), b
      with dividesB, dividesB
    rw ← equation
    exact divides.add at gcd(a, b) * gcd(a, b), a * a, b * b with left, right
  have dividesC := divides.of_mul_self at gcd(a, b), c with squareDividesC
  obtain ⟨z, zNat, cShape⟩ := dividesC
  have reduced : x * x + y * y = z * z
    have scaled : gcd(a, b) * gcd(a, b) * (x * x + y * y) =
        gcd(a, b) * gcd(a, b) * (z * z)
      have leftSide : gcd(a, b) * gcd(a, b) * (x * x + y * y) = a * a + b * b
        ring [aShape, bShape]
      have rightSide : gcd(a, b) * gcd(a, b) * (z * z) = c * c by ring [cShape]
      rw [leftSide, rightSide]
      exact equation
    have squareNonzero := mul.nonzero at gcd(a, b), gcd(a, b)
      with gcdNonzero, gcdNonzero
    exact mul.left_cancel_positive at x * x + y * y, z * z,
      gcd(a, b) * gcd(a, b) with squareNonzero, scaled
  have positiveX : 0 < x
    have nonzero : x ≠ 0
      intro zero
      rw zero at aShape
      have value : gcd(a, b) * 0 = 0 by ring
      rw value at aShape
      arith
    exact Less.positive_of_ne_zero at x with nonzero
  have positiveY : 0 < y
    have nonzero : y ≠ 0
      intro zero
      rw zero at bShape
      have value : gcd(a, b) * 0 = 0 by ring
      rw value at bShape
      arith
    exact Less.positive_of_ne_zero at y with nonzero
  have primitive : PrimitivePythagorean(x, y, z)
    constructor
    . refine ⟨positiveX, positiveY, _⟩
      rw [pow.two, pow.two, pow.two]
      exact reduced
    . exact cofactors
  refine ⟨x, xNat, y, yNat, z, zNat, primitive, _, _, _⟩
  . symm
    exact aShape
  . symm
    exact bShape
  . symm
    exact cShape