-- Interchanging a double sum over divisors: for a weight F on pairs of
-- naturals and n ≠ 0, ∑_{d ∣ n} ∑_{e ∣ n/d} F(d, e) = ∑_{e ∣ n} ∑_{d ∣ n/e}
-- F(d, e), both being the sum of F over the pairs of divisors whose product
-- divides n, computed by rows and by columns of a matrix on
-- divisors(n) × divisors(n) that vanishes off those pairs
-- 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.Combinatorics.DivisorInterchange
import DivisorSums

-- The row e ↦ F(d, e) and the column d ↦ F(d, e) of a weight on pairs.
defgraph rowTerms (F d : Set) (e ∈ ℕ) : ℤ := F((d, e))
requires weights : F : ℕ × ℕ → ℤ
requires rowIndex : d ∈ ℕ
proof
  have pair := product.pair with rowIndex, member
  exact FunctionBetween.apply_closed with weights, pair

defgraph columnTerms (F e : Set) (d ∈ ℕ) : ℤ := F((d, e))
requires weights : F : ℕ × ℕ → ℤ
requires columnIndex : e ∈ ℕ
proof
  have pair := product.pair with member, columnIndex
  exact FunctionBetween.apply_closed with weights, pair

-- The inner sums d ↦ ∑_{e ∣ n/d} F(d, e) and e ↦ ∑_{d ∣ n/e} F(d, e).
defgraph rowDivisorSums (F : Set) (n : Nat) (d ∈ ℕ) : ℤ :=
  divisorSum(quotientNat(n, d), rowTerms(F, d))
requires weights : F : ℕ × ℕ → ℤ
proof
  have dNat := nat.omega.backward with member
  have quotientNatural := quotientNat.natural at d, n
  have rows := rowTerms.function at F, d with weights, member
  exact divisorSum.integer at quotientNat(n, d), rowTerms(F, d) with rows

defgraph columnDivisorSums (F : Set) (n : Nat) (e ∈ ℕ) : ℤ :=
  divisorSum(quotientNat(n, e), columnTerms(F, e))
requires weights : F : ℕ × ℕ → ℤ
proof
  have eNat := nat.omega.backward with member
  have quotientNatural := quotientNat.natural at e, n
  have columns := columnTerms.function at F, e with weights, member
  exact divisorSum.integer at quotientNat(n, e), columnTerms(F, e) with
    columns

-- The pairs of divisors of n whose product divides n.
def dividingPairs (n : Nat) : Set :=
  { q ∈ divisors(n) × divisors(n) |
    ∃ d ∈ divisors(n), ∃ e ∈ divisors(n), q = (d, e) ∧ mul(d, e) ∣ n }

theorem dividingPairs.pair (n : Nat, d e : Set)
  shows (d, e) ∈ dividingPairs(n) ↔
    d ∈ divisors(n) ∧ e ∈ divisors(n) ∧ mul(d, e) ∣ n
proof
  have characterization := dividingPairs.membership at n, (d, e)
  constructor
  . intro member
    have unpacked := characterization.forward with member
    obtain ⟨pairMember, x, hx, y, hy, shape, factor⟩ := unpacked
    obtain ⟨dEq, eEq⟩ := orderedPair.injective with shape
    rw [dEq, eEq]
    exact ⟨hx, hy, factor⟩
  . intro triple
    rcases triple with ⟨hd, he, factor⟩
    have pairMember := product.pair with hd, he
    apply characterization.backward
    refine ⟨pairMember, d, hd, e, he, rfl, factor⟩

theorem dividingPairs.symmetric (n : Nat, d e : Set)
  requires member : (d, e) ∈ dividingPairs(n)
  shows (e, d) ∈ dividingPairs(n)
proof
  have triple := dividingPairs.pair.forward with member
  rcases triple with ⟨hd, he, factor⟩
  have dNat := divisors.natural with hd
  have eNat := divisors.natural with he
  have commuted := mul.commutative at e, d
  rw commuted at factor
  exact dividingPairs.pair.backward with ⟨he, hd, factor⟩

theorem dividingPairs.iff_quotient (n : Nat, d e : Set)
  requires nonzero : n ≠ 0
  requires member : d ∈ divisors(n)
  shows (d, e) ∈ dividingPairs(n) ↔ e ∈ divisors(quotientNat(n, d))
