-- Points of the cube padded to real sequences and masked by a set of
-- coordinates, the terms of the block sums used to retract onto simplices
-- 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.CubeMask
import CubeFunctions
import Real.Natural
import Set.Indicator

-- Two cube functions that agree on the cube share continuity.
theorem CubeFunctionContinuous.congr (g h n : Set)
  requires gFunction : CubeFunction(g, n)
  requires hContinuous : CubeFunctionContinuous(h, n)
  requires agree : ∀ x ∈ unitCube(n), g(x) = h(x)
  shows CubeFunctionContinuous(g, n)
proof
  constructor
  . exact gFunction
  . intros x hx E EReal Epos
    have hAt := CubeFunctionContinuous.at with hContinuous, hx
    obtain ⟨D, DReal, DPos, close⟩ := hAt with EReal, Epos
    witness D
    constructor
    . exact DReal
    . constructor
      . exact DPos
      . intros y hy near
        have gy := agree with hy
        have gx := agree with hx
        rw [gy, gx]
        exact close with hy, near

-- A point of the cube padded by zero to a real sequence.
def paddedPoint (n x : Set) : Set := x ∪ constantFunction(ℕ ∖ n, ℝ, (0:R))

theorem naturals.split (n : Nat)
  shows n ∪ (ℕ ∖ n) = ℕ
proof
  ext z
  constructor
  . intro hz
    have split := union2.membership.forward with hz
    cases split with inN | inRest
    . have zNat := natural.member at n with inN
      exact nat.omega.forward with zNat
    . have parts := difference.membership.forward with inRest
      exact parts.left
  . intro hz
    apply union2.membership.backward
    have alternatives : z ∈ n ∨ ¬z ∈ n by tauto
    cases alternatives with inN | notN
    . left
      exact inN
    . right
      exact difference.membership.backward with ⟨hz, notN⟩

theorem paddedPoint.sequence (n : Nat) (x : Set)
  requires xFunction : x : n → ℝ
  shows RealSequence(paddedPoint(n, x))
proof
  have zeroIn : (0:R) ∈ ℝ by typecheck
  have constant := constantFunction.function at ℕ ∖ n, ℝ, (0:R) with zeroIn
  have disjoint := Disjoint.symmetric with (difference.disjoint at ℕ, n)
  have union := FunctionBetween.union_disjoint with xFunction, constant, disjoint
  have domain := naturals.split at n
  have codomain := union2.idempotent at ℝ
  rw [domain, codomain] at union
  have same : paddedPoint(n, x) = x ∪ constantFunction(ℕ ∖ n, ℝ, (0:R))
    unfold paddedPoint
    rfl
  unfold RealSequence
  rwa same

theorem paddedPoint.value (n : Nat) (x i : Set)
  requires xFunction : x : n → ℝ
  requires hi : i ∈ n
  shows paddedPoint(n, x)(i) = x(i)
proof
  have zeroIn : (0:R) ∈ ℝ by typecheck
  have constant := constantFunction.function at ℕ ∖ n, ℝ, (0:R) with zeroIn
  have disjoint := Disjoint.symmetric with (difference.disjoint at ℕ, n)
  have same : paddedPoint(n, x) = x ∪ constantFunction(ℕ ∖ n, ℝ, (0:R))
    unfold paddedPoint
    rfl
  rw same
  exact FunctionBetween.union_apply_left with xFunction, constant, disjoint, hi

-- The i-th term of the sum over S: the coordinate when i ∈ S, else zero.
defgraph maskedPoint (n S x : Set) (i ∈ ℕ) : ℝ := (realNat(indicator(S, i)):R) * (paddedPoint(n, x)(i):R)
requires nNat : Nat(n)
requires xFunction : x : n → ℝ
proof
  have sequence := paddedPoint.sequence with xFunction
  have iNat : Nat(i) by typecheck
  have padded := RealSequence.value at paddedPoint(n, x), i with sequence
  have weightNat := indicator.natural at S, i
  have weight := realNat.real at indicator(S, i)
  exact realMulValue.real at realNat(indicator(S, i)), paddedPoint(n, x)(i)

theorem maskedPoint.sequence (n : Nat) (S x : Set)
  requires xFunction : x : n → ℝ
  shows RealSequence(maskedPoint(n, S, x))
proof
  exact maskedPoint.function at n, S, x with *

theorem maskedPoint.coordinate (n : Nat) (S x i : Set)
  requires xFunction : x : n → ℝ
  requires hi : i ∈ n
  shows maskedPoint(n, S, x)(i) = (realNat(indicator(S, i)):R) * (x(i):R)
proof
  have iNat := natural.member at n with hi
  have iIn := nat.omega.forward with iNat
  have value := maskedPoint.value at n, S, x, i with *
  have padded := paddedPoint.value with xFunction, hi
  rw [value, padded]
  rfl