-- More modular arithmetic: remainders of multiples and sums, powers, the
-- exhaustive small remainders, congruence modulo n as a predicate with its
-- equivalence and compatibility laws, cancellation of a coprime factor, and
-- the congruence forms of Euler's and Fermat's theorems
-- 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.ModularMore
import Euler
import DivisibilityMore

theorem mod.self (n : Nat)
  shows n % n = 0
proof
  have reflexive := divides.reflexive at n
  have characterization := divides.iff_mod_zero at n, n
  exact characterization.forward with reflexive

theorem mod.mul_self_left (n k : Nat)
  shows (n * k) % n = 0
proof
  have multiple := divides.mul_left_self at n, k
  have characterization := divides.iff_mod_zero at n, n * k
  exact characterization.forward with multiple

theorem mod.mul_self_right (n k : Nat)
  shows (k * n) % n = 0
proof
  have multiple := divides.mul_right_self at n, k
  have characterization := divides.iff_mod_zero at n, k * n
  exact characterization.forward with multiple

theorem mod.of_divides (d a : Nat)
  requires hda : d ∣ a
  shows a % d = 0
proof
  have characterization := divides.iff_mod_zero at d, a
  exact characterization.forward with hda

theorem mod.add_self (n a : Nat)
  requires nonzero : n ≠ 0
  shows (a + n) % n = a % n
proof
  have shifted := mod.add_multiple at n, a, 1 with nonzero
  have one : a + n * 1 = a + n by ring
  rw one at shifted
  exact shifted

theorem mod.self_add (n a : Nat)
  requires nonzero : n ≠ 0
  shows (n + a) % n = a % n
proof
  have commute : n + a = a + n by ring
  rw commute
  exact mod.add_self at n, a with nonzero

theorem mod.multiple_add (n a s : Nat)
  requires nonzero : n ≠ 0
  shows (n * s + a) % n = a % n
proof
  have commute : n * s + a = a + n * s by ring
  rw commute
  exact mod.add_multiple at n, a, s with nonzero

theorem mod.le (n a : Nat)
  shows a % n ≤ a
proof
  have split : n = 0 ∨ n ≠ 0 by tauto
  cases split with zero | nonzero
  . rw zero
    have value := mod.zero_right at a
    rw value
    exact LessEqual.reflexive at a
  . obtain ⟨q, qNat, equation⟩ := mod.equation at a, n with nonzero
    witness n * q
    symm
    linear_combination equation

theorem mod.eq_self_iff (n a : Nat)
  requires nonzero : n ≠ 0
  shows a % n = a ↔ a < n
proof
  constructor
  . intro same
    have bound := mod.less at a, n with nonzero
    rw same at bound
    exact bound
  . intro below
    exact mod.of_less with below

theorem mod.one_left (n : Nat)
  requires above : 1 < n
  shows 1 % n = 1
proof
  exact mod.of_less with above

theorem mod.add_mod_left (n a b : Nat)
  requires nonzero : n ≠ 0
  shows (a % n + b) % n = (a + b) % n
proof
  have first := mod.add at n, a % n, b with nonzero
  have second := mod.add at n, a, b with nonzero
  have idempotent := mod.mod at n, a with nonzero
  rw [first, idempotent, second]
  rfl

theorem mod.add_mod_right (n a b : Nat)
  requires nonzero : n ≠ 0
  shows (a + b % n) % n = (a + b) % n
proof
  have first := mod.add at n, a, b % n with nonzero
  have second := mod.add at n, a, b with nonzero
  have idempotent := mod.mod at n, b with nonzero
  rw [first, idempotent, second]
  rfl

theorem mod.pow (n a k : Nat)
  requires nonzero : n ≠ 0
  shows (a ^ k) % n = ((a % n) ^ k) % n
proof
  induction k
  | zero =>
    have left := pow.zero at a
    have right := pow.zero at a % n
    rw [left, right]
    rfl
  | j + 1 using ih =>
    have left := pow.add_one at a, j
    have right := pow.add_one at a % n, j
    rw [left, right]
    have first := mod.mul at n, a, a ^ j with nonzero
    rw [first, ih]
    exact mod.mul_right at n, a % n, (a % n) ^ j with nonzero