proof
  have dNat := divisors.natural with member
  have factor := divisors.divides with member
  have quotientNonzero := divides.quotient_ne_zero with nonzero, factor
  have inner := divisors.quotient_subset with nonzero, factor
  constructor
  . intro pairMember
    have triple := dividingPairs.pair.forward with pairMember
    rcases triple with ⟨_, he, product⟩
    have eNat := divisors.natural with he
    have characterization := divides.mul_iff_quotient at d, e, n with nonzero,
      factor
    have quotientDivides := characterization.forward with product
    have below := divides.le with quotientDivides, quotientNonzero
    have membership := divisors.membership_le at quotientNat(n, d), e
    exact membership.backward with ⟨below, quotientDivides⟩
  . intro quotientMember
    have he := inner with e, quotientMember
    have eNat := divisors.natural with he
    have quotientDivides := divisors.divides with quotientMember
    have characterization := divides.mul_iff_quotient at d, e, n with nonzero,
      factor
    have product := characterization.backward with quotientDivides
    exact dividingPairs.pair.backward with ⟨member, he, product⟩

-- The matrix F(d, e) · [d · e ∣ n] on divisors(n) × divisors(n).
defgraph divisorMatrix (F : Set) (n : Nat) (q ∈ divisors(n) × divisors(n)) :
    ℤ := integerMulValue(F(q), integerOfNat(indicator(dividingPairs(n), q)))
requires weights : F : ℕ × ℕ → ℤ
proof
  have naturals := divisors.subset_naturals at n
  have subset := product.monotone with naturals, naturals
  have inNaturals := subset with q, member
  have value := FunctionBetween.apply_closed with weights, inNaturals
  have indicatorNat := indicator.natural at dividingPairs(n), q
  have mask := integerOfNat.integer at indicator(dividingPairs(n), q)
  exact integerMulValue.integer at F(q),
    integerOfNat(indicator(dividingPairs(n), q)) with *

theorem divisorMatrix.inside (F : Set) (n : Nat) (d e : Set)
  requires weights : F : ℕ × ℕ → ℤ
  requires member : (d, e) ∈ dividingPairs(n)
  shows divisorMatrix(F, n)((d, e)) = F((d, e))
proof
  have triple := dividingPairs.pair.forward with member
  rcases triple with ⟨hd, he, _⟩
  have naturals := divisors.subset_naturals at n
  have pairMember := product.pair with hd, he
  have dMember := naturals with d, hd
  have eMember := naturals with e, he
  have inNaturals := product.pair with dMember, eMember
  have value := FunctionBetween.apply_closed with weights, inNaturals
  have entry := divisorMatrix.value at F, n, (d, e) with weights, pairMember
  have one := indicator.member with member
  rw [entry, one, integerOfNat.one]
  ring

theorem divisorMatrix.outside (F : Set) (n : Nat) (d e : Set)
  requires weights : F : ℕ × ℕ → ℤ
  requires pairMember : (d, e) ∈ divisors(n) × divisors(n)
  requires absent : ¬((d, e) ∈ dividingPairs(n))
  shows divisorMatrix(F, n)((d, e)) = integerZero
proof
  have naturals := divisors.subset_naturals at n
  have subset := product.monotone with naturals, naturals
  have inNaturals := subset with (d, e), pairMember
  have value := FunctionBetween.apply_closed with weights, inNaturals
  have entry := divisorMatrix.value at F, n, (d, e) with weights, pairMember
  have zero := indicator.absent with absent
  rw [entry, zero, integerOfNat.zero]
  ring

-- Row d of the matrix sums to ∑_{e ∣ n/d} F(d, e).
theorem divisorMatrix.row (F : Set) (n : Nat) (d : Set)
  requires nonzero : n ≠ 0
  requires weights : F : ℕ × ℕ → ℤ
  requires member : d ∈ divisors(n)
  shows finiteRowSums(divisors(n), divisors(n), ℤ, integerAdd, integerZero,
      integerNeg, integerMul, integerOne, divisorMatrix(F, n))(d) =
    rowDivisorSums(F, n)(d)
