-- Summing grouped coefficients recovers the original finite sum (ZF).
-- 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.Algebra.Ring.Regrouping
import FiniteDoubleSum
import Pushforward

defunique groupingEntry (D C A zero h w t : Set) (a ∈ A) :=
  ∃ d ∈ D, ∃ c ∈ C, t = (d, c) ∧ a = singleValue(A, zero, w(d), h(d), c)
requires zeroIn : zero ∈ A
requires weights : w : D → A
requires pair : t ∈ D × C
proof
  obtain ⟨bound, d, hd, c, hc, shape⟩ := product.membership.forward at
    D, C, t with pair
  have coefficient := FunctionBetween.apply_closed at w, D, A, d with *
  have value := singleValue.closed at A, zero, w(d), h(d), c with *
  refine ⟨singleValue(A, zero, w(d), h(d), c), value, d, hd, c, hc, shape, rfl⟩
uniqueness
  intros u hu first v hv second
  obtain ⟨d, hd, c, hc, shapeU, valueU⟩ := first
  obtain ⟨e, he, b, hb, shapeV, valueV⟩ := second
  have equal : (d, c) = (e, b)
    cc
  have coordinates := orderedPair.injective at d, c, e, b with equal
  cc

defgraph groupingMatrix (D C A zero h w : Set) (t ∈ D × C) : A := groupingEntry(D, C, A, zero, h, w, t)
requires zeroIn : zero ∈ A
requires weights : w : D → A
proof
  have selected := groupingEntry.selected at D, C, A, zero, h, w, t with *
  exact selected.left

theorem groupingMatrix.pair (D C A zero h w d c : Set)
  requires zeroIn : zero ∈ A
  requires weights : w : D → A
  requires row : d ∈ D
  requires column : c ∈ C
  shows groupingMatrix(D, C, A, zero, h, w)((d, c)) = singleValue(A, zero, w(d), h(d), c)
proof
  have pair := product.pair at d, c, D, C with *
  have selected := groupingEntry.selected at D, C, A, zero, h, w, (d, c) with *
  obtain ⟨e, he, b, hb, shape, output⟩ := selected.right
  have coordinates := orderedPair.injective with shape
  have value := groupingMatrix.value at D, C, A, zero, h, w, (d, c) with *
  cc

section Rings
variable (D C A add zero neg mul one h w : Set)
hypothesis ring : IsCommRing(A, add, zero, neg, mul, one)
hypothesis weights : w : D → A

theorem groupingMatrix.rows
  requires finiteC : Finite(C)
  requires indexMap : h : D → C
  shows finiteRowSums(D, C, A, add, zero, neg, mul, one, groupingMatrix(D, C, A, zero, h, w)) = w
proof
  have ordinary := IsCommRing.ring with ring
  have zeroIn := IsRing.zero_in_carrier with ordinary
  have matrix := groupingMatrix.function at D, C, A, zero, h, w with *
  have rows := finiteRowSums.function at D, C, A, add, zero, neg, mul, one, groupingMatrix(D, C, A, zero, h, w) with *
  apply FunctionBetween.ext with rows, weights
  intros d hd
  have coefficient := FunctionBetween.apply_closed at w, D, A, d with *
  have index := FunctionBetween.apply_closed at h, D, C, d with *
  have row := finiteRowWeights.function at D, C, A, groupingMatrix(D, C, A, zero, h, w), d with *
  have single := coefficientSingle.function at C, A, zero, w(d), h(d) with *
  have same : finiteRowWeights(D, C, A, groupingMatrix(D, C, A, zero, h, w), d) = coefficientSingle(C, A, zero, w(d), h(d))
    apply FunctionBetween.ext with row, single
    intros c hc
    have left := finiteRowWeights.value at D, C, A, groupingMatrix(D, C, A, zero, h, w), d, c with *
    have entry := groupingMatrix.pair at D, C, A, zero, h, w, d, c with *
    have right := coefficientSingle.value at C, A, zero, w(d), h(d), c with *
    cc
  have sum := finiteRingSum.single with ring, coefficient, index
  have value := finiteRowSums.value at D, C, A, add, zero, neg, mul, one, groupingMatrix(D, C, A, zero, h, w), d with *
  cc

theorem groupingMatrix.columns
  requires finiteD : Finite(D)
  shows finiteRowSums(C, D, A, add, zero, neg, mul, one, transposedMatrix(D, C, A, groupingMatrix(D, C, A, zero, h, w))) = finitePushforward(D, C, A, add, zero, neg, mul, one, h, w)