theorem natural.below_two (r : Nat)
  requires below : r < 2
  shows r = 0 ∨ r = 1
proof
  have le1 : r ≤ 1 by arith
  have alternatives := LessEqual.cases at r, 1 with le1
  cases alternatives with lt1 | eq1
  . left
    have le0 : r ≤ 0 by arith
    exact LessEqual.eq_zero with le0
  . right
    exact eq1

theorem natural.below_three (r : Nat)
  requires below : r < 3
  shows r = 0 ∨ r = 1 ∨ r = 2
proof
  have le2 : r ≤ 2 by arith
  have alternatives := LessEqual.cases at r, 2 with le2
  cases alternatives with lt2 | eq2
  . have small := natural.below_two with lt2
    cases small with zero | one
    . left
      exact zero
    . right
      left
      exact one
  . right
    right
    exact eq2

theorem natural.below_four (r : Nat)
  requires below : r < 4
  shows r = 0 ∨ r = 1 ∨ r = 2 ∨ r = 3
proof
  have le3 : r ≤ 3 by arith
  have alternatives := LessEqual.cases at r, 3 with le3
  cases alternatives with lt3 | eq3
  . have small := natural.below_three with lt3
    rcases small with zero | one | two
    . left
      exact zero
    . right
      left
      exact one
    . right
      right
      left
      exact two
  . right
    right
    right
    exact eq3

theorem mod.two_cases (a : Nat)
  shows a % 2 = 0 ∨ a % 2 = 1
proof
  have two : 2 ≠ 0 by arith
  have bound := mod.less at a, 2 with two
  exact natural.below_two with bound

theorem mod.three_cases (a : Nat)
  shows a % 3 = 0 ∨ a % 3 = 1 ∨ a % 3 = 2
proof
  have three : 3 ≠ 0 by arith
  have bound := mod.less at a, 3 with three
  exact natural.below_three with bound

theorem mod.four_cases (a : Nat)
  shows a % 4 = 0 ∨ a % 4 = 1 ∨ a % 4 = 2 ∨ a % 4 = 3
proof
  have four : 4 ≠ 0 by arith
  have bound := mod.less at a, 4 with four
  exact natural.below_four with bound

-- Adding d preserves the remainder exactly when n divides d.
theorem mod.add_eq_iff (n a d : Nat)
  shows (a + d) % n = a % n ↔ n ∣ d
proof
  have split : n = 0 ∨ n ≠ 0 by tauto
  cases split with zero | nonzero
  . rw zero
    have left := mod.zero_right at a + d
    have right := mod.zero_right at a
    rw [left, right]
    constructor
    . intro same
      have le0 : d ≤ 0 by arith
      have dZero := LessEqual.eq_zero with le0
      rw dZero
      exact divides.zero at 0
    . intro h
      have dZero := zero.divides_only_zero with h
      rw dZero
      ring
  . constructor
    . intro same
      exact mod.add_cancel at n, a, d with nonzero, same
    . intro h
      obtain ⟨s, sNat, equation⟩ := h
      rw ← equation
      exact mod.add_multiple at n, a, s with nonzero

theorem mod.eq_iff_divides_natDist (n a b : Nat)
  shows a % n = b % n ↔ n ∣ natDist(a, b)
proof
  have total := LessEqual.total at a, b
  cases total with below | above
  . obtain ⟨d, dNat, split⟩ := LessEqual.exists_add with below
    have distance : natDist(a, b) = d
      have value := natDist.of_le at a, b with below
      have sum := split
      symm at sum
      have difference := sub.of_add_eq_left at a, d, b with sum
      rw [value, difference]
      rfl
    rw [distance, split]
    have shifted := mod.add_eq_iff at n, a, d
    constructor
    . intro same
      symm at same
      exact shifted.forward with same
    . intro h
      symm
      exact shifted.backward with h
  . obtain ⟨d, dNat, split⟩ := LessEqual.exists_add with above
    have distance : natDist(a, b) = d
      have value := natDist.of_ge at a, b with above
      have sum := split
      symm at sum
      have difference := sub.of_add_eq_left at b, d, a with sum
      rw [value, difference]
      rfl
    rw [distance, split]
    exact mod.add_eq_iff at n, b, d

