-- Block sums of a point of the cube: each term is the coordinate or zero,
-- the sum is nonnegative and at least any selected coordinate, and it
-- transforms linearly under a common affine change of the selected coordinates
-- 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.FixedPoint.BlockSums
import CubeAssembly

-- A masked term is the coordinate on the block and zero off it.
theorem maskedPoint.on_block (n : Nat) (S x i : Set)
  requires xFunction : x : n → ℝ
  requires hi : i ∈ n
  requires inBlock : i ∈ S
  shows maskedPoint(n, S, x)(i) = x(i)
proof
  have value := maskedPoint.coordinate at n, S, x, i with xFunction, hi
  have one := indicator.member with inBlock
  have xi := FunctionBetween.apply_closed with xFunction, hi
  rw [value, one, realNat.one]
  exact realMulValue.one_left

theorem maskedPoint.off_block (n : Nat) (S x i : Set)
  requires xFunction : x : n → ℝ
  requires hi : i ∈ n
  requires outside : ¬i ∈ S
  shows maskedPoint(n, S, x)(i) = (0:R)
proof
  have value := maskedPoint.coordinate at n, S, x, i with xFunction, hi
  have zero := indicator.absent with outside
  have xi := FunctionBetween.apply_closed with xFunction, hi
  rw [value, zero, realNat.zero]
  exact realMulValue.zero_left

theorem maskedPoint.nonnegative (n : Nat) (S x i : Set)
  requires hx : x ∈ unitCube(n)
  requires hi : i ∈ n
  shows (0:R) ≤ (maskedPoint(n, S, x)(i):R)
proof
  have xFunction := unitCube.point with hx
  have alternatives : i ∈ S ∨ ¬i ∈ S by tauto
  cases alternatives with inBlock | outside
  . have same := maskedPoint.on_block with xFunction, hi, inBlock
    have interval := unitCube.coordinate with hx, hi
    obtain ⟨xReal, lower, upper⟩ := closedInterval.membership.forward with interval
    rwa same
  . have same := maskedPoint.off_block with xFunction, hi, outside
    rw same
    linarith

-- Partial sums of nonnegative terms are nonnegative and bound every term.
theorem realPartialSum.nonnegative_prefix (u : Set) (m : Nat)
  requires sequence : RealSequence(u)
  shows (∀ k ∈ m, (0:R) ≤ (u(k):R)) → (0:R) ≤ (realPartialSum(u, m):R)
proof
  induction m
  | zero =>
    intro terms
    have zero := realPartialSum.zero at u
    rw zero
    linarith
  | succ(j) using ih =>
    intro terms
    have jNat : Nat(j) by typecheck
    have last := terms with (bootstrap.succ.self at j)
    have smaller : ∀ k ∈ j, (0:R) ≤ (u(k):R)
      intros k hk
      have kNat := natural.member at j with hk
      have less := Less.of_membership with hk
      have below : k < succ(j) by arith
      have inSucc : k ∈ succ(j)
        exact membership.of_Less with below
      exact terms with inSucc
    have previous := ih with smaller
    have step := realPartialSum.succ at u, j
    have lastReal := RealSequence.value at u, j with sequence
    rw step
    linarith

theorem realPartialSum.term_le_prefix (u : Set) (m : Nat)
  requires sequence : RealSequence(u)
  shows (∀ k ∈ m, (0:R) ≤ (u(k):R)) → ∀ j ∈ m, (u(j):R) ≤ (realPartialSum(u, m):R)
proof
  induction m
  | zero =>
    intros terms j hj
    exfalso
    rw nat.zero at hj
    simpa using hj
  | succ(i) using ih =>
    intros terms j hj
    have iNat : Nat(i) by typecheck
    have smaller : ∀ k ∈ i, (0:R) ≤ (u(k):R)
      intros k hk
      have kNat := natural.member at i with hk
      have less := Less.of_membership with hk
      have below : k < succ(i) by arith
      have inSucc : k ∈ succ(i)
        exact membership.of_Less with below
      exact terms with inSucc
    have nonnegative := realPartialSum.nonnegative_prefix at u, i with sequence
    have prefixNonneg := nonnegative with smaller
    have step := realPartialSum.succ at u, i
    have lastReal := RealSequence.value at u, i with sequence
    have last := terms with (bootstrap.succ.self at i)
    have split := bootstrap.succ.membership.forward with hj
    rw step
    cases split with below | same
    . have bounds := ih with smaller
      have bound := bounds with below
      have jNat := natural.member at i with below
      have jReal := RealSequence.value at u, j with sequence
      linarith
    . rw same
      linarith

-- Block sums are nonnegative and at least any selected coordinate.
theorem sumCube.terms_nonnegative (n : Nat) (S x : Set)
  requires hx : x ∈ unitCube(n)
  shows ∀ k ∈ n, (0:R) ≤ (maskedPoint(n, S, x)(k):R)
proof
  intros k hk
  exact maskedPoint.nonnegative with hx, hk

theorem sumCube.nonnegative (n : Nat) (S x : Set)
  requires hx : x ∈ unitCube(n)
  shows (0:R) ≤ (sumCube(n, S)(x):R)