proof
  have ordinary := IsCommRing.ring with ring
  have zeroIn := IsRing.zero_in_carrier with ordinary
  have matrix := groupingMatrix.function at D, C, A, zero, h, w with *
  have transpose := transposedMatrix.function at D, C, A, groupingMatrix(D, C, A, zero, h, w) with *
  have rows := finiteRowSums.function at C, D, A, add, zero, neg, mul, one, transposedMatrix(D, C, A, groupingMatrix(D, C, A, zero, h, w)) with *
  have grouped := finitePushforward.function at D, C, A, add, zero, neg, mul, one, h, w with *
  apply FunctionBetween.ext with rows, grouped
  intros c hc
  have row := finiteRowWeights.function at C, D, A, transposedMatrix(D, C, A, groupingMatrix(D, C, A, zero, h, w)), c with *
  have selected := pushforwardWeights.function at D, A, add, zero, neg, mul, one, h, w, c with *
  have same : finiteRowWeights(C, D, A, transposedMatrix(D, C, A, groupingMatrix(D, C, A, zero, h, w)), c) = pushforwardWeights(D, A, add, zero, neg, mul, one, h, w, c)
    apply FunctionBetween.ext with row, selected
    intros d hd
    have left := finiteRowWeights.value at C, D, A, transposedMatrix(D, C, A, groupingMatrix(D, C, A, zero, h, w)), c, d with *
    have swapped := transposedMatrix.pair with hd, hc
    have entry := groupingMatrix.pair at D, C, A, zero, h, w, d, c with *
    have right := pushforwardWeights.value at D, A, add, zero, neg, mul, one, h, w, c, d with *
    cc
  have left := finiteRowSums.value at C, D, A, add, zero, neg, mul, one, transposedMatrix(D, C, A, groupingMatrix(D, C, A, zero, h, w)), c with *
  have right := finitePushforward.value at D, C, A, add, zero, neg, mul, one, h, w, c with *
  cc

theorem finitePushforward.sum
  requires finiteD : Finite(D)
  requires finiteC : Finite(C)
  requires indexMap : h : D → C
  shows finiteRingSum(C, A, add, zero, neg, mul, one, finitePushforward(D, C, A, add, zero, neg, mul, one, h, w)) = finiteRingSum(D, A, add, zero, neg, mul, one, w)
proof
  have ordinary := IsCommRing.ring with ring
  have zeroIn := IsRing.zero_in_carrier with ordinary
  have matrix := groupingMatrix.function at D, C, A, zero, h, w with *
  have rows := groupingMatrix.rows with ring, weights, finiteC, indexMap
  have columns := groupingMatrix.columns at D, C, A, add, zero, neg, mul, one, h, w with *
  have interchange := finiteDoubleSum.interchange with ring, finiteD, finiteC, matrix
  unfold finiteDoubleSum at interchange
  cc

-- Only the finite image needs to be summed, even if C itself is infinite.
theorem finitePushforward.sum_image
  requires finiteD : Finite(D)
  requires indexMap : h : D → C
  shows finiteRingSum(functionImage(h, D, C), A, add, zero, neg, mul, one,
    finitePushforward(D, functionImage(h, D, C), A, add, zero, neg, mul, one, h, w)) =
    finiteRingSum(D, A, add, zero, neg, mul, one, w)
proof
  have reflexive : D ⊆ D by simp [Subset]
  have finiteImage := Finite.function_image at D, C, h, D with *
  have corestricted := functionImage.corestrict at h, D, C with *
  exact finitePushforward.sum with *

end Rings
theorem finitePushforward.constant_index (D A add zero neg mul one w c : Set)
  requires ring : IsCommRing(A, add, zero, neg, mul, one)
  requires finite : Finite(D)
  requires weights : w : D → A
  shows finitePushforward(D, {c}, A, add, zero, neg, mul, one, constantFunction(D, {c}, c), w)(c) =
    finiteRingSum(D, A, add, zero, neg, mul, one, w)
proof
  have member : c ∈ {c} by simp [singleton.membership]
  have indexMap := constantFunction.function at D, {c}, c with member
  have finiteC := Finite.singleton at c
  have grouped := finitePushforward.function at D, {c}, A, add, zero, neg, mul, one, constantFunction(D, {c}, c), w with *
  have total := finitePushforward.sum at D, {c}, A, add, zero, neg, mul, one, constantFunction(D, {c}, c), w with *
  have singleton := finiteRingSum.singleton with *
  cc