-- The modular inverse of a unit mod a prime, chosen once: it is a unit, it
-- is an involution, and its only fixed points are 1 and p − 1
-- 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.InverseInvolution
import Euler

defunique modInverse (p x : Nat) (b ∈ p) := mul(x, b) % p = 1
requires prime : Prime(p)
requires unit : gcd(x, p) = 1
proof
  have aboveOne := Prime.greater_than_one at p with prime
  obtain ⟨b, bNat, bBelow, inverse⟩ := coprime.inverse at x, p with aboveOne, unit
  witness b
  constructor
  . exact membership.of_Less at b, p with *, *, bBelow
  . exact inverse

uniqueness
  intros u uMember first v vMember second
  have nonzero := Prime.nonzero at p with prime
  have uNat := natural.member at p with uMember
  have vNat := natural.member at p with vMember
  have uBelow := Less.of_membership at u, p
  have vBelow := Less.of_membership at v, p
  have same : (x * u) % p = (x * v) % p
    rw [first, second]
    rfl
  have residues := coprime.mul_cancel_mod at x, p, u, v with nonzero, unit, same
  have uSelf := mod.of_less at p, u with (uBelow with uMember)
  have vSelf := mod.of_less at p, v with (vBelow with vMember)
  rw [uSelf, vSelf] at residues
  exact residues

theorem modInverse.below (p x : Nat)
  requires prime : Prime(p)
  requires unit : gcd(x, p) = 1
  shows modInverse(p, x) < p
proof
  have selected := modInverse.selected with prime, unit
  have bNat := natural.member at p with selected.left
  have below := Less.of_membership at modInverse(p, x), p
  exact below with selected.left

theorem modInverse.natural (p x : Nat)
  requires prime : Prime(p)
  requires unit : gcd(x, p) = 1
  shows Nat(modInverse(p, x))
proof
  have selected := modInverse.selected with prime, unit
  exact natural.member at p with selected.left

theorem modInverse.product (p x : Nat)
  requires prime : Prime(p)
  requires unit : gcd(x, p) = 1
  shows mul(x, modInverse(p, x)) % p = 1
proof
  have selected := modInverse.selected with prime, unit
  exact selected.right

theorem modInverse.unit (p x : Nat)
  requires prime : Prime(p)
  requires unit : gcd(x, p) = 1
  shows gcd(modInverse(p, x), p) = 1
proof
  have nonzero := Prime.nonzero with prime
  have bNat := modInverse.natural with prime, unit
  have product := modInverse.product with prime, unit
  exact units.inverse_coprime with nonzero, product

-- Inverses are unique among residues below p.
theorem modInverse.unique_below (p x b : Nat)
  requires prime : Prime(p)
  requires unit : gcd(x, p) = 1
  requires below : b < p
  requires product : (x * b) % p = 1
  shows b = modInverse(p, x)
proof
  have nonzero := Prime.nonzero with prime
  have inverseNat := modInverse.natural with prime, unit
  have inverseBelow := modInverse.below with prime, unit
  have inverseProduct := modInverse.product with prime, unit
  have same : (x * b) % p = (x * modInverse(p, x)) % p
    rw [product, inverseProduct]
    rfl
  have residues := coprime.mul_cancel_mod with nonzero, unit, same
  have bSelf := mod.of_less with below
  have inverseSelf := mod.of_less with inverseBelow
  rwa [bSelf, inverseSelf] at residues

theorem modInverse.involution (p x : Nat)
  requires prime : Prime(p)
  requires unit : gcd(x, p) = 1
  requires below : x < p
  shows modInverse(p, modInverse(p, x)) = x
proof
  have inverseNat := modInverse.natural with prime, unit
  have inverseUnit := modInverse.unit with prime, unit
  have product := modInverse.product with prime, unit
  have commute := mul.commutative at modInverse(p, x), x with *
  rw commute at product
  have flipped : mul(modInverse(p, x), x) % p = 1 by exact product
  have same := modInverse.unique_below with prime, inverseUnit, below, flipped
  symm at same
  exact same

theorem modInverse.positive (p x : Nat)
  requires prime : Prime(p)
  requires unit : gcd(x, p) = 1
  shows 0 < modInverse(p, x)
proof
  have aboveOne := Prime.greater_than_one with prime
  have inverseNat := modInverse.natural with prime, unit
  have product := modInverse.product with prime, unit
  have trichotomy := Less.trichotomy at 0, modInverse(p, x)
  rcases trichotomy with positive | zero | negative
  . exact positive
  . exfalso
    rw ← zero at product
    have vanish : x * 0 = 0 by ring
    rw [vanish, mod.zero_left] at product
    have impossible : ¬(0:N) = 1 by arith
    exact impossible with product
  . exfalso
    have impossible : ¬modInverse(p, x) < 0 by arith
    exact impossible with negative

-- A residue below p is its own inverse only when it is 1 or p − 1.
theorem modInverse.fixed (p x : Nat)
  requires prime : Prime(p)
  requires positive : 0 < x
  requires below : x < p
  requires fixed : modInverse(p, x) = x
  shows x = 1 ∨ succ(x) = p
proof
  have nonzero := Prime.nonzero with prime
  have unit := prime.coprime_below with prime, positive, below
  have product := modInverse.product with prime, unit
  rw fixed at product
  have shape := natural.zero_or_successor at x
  cases shape with zero | successor
  . exfalso
    rw zero at positive
    exact Less.irreflexive with *, positive
  . obtain ⟨y, yNat, form⟩ := successor
    obtain ⟨q, qNat, equation⟩ := mod.equation at x * x, p with nonzero
    rw [product, form] at equation
    have expand : succ(y) * succ(y) = y * (y + 2) + 1 by ring
    rw expand at equation
    have cancel := add.right_cancel at y * (y + 2), p * q, 1
    have reduced := cancel with equation
    have dividesProduct : p ∣ y * (y + 2)
      witness q
      rw reduced
      rfl
    have split := Prime.divides_mul with prime, dividesProduct
    cases split with dividesY | dividesNext
    . left
      have yShape := natural.zero_or_successor at y
      cases yShape with yZero | ySucc
      . rw [form, yZero]
        rfl
      . exfalso
        have yNonzero : y ≠ 0
          intro zero
          obtain ⟨z, zNat, yForm⟩ := ySucc
          rw yForm at zero
          have impossible : succ(z) ≠ 0 by arith
          exact impossible with zero
        have le := divides.le with dividesY, yNonzero
        rw form at below
        have impossible : ¬p ≤ y by arith
        exact impossible with le
    . right
      have nextNonzero : y + 2 ≠ 0 by arith
      have le := divides.le with dividesNext, nextNonzero
      rw form at below
      have upper : y + 2 ≤ p
        have unfold : y + 2 = succ(succ(y)) by ring
        rw unfold
        exact below
      have same := LessEqual.antisymmetric with upper, le
      rw form
      have unfold : succ(succ(y)) = y + 2 by ring
      rwa unfold