proof
  have xFunction := unitCube.point with hx
  have sequence := maskedPoint.sequence at n, S, x with xFunction
  have value := sumCube.value at n, S, x with hx
  have terms := sumCube.terms_nonnegative at n, S, x with hx
  have bound := realPartialSum.nonnegative_prefix with sequence
  rw value
  exact bound with terms

theorem sumCube.ge_coordinate (n : Nat) (S x j : Set)
  requires hx : x ∈ unitCube(n)
  requires hj : j ∈ n
  requires inBlock : j ∈ S
  shows (x(j):R) ≤ (sumCube(n, S)(x):R)
proof
  have xFunction := unitCube.point with hx
  have sequence := maskedPoint.sequence at n, S, x with xFunction
  have value := sumCube.value at n, S, x with hx
  have terms := sumCube.terms_nonnegative at n, S, x with hx
  have bounds := realPartialSum.term_le_prefix with sequence
  have forAll := bounds with terms
  have bound := forAll with hj
  have same := maskedPoint.on_block with xFunction, hj, inBlock
  rw value
  rwa same at bound

-- When the selected coordinates of z are an affine image of those of x and a,
-- the block sum of z is the same affine image of the block sums of x and a.
theorem sumCube.affine (n : Nat) (S x a z : Set) (E C : Real)
  requires hx : x ∈ unitCube(n)
  requires ha : a ∈ unitCube(n)
  requires hz : z ∈ unitCube(n)
  requires shifted : ∀ j ∈ n, j ∈ S → z(j) = ((x(j):R) + E * (a(j):R)) * C
  shows sumCube(n, S)(z) = ((sumCube(n, S)(x):R) + E * (sumCube(n, S)(a):R)) * C
proof
  have xFunction := unitCube.point with hx
  have aFunction := unitCube.point with ha
  have zFunction := unitCube.point with hz
  have maskedX := maskedPoint.sequence at n, S, x with xFunction
  have maskedA := maskedPoint.sequence at n, S, a with aFunction
  have maskedZ := maskedPoint.sequence at n, S, z with zFunction
  have constantE := constantSequence.sequence at E
  have constantC := constantSequence.sequence at C
  have scaledA := mulSequence.sequence with constantE, maskedA
  have inner := addSequence.sequence with maskedX, scaledA
  have agree : ∀ k ∈ ℕ, k < n → maskedPoint(n, S, z)(k) = mulSequence(constantSequence(C), addSequence(maskedPoint(n, S, x), mulSequence(constantSequence(E), maskedPoint(n, S, a))))(k)
    intros k hk hkn
    have kNat : Nat(k) by typecheck
    have hkIn : k ∈ n by exact membership.of_Less with hkn
    have xReal := RealSequence.value at maskedPoint(n, S, x), k with maskedX
    have aReal := RealSequence.value at maskedPoint(n, S, a), k with maskedA
    have eReal := RealSequence.value at constantSequence(E), k with constantE
    have cReal := RealSequence.value at constantSequence(C), k with constantC
    have eValue := constantSequence.value at E, k
    have cValue := constantSequence.value at C, k
    have scaledValue := mulSequence.value with constantE, maskedA, eReal, aReal
    have scaledReal := RealSequence.value at mulSequence(constantSequence(E), maskedPoint(n, S, a)), k with scaledA
    have innerValue := addSequence.value with maskedX, scaledA, xReal, scaledReal
    have innerReal := RealSequence.value at addSequence(maskedPoint(n, S, x), mulSequence(constantSequence(E), maskedPoint(n, S, a))), k with inner
    have outerValue := mulSequence.value with constantC, inner, cReal, innerReal
    rw [outerValue, innerValue, scaledValue, eValue, cValue]
    have alternatives : k ∈ S ∨ ¬k ∈ S by tauto
    cases alternatives with inBlock | outside
    . have zSame := maskedPoint.on_block with zFunction, hkIn, inBlock
      have xSame := maskedPoint.on_block with xFunction, hkIn, inBlock
      have aSame := maskedPoint.on_block with aFunction, hkIn, inBlock
      have zValue := shifted with hkIn, inBlock
      have xk := unitCube.coordinate_real with hx, hkIn
      have ak := unitCube.coordinate_real with ha, hkIn
      rw [zSame, xSame, aSame, zValue]
      ring
    . have zSame := maskedPoint.off_block with zFunction, hkIn, outside
      have xSame := maskedPoint.off_block with xFunction, hkIn, outside
      have aSame := maskedPoint.off_block with aFunction, hkIn, outside
      rw [zSame, xSame, aSame]
      ring
  have outer := mulSequence.sequence with constantC, inner
  have equal := realPartialSum.eq_prefix with maskedZ, outer, agree
  have scaleC := realPartialSum.scale at addSequence(maskedPoint(n, S, x), mulSequence(constantSequence(E), maskedPoint(n, S, a))), n, C with inner
  have addXA := realPartialSum.add with maskedX, scaledA
  have scaleE := realPartialSum.scale at maskedPoint(n, S, a), n, E with maskedA
  have zValue := sumCube.value at n, S, z with hz
  have xValue := sumCube.value at n, S, x with hx
  have aValue := sumCube.value at n, S, a with ha
  rw [zValue, xValue, aValue, equal, scaleC, addXA, scaleE]
  ring