-- The sum of an indicator over a finite set counts the members it marks
-- 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.Cardinal.IndicatorSum
import NaturalSum
import Set.Indicator

defgraph indicatorMap (A V : Set) (v ∈ V) : ℕ := indicator(A, v)
proof
  have natural := indicator.natural at A, v
  exact nat.omega.forward with natural

theorem indicatorMap.natural (A V : Set)
  shows ∀ v ∈ V, indicatorMap(A, V)(v) ∈ ℕ
proof
  intros v hv
  rw indicatorMap.value with hv
  exact nat.omega.forward with (indicator.natural at A, v)

-- Meeting a point counts it exactly when the point is a member.
theorem cardinality.intersection_singleton (X a : Set)
  shows cardinality(X ∩ {a}) = indicator(X, a)
proof
  by_cases member : a ∈ X
  . have same : X ∩ {a} = {a}
      ext y
      simp
    have one := indicator.member at X, a with member
    rw [same, cardinality.singleton, one]
    rfl
  . have same : X ∩ {a} = ∅
      ext y
      simp
    have zero := indicator.absent at X, a with member
    rw [same, zero]
    exact cardinality.empty

theorem natSum.indicator_of_size (A : Set, n : Nat)
  shows ∀ V, Finite(V) → cardinality(V) = n → natSum(V, indicatorMap(A, V)) = cardinality(V ∩ A)
proof
  induction n
  | zero =>
    intros V finite size
    have has := cardinality.correct with finite
    rw size at has
    have empty := HasCardinality.empty_of_zero at V with has
    rw [empty, natSum.empty, intersection.left_empty]
    rw ← nat.zero
    rw cardinality.natural
    rfl
  | succ(k) using ih =>
    intros V finite size
    have has := cardinality.correct with finite
    rw size at has
    obtain ⟨a, ha⟩ := HasCardinality.exists_member with has
    have removed := cardinality.remove with finite, ha
    rw size at removed
    have smaller := Naturals.injective.succ with removed
    have restSize : cardinality(V ∖ {a}) = k
      symm
      exact smaller
    have restInside := difference.subset_left at V, {a}
    have restFinite := Finite.subset with finite, restInside
    have rest := ih with V ∖ {a}, restFinite, restSize
    have fresh : ¬a ∈ V ∖ {a}
      intro inside
      have parts := difference.singleton_membership.forward with inside
      apply parts.right
      rfl
    have restore := insert.restore with ha
    have graph := indicatorMap.function at A, V
    have restGraph := indicatorMap.function at A, V ∖ {a}
    have weights : ∀ v ∈ V ∖ {a}, indicatorMap(A, V)(v) ∈ ℕ
      intros v hv
      exact FunctionBetween.apply_closed with graph, (restInside with v, hv)
    have weightA := FunctionBetween.apply_closed with graph, ha
    have agree : ∀ v ∈ V ∖ {a}, indicatorMap(A, V ∖ {a})(v) = indicatorMap(A, V)(v)
      intros v hv
      have valueRest := indicatorMap.value at A, V ∖ {a}, v with hv
      have value := indicatorMap.value at A, V, v with (restInside with v, hv)
      rw [valueRest, value]
      rfl
    have sameSum := natSum.congr with agree
    have inserted := natSum.insert with restFinite, weights, weightA, fresh
    rw restore at inserted
    have valueA := indicatorMap.value at A, V, a with ha
    have restIntersection := Finite.subset with restFinite, (intersection.subset_left at V ∖ {a}, A)
    have split : a ∈ A ∨ ¬a ∈ A by tauto
    cases split with inA | outA
    . have one := indicator.member with inA
      have decomposition : V ∩ A = insert(a, (V ∖ {a}) ∩ A)
        ext x
        constructor
        . intro hx
          have both := intersection.membership.forward with hx
          have split2 : x = a ∨ x ≠ a by tauto
          apply insert.membership.backward
          cases split2 with same | different
          . left
            exact same
          . right
            have inRest : x ∈ V ∖ {a} := difference.singleton_membership.backward with ⟨both.left, different⟩
            exact intersection.membership.backward with ⟨inRest, both.right⟩
        . intro hx
          have split2 := insert.membership.forward with hx
          apply intersection.membership.backward
          cases split2 with same | inRest
          . rw same
            exact ⟨ha, inA⟩
          . have both := intersection.membership.forward with inRest
            have parts := difference.singleton_membership.forward with both.left
            exact ⟨parts.left, both.right⟩
      have freshIntersection : ¬a ∈ (V ∖ {a}) ∩ A
        intro inside
        have both := intersection.membership.forward with inside
        exact fresh with both.left
      have counted := cardinality.insert with restIntersection, freshIntersection
      rw [inserted, ← sameSum, rest, valueA, one, decomposition, counted]
      have correct := cardinality.correct with restIntersection
      have cardNat := correct.left
      rw add.right_one
      rfl
    . have zero := indicator.absent with outA
      have decomposition : V ∩ A = (V ∖ {a}) ∩ A
        ext x
        constructor
        . intro hx
          have both := intersection.membership.forward with hx
          have different : x ≠ a
            intro same
            rw same at both
            exact outA with both.right
          have inRest : x ∈ V ∖ {a} := difference.singleton_membership.backward with ⟨both.left, different⟩
          exact intersection.membership.backward with ⟨inRest, both.right⟩
        . intro hx
          have both := intersection.membership.forward with hx
          have parts := difference.singleton_membership.forward with both.left
          exact intersection.membership.backward with ⟨parts.left, both.right⟩
      rw [inserted, ← sameSum, rest, valueA, zero, decomposition]
      have correct := cardinality.correct with restIntersection
      have cardNat := correct.left
      rw add.right_unit
      rfl


-- The same for a finite set, with no size parameter.
theorem natSum.indicator (V A : Set)
  requires finite : Finite(V)
  shows natSum(V, indicatorMap(A, V)) = cardinality(V ∩ A)
proof
  have correct := cardinality.correct with finite
  have natural := correct.left
  have size : cardinality(V) = cardinality(V) by rfl
  exact natSum.indicator_of_size at A, cardinality(V)
    with natural, V, finite, size