proof
  have ring := integers.comm_ring
  have finite := divisors.finite at n
  have naturals := divisors.subset_naturals at n
  have matrix := divisorMatrix.function at F, n with weights
  have rowValue := finiteRowSums.value at divisors(n), divisors(n), ℤ,
    integerAdd, integerZero, integerNeg, integerMul, integerOne,
      divisorMatrix(F, n), d with ring, finite, matrix, member
  have dNat := divisors.natural with member
  have dMember := naturals with d, member
  have factor := divisors.divides with member
  have rows := rowTerms.function at F, d with weights, dMember
  have target := rowDivisorSums.value at F, n, d with weights
  have inner := divisors.quotient_subset with nonzero, factor
  have innerNaturals := divisors.subset_naturals at quotientNat(n, d)
  have rowWeights := finiteRowWeights.function at divisors(n), divisors(n), ℤ,
    divisorMatrix(F, n), d with matrix, member
  have restrictedRows := restrict.domain_function with rows, innerNaturals
  have agreement : ∀ s ∈ divisors(quotientNat(n, d)),
      finiteRowWeights(divisors(n), divisors(n), ℤ, divisorMatrix(F, n), d)(s) =
        restrict(rowTerms(F, d), divisors(quotientNat(n, d)))(s)
    intros s hs
    have sMemberD := inner with s, hs
    have sMember := naturals with s, sMemberD
    have entry := finiteRowWeights.value at divisors(n), divisors(n), ℤ,
      divisorMatrix(F, n), d, s with matrix, member, sMemberD
    have characterization := dividingPairs.iff_quotient at n, d, s with nonzero,
      member
    have pairIn := characterization.backward with hs
    have inside := divisorMatrix.inside with weights, pairIn
    have termValue := rowTerms.value at F, d, s with weights, dMember
    have restricted := restrict.domain_apply with rows, innerNaturals, hs
    cc
  have vanishes : ∀ s ∈ divisors(n), ¬(s ∈ divisors(quotientNat(n, d))) →
      finiteRowWeights(divisors(n), divisors(n), ℤ, divisorMatrix(F, n), d)(s) =
        integerZero
    intros s hs outside
    have entry := finiteRowWeights.value at divisors(n), divisors(n), ℤ,
      divisorMatrix(F, n), d, s with matrix, member, hs
    have characterization := dividingPairs.iff_quotient at n, d, s with nonzero,
      member
    have notPair : ¬((d, s) ∈ dividingPairs(n))
      intro pairIn
      apply outside
      exact characterization.forward with pairIn
    have pairMember := product.pair with member, hs
    have zero := divisorMatrix.outside with weights, pairMember, notPair
    cc
  have extension := finiteRingSum.zero_extension with ring, finite, rowWeights,
    restrictedRows, inner, agreement, vanishes
  rw [rowValue, target]
  simp [divisorSum, integerSetSum]

-- Row e of the transposed matrix sums to ∑_{d ∣ n/e} F(d, e).
theorem divisorMatrix.column (F : Set) (n : Nat) (e : Set)
  requires nonzero : n ≠ 0
  requires weights : F : ℕ × ℕ → ℤ
  requires member : e ∈ divisors(n)
  shows finiteRowSums(divisors(n), divisors(n), ℤ, integerAdd, integerZero,
      integerNeg, integerMul, integerOne, transposedMatrix(divisors(n),
        divisors(n), ℤ, divisorMatrix(F, n)))(e) =
    columnDivisorSums(F, n)(e)
proof
  have ring := integers.comm_ring
  have finite := divisors.finite at n
  have naturals := divisors.subset_naturals at n
  have matrix := divisorMatrix.function at F, n with weights
  have transposed := transposedMatrix.function at divisors(n), divisors(n), ℤ,
    divisorMatrix(F, n) with matrix
  have rowValue := finiteRowSums.value at divisors(n), divisors(n), ℤ,
    integerAdd, integerZero, integerNeg, integerMul, integerOne,
      transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n)), e
        with ring, finite, transposed, member
  have eNat := divisors.natural with member
  have eMember := naturals with e, member
  have factor := divisors.divides with member
  have columns := columnTerms.function at F, e with weights, eMember
  have target := columnDivisorSums.value at F, n, e with weights
  have inner := divisors.quotient_subset with nonzero, factor
  have innerNaturals := divisors.subset_naturals at quotientNat(n, e)
  have rowWeights := finiteRowWeights.function at divisors(n), divisors(n), ℤ,
    transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n)), e with
      transposed, member
  have restrictedColumns := restrict.domain_function with columns,
    innerNaturals
  have agreement : ∀ s ∈ divisors(quotientNat(n, e)),
      finiteRowWeights(divisors(n), divisors(n), ℤ,
        transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n)),
          e)(s) =
        restrict(columnTerms(F, e), divisors(quotientNat(n, e)))(s)
    intros s hs
    have sMemberD := inner with s, hs
    have sMember := naturals with s, sMemberD
    have entry := finiteRowWeights.value at divisors(n), divisors(n), ℤ,
      transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n)), e, s
        with transposed, member, sMemberD
    have swapped := transposedMatrix.pair at divisors(n), divisors(n), ℤ,
      divisorMatrix(F, n), s, e with matrix, sMemberD, member
    have characterization := dividingPairs.iff_quotient at n, e, s with nonzero,
      member
    have pairIn := characterization.backward with hs
    have symmetric := dividingPairs.symmetric with pairIn
    have inside := divisorMatrix.inside with weights, symmetric
    have termValue := columnTerms.value at F, e, s with weights, eMember
    have restricted := restrict.domain_apply with columns, innerNaturals, hs
    cc
  have vanishes : ∀ s ∈ divisors(n), ¬(s ∈ divisors(quotientNat(n, e))) →
      finiteRowWeights(divisors(n), divisors(n), ℤ,
        transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n)),
          e)(s) = integerZero
    intros s hs outside
    have entry := finiteRowWeights.value at divisors(n), divisors(n), ℤ,
      transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n)), e, s
        with transposed, member, hs
    have swapped := transposedMatrix.pair at divisors(n), divisors(n), ℤ,
      divisorMatrix(F, n), s, e with matrix, hs, member
    have characterization := dividingPairs.iff_quotient at n, e, s with nonzero,
      member
    have notPair : ¬((s, e) ∈ dividingPairs(n))
      intro pairIn
      apply outside
      have symmetric := dividingPairs.symmetric with pairIn
      exact characterization.forward with symmetric
    have pairMember := product.pair with hs, member
    have zero := divisorMatrix.outside with weights, pairMember, notPair
    cc
  have extension := finiteRingSum.zero_extension with ring, finite, rowWeights,
    restrictedColumns, inner, agreement, vanishes
  rw [rowValue, target]
  simp [divisorSum, integerSetSum]

-- Summing by rows and by columns agrees:
-- ∑_{d ∣ n} ∑_{e ∣ n/d} F(d, e) = ∑_{e ∣ n} ∑_{d ∣ n/e} F(d, e).
theorem divisorSum.interchange (F : Set) (n : Nat)
  requires nonzero : n ≠ 0
  requires weights : F : ℕ × ℕ → ℤ
  shows divisorSum(n, rowDivisorSums(F, n)) =
    divisorSum(n, columnDivisorSums(F, n))
proof
  have ring := integers.comm_ring
  have finite := divisors.finite at n
  have naturals := divisors.subset_naturals at n
  have matrix := divisorMatrix.function at F, n with weights
  have transposed := transposedMatrix.function at divisors(n), divisors(n), ℤ,
    divisorMatrix(F, n) with matrix
  have rowsFunction := rowDivisorSums.function at F, n with weights
  have columnsFunction := columnDivisorSums.function at F, n with weights
  have restrictedRows := restrict.domain_function with rowsFunction, naturals
  have restrictedColumns := restrict.domain_function with columnsFunction,
    naturals
  have rowSums := finiteRowSums.function at divisors(n), divisors(n), ℤ,
    integerAdd, integerZero, integerNeg, integerMul, integerOne,
      divisorMatrix(F, n) with ring, finite, matrix
  have columnSums := finiteRowSums.function at divisors(n), divisors(n), ℤ,
    integerAdd, integerZero, integerNeg, integerMul, integerOne,
      transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n)) with
        ring, finite, transposed
  have interchange := finiteDoubleSum.interchange with ring, finite, finite,
    matrix
  have leftAgree : ∀ d ∈ divisors(n), restrict(rowDivisorSums(F, n),
      divisors(n))(d) = finiteRowSums(divisors(n), divisors(n), ℤ, integerAdd,
        integerZero, integerNeg, integerMul, integerOne, divisorMatrix(F, n))(d)
    intros d hd
    have row := divisorMatrix.row with nonzero, weights, hd
    have restricted := restrict.domain_apply with rowsFunction, naturals, hd
    cc
  have rightAgree : ∀ e ∈ divisors(n), restrict(columnDivisorSums(F, n),
      divisors(n))(e) = finiteRowSums(divisors(n), divisors(n), ℤ, integerAdd,
        integerZero, integerNeg, integerMul, integerOne,
          transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n)))(e)
    intros e he
    have column := divisorMatrix.column with nonzero, weights, he
    have restricted := restrict.domain_apply with columnsFunction, naturals, he
    cc
  have leftSum := finiteRingSum.congr with ring, finite, restrictedRows,
    rowSums, leftAgree
  have rightSum := finiteRingSum.congr with ring, finite, restrictedColumns,
    columnSums, rightAgree
  have leftDouble : finiteDoubleSum(divisors(n), divisors(n), ℤ, integerAdd,
      integerZero, integerNeg, integerMul, integerOne, divisorMatrix(F, n)) =
    finiteRingSum(divisors(n), ℤ, integerAdd, integerZero, integerNeg,
      integerMul, integerOne, finiteRowSums(divisors(n), divisors(n), ℤ,
        integerAdd, integerZero, integerNeg, integerMul, integerOne,
          divisorMatrix(F, n)))
    unfold finiteDoubleSum
    rfl
  have rightDouble : finiteDoubleSum(divisors(n), divisors(n), ℤ, integerAdd,
      integerZero, integerNeg, integerMul, integerOne,
        transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n))) =
    finiteRingSum(divisors(n), ℤ, integerAdd, integerZero, integerNeg,
      integerMul, integerOne, finiteRowSums(divisors(n), divisors(n), ℤ,
        integerAdd, integerZero, integerNeg, integerMul, integerOne,
          transposedMatrix(divisors(n), divisors(n), ℤ, divisorMatrix(F, n))))
    unfold finiteDoubleSum
    rfl
  unfold divisorSum
  unfold integerSetSum
  cc