predicate CongruentMod (a b n : Nat) := a % n = b % n

theorem CongruentMod.reflexive (a n : Nat)
  shows CongruentMod(a, a, n)
proof
  unfold CongruentMod
  rfl

theorem CongruentMod.symmetric (a b n : Nat)
  requires congruent : CongruentMod(a, b, n)
  shows CongruentMod(b, a, n)
proof
  unfold CongruentMod at congruent
  unfold CongruentMod
  symm
  exact congruent

theorem CongruentMod.transitive (a b c n : Nat)
  requires first : CongruentMod(a, b, n)
  requires second : CongruentMod(b, c, n)
  shows CongruentMod(a, c, n)
proof
  unfold CongruentMod at first
  unfold CongruentMod at second
  unfold CongruentMod
  rw [first, second]
  rfl

theorem CongruentMod.of_eq (a b n : Nat)
  requires same : a = b
  shows CongruentMod(a, b, n)
proof
  rw same
  exact CongruentMod.reflexive at b, n

theorem CongruentMod.zero_modulus (a b : Nat)
  shows CongruentMod(a, b, 0) ↔ a = b
proof
  have left := mod.zero_right at a
  have right := mod.zero_right at b
  constructor
  . intro congruent
    unfold CongruentMod at congruent
    rw [left, right] at congruent
    exact congruent
  . intro same
    unfold CongruentMod
    rw [left, right]
    exact same

theorem CongruentMod.one_modulus (a b : Nat)
  shows CongruentMod(a, b, 1)
proof
  unfold CongruentMod
  have left := mod.one at a
  have right := mod.one at b
  rw [left, right]
  rfl

theorem CongruentMod.iff_divides_natDist (a b n : Nat)
  shows CongruentMod(a, b, n) ↔ n ∣ natDist(a, b)
proof
  have characterization := mod.eq_iff_divides_natDist at n, a, b
  constructor
  . intro congruent
    unfold CongruentMod at congruent
    exact characterization.forward with congruent
  . intro distance
    unfold CongruentMod
    exact characterization.backward with distance

theorem CongruentMod.divides_natDist (a b n : Nat)
  requires congruent : CongruentMod(a, b, n)
  shows n ∣ natDist(a, b)
proof
  have characterization := CongruentMod.iff_divides_natDist at a, b, n
  exact characterization.forward with congruent

theorem CongruentMod.of_divides_natDist (a b n : Nat)
  requires distance : n ∣ natDist(a, b)
  shows CongruentMod(a, b, n)
proof
  have characterization := CongruentMod.iff_divides_natDist at a, b, n
  exact characterization.backward with distance

theorem CongruentMod.of_add (a d n : Nat)
  requires multiple : n ∣ d
  shows CongruentMod(a + d, a, n)
proof
  unfold CongruentMod
  have shifted := mod.add_eq_iff at n, a, d
  exact shifted.backward with multiple

theorem CongruentMod.zero_right_iff (a n : Nat)
  shows CongruentMod(a, 0, n) ↔ n ∣ a
proof
  have characterization := CongruentMod.iff_divides_natDist at a, 0, n
  have value := natDist.zero_right at a
  rw value at characterization
  exact characterization

theorem CongruentMod.mod_right (a n : Nat)
  shows CongruentMod(a, a % n, n)
proof
  unfold CongruentMod
  have split : n = 0 ∨ n ≠ 0 by tauto
  cases split with zero | nonzero
  . rw zero
    have value := mod.zero_right at a
    rw value
    rw value
    rfl
  . have idempotent := mod.mod at n, a with nonzero
    rw idempotent
    rfl

theorem CongruentMod.mod_left (a n : Nat)
  shows CongruentMod(a % n, a, n)
proof
  have right := CongruentMod.mod_right at a, n
  exact CongruentMod.symmetric with right

theorem CongruentMod.divides_modulus (a b m n : Nat)
  requires hmn : m ∣ n
  requires congruent : CongruentMod(a, b, n)
  shows CongruentMod(a, b, m)
proof
  have distance := CongruentMod.divides_natDist with congruent
  have smaller := divides.transitive at m, n, natDist(a, b) with hmn, distance
  exact CongruentMod.of_divides_natDist with smaller

theorem CongruentMod.add (a b c d n : Nat)
  requires first : CongruentMod(a, b, n)
  requires second : CongruentMod(c, d, n)
  shows CongruentMod(a + c, b + d, n)
proof
  have split : n = 0 ∨ n ≠ 0 by tauto
  cases split with zero | nonzero
  . rw zero at first
    rw zero at second
    rw zero
    have firstEq := CongruentMod.zero_modulus at a, b
    have secondEq := CongruentMod.zero_modulus at c, d
    have ab := firstEq.forward with first
    have cd := secondEq.forward with second
    rw [ab, cd]
    exact CongruentMod.reflexive at b + d, 0
  . unfold CongruentMod at first
    unfold CongruentMod at second
    unfold CongruentMod
    have left := mod.add at n, a, c with nonzero
    have right := mod.add at n, b, d with nonzero
    rw [left, right, first, second]
    rfl

theorem CongruentMod.mul (a b c d n : Nat)
  requires first : CongruentMod(a, b, n)
  requires second : CongruentMod(c, d, n)
  shows CongruentMod(a * c, b * d, n)
proof
  have split : n = 0 ∨ n ≠ 0 by tauto
  cases split with zero | nonzero
  . rw zero at first
    rw zero at second
    rw zero
    have firstEq := CongruentMod.zero_modulus at a, b
    have secondEq := CongruentMod.zero_modulus at c, d
    have ab := firstEq.forward with first
    have cd := secondEq.forward with second
    rw [ab, cd]
    exact CongruentMod.reflexive at b * d, 0
  . unfold CongruentMod at first
    unfold CongruentMod at second
    unfold CongruentMod
    have left := mod.mul at n, a, c with nonzero
    have right := mod.mul at n, b, d with nonzero
    rw [left, right, first, second]
    rfl

theorem CongruentMod.add_left (a b c n : Nat)
  requires congruent : CongruentMod(a, b, n)
  shows CongruentMod(c + a, c + b, n)
proof
  have reflexive := CongruentMod.reflexive at c, n
  exact CongruentMod.add at c, c, a, b, n with reflexive, congruent

theorem CongruentMod.add_right (a b c n : Nat)
  requires congruent : CongruentMod(a, b, n)
  shows CongruentMod(a + c, b + c, n)
proof
  have reflexive := CongruentMod.reflexive at c, n
  exact CongruentMod.add at a, b, c, c, n with congruent, reflexive

theorem CongruentMod.mul_left (a b c n : Nat)
  requires congruent : CongruentMod(a, b, n)
  shows CongruentMod(c * a, c * b, n)
proof
  have reflexive := CongruentMod.reflexive at c, n
  exact CongruentMod.mul at c, c, a, b, n with reflexive, congruent

theorem CongruentMod.mul_right (a b c n : Nat)
  requires congruent : CongruentMod(a, b, n)
  shows CongruentMod(a * c, b * c, n)
proof
  have reflexive := CongruentMod.reflexive at c, n
  exact CongruentMod.mul at a, b, c, c, n with congruent, reflexive

theorem CongruentMod.pow (a b k n : Nat)
  requires congruent : CongruentMod(a, b, n)
  shows CongruentMod(a ^ k, b ^ k, n)
proof
  induction k
  | zero =>
    have left := pow.zero at a
    have right := pow.zero at b
    rw [left, right]
    exact CongruentMod.reflexive at 1, n
  | j + 1 using ih =>
    have left := pow.add_one at a, j
    have right := pow.add_one at b, j
    rw [left, right]
    exact CongruentMod.mul at a, b, a ^ j, b ^ j, n with congruent, ih

theorem CongruentMod.add_right_cancel (a b c n : Nat)
  requires congruent : CongruentMod(a + c, b + c, n)
  shows CongruentMod(a, b, n)
proof
  have distance := CongruentMod.divides_natDist with congruent
  have shift := natDist.add_right at a, b, c
  rw shift at distance
  exact CongruentMod.of_divides_natDist with distance

theorem CongruentMod.add_left_cancel (a b c n : Nat)
  requires congruent : CongruentMod(c + a, c + b, n)
  shows CongruentMod(a, b, n)
proof
  have distance := CongruentMod.divides_natDist with congruent
  have shift := natDist.add_left at a, b, c
  rw shift at distance
  exact CongruentMod.of_divides_natDist with distance

theorem CongruentMod.mul_cancel_coprime (a x y n : Nat)
  requires coprimeAN : gcd(a, n) = 1
  requires congruent : CongruentMod(a * x, a * y, n)
  shows CongruentMod(x, y, n)
proof
  have split : n = 0 ∨ n ≠ 0 by tauto
  cases split with zero | nonzero
  . rw zero at congruent
    rw zero at coprimeAN
    rw zero
    have value := gcd.zero_right at a
    rw value at coprimeAN
    have products := CongruentMod.zero_modulus at a * x, a * y
    have same := products.forward with congruent
    rw coprimeAN at same
    have left : 1 * x = x by ring
    have right : 1 * y = y by ring
    rw [left, right] at same
    have characterization := CongruentMod.zero_modulus at x, y
    exact characterization.backward with same
  . unfold CongruentMod at congruent
    unfold CongruentMod
    exact coprime.mul_cancel_mod at a, n, x, y
      with nonzero, coprimeAN, congruent

theorem CongruentMod.euler (a n : Nat)
  requires aboveOne : 1 < n
  requires coprimeAN : gcd(a, n) = 1
  shows CongruentMod(a ^ totient(n), 1, n)
proof
  unfold CongruentMod
  have power := euler.theorem with aboveOne, coprimeAN
  have one := mod.one_left at n with aboveOne
  rw [power, one]
  rfl

theorem CongruentMod.pow_totient_period (a k n : Nat)
  requires aboveOne : 1 < n
  requires coprimeAN : gcd(a, n) = 1
  shows CongruentMod(a ^ (k + totient(n)), a ^ k, n)
proof
  have expand := pow.add at a, totient(n), k
  rw expand
  have unit := CongruentMod.euler with aboveOne, coprimeAN
  have scaled := CongruentMod.mul_left at a ^ totient(n), 1, a ^ k, n with unit
  have one : a ^ k * 1 = a ^ k by ring
  rw one at scaled
  exact scaled

theorem CongruentMod.fermat (a p : Nat)
  requires prime : Prime(p)
  requires notDivides : ¬(p ∣ a)
  shows CongruentMod(a ^ (p - 1), 1, p)
proof
  unfold CongruentMod
  have power := fermat.little with prime, notDivides
  have aboveOne := Prime.greater_than_one with prime
  have one := mod.one_left at p with aboveOne
  rw [power, one]
  rfl

theorem CongruentMod.fermat_pow (a p : Nat)
  requires prime : Prime(p)
  shows CongruentMod(a ^ p, a, p)
proof
  have aboveOne := Prime.greater_than_one with prime
  have positive : 1 ≤ p by arith
  have predecessor := sub.add_cancel at p, 1 with positive
  have expand := pow.add_one at a, p - 1
  rw predecessor at expand
  rw expand
  have split : p ∣ a ∨ ¬(p ∣ a) by tauto
  cases split with yes | no
  . have left : p ∣ a * a ^ (p - 1)
      exact divides.multiply_right at p, a, a ^ (p - 1) with yes
    have zeroLeft := CongruentMod.zero_right_iff at a * a ^ (p - 1), p
    have zeroRight := CongruentMod.zero_right_iff at a, p
    have first := zeroLeft.backward with left
    have second := zeroRight.backward with yes
    have swapped := CongruentMod.symmetric with second
    exact CongruentMod.transitive at a * a ^ (p - 1), 0, a, p
      with first, swapped
  . have unit := CongruentMod.fermat with prime, no
    have scaled := CongruentMod.mul_left at a ^ (p - 1), 1, a, p with unit
    have one : a * 1 = a by ring
    rw one at scaled
    exact scaled