LowerBounds.lean

/-
  Copyright (c) 2026 Joe Gallagher / Mintpath LLC.
  Released under MIT license as described in the file LICENSE.
  Authors: Joe Gallagher

  Formalization of lower bound techniques from Section V of AUY83.

  Includes:
  - Lower bound techniques for fixed partitions
  - The rectangle property and its consequences
  - Theorem 11: lower bound on protocol partition number for inner product
-/
import Mathlib
import AUY83.Formulas
import AUY83.GF2

open Finset

namespace AUY83

/-! ## Lower Bound Techniques for Fixed Partitions (Section V) -/

/-- A **fooling set for a Boolean function** relative to a variable partition. -/
structure BoolFoolingSet (n : ℕ) (f : BoolFun n) (S : VarPartition n) where
  pairs : Finset (Assignment n × Assignment n)
  all_one : ∀ p ∈ pairs, commMatrixEntry f S p.1 p.2 = true
  fooling : ∀ p ∈ pairs, ∀ q ∈ pairs, p ≠ q →
    commMatrixEntry f S p.1 q.2 = false ∨
    commMatrixEntry f S q.1 p.2 = false

/-
**Fooling set lower bound**: The size of any fooling set is a lower bound on `pp₁(f, S)`.
-/
theorem foolingSet_le_pp₁ {n : ℕ} (f : BoolFun n) (S : VarPartition n)
    (F : BoolFoolingSet n f S) :
    F.pairs.card ≤ pp₁ f S := by
  refine' le_csInf _ _;
  · obtain ⟨C, hC⟩ : ∃ C : RectangleCover n f S, True := by
      obtain ⟨T, hT⟩ : ∃ T : ProtocolTree n, T.IsValid f S := by
        obtain ⟨F, hF⟩ := exists_formula n f;
        exact ⟨ _, formulaToProtocol_valid F f S hF ⟩;
      exact ⟨ protocol_gives_rectangle_cover f S T hT |> Classical.choose, trivial ⟩;
    exact ⟨ _, ⟨ C, rfl ⟩ ⟩;
  · simp +zetaDelta at *;
    intro C
    set rects := C.rects with hrects;
    have h_distinct_rectangles : ∀ p ∈ F.pairs, ∃ R ∈ rects, R.covers p.1 p.2 ∧ ∀ q ∈ F.pairs, q ≠ p → ¬(R.covers q.1 q.2) := by
      intro p hp
      obtain ⟨R, hR⟩ : ∃ R ∈ rects, R.covers p.1 p.2 := by
        exact C.covers_all p.1 p.2 ( F.all_one p hp );
      use R;
      simp_all +decide [ Rectangle.covers ];
      intro a b hab hne ha hb; have := F.fooling _ hp _ hab; simp_all +decide [ commMatrixEntry ] ;
      have := C.all_one R hR.1; simp_all +decide [ Rectangle.IsOne ] ;
      have := F.all_one p hp; have := F.all_one ( a, b ) hab; simp_all +decide [ Rectangle.IsMonochromatic ] ;
      unfold commMatrixEntry at this; aesop;
    choose! R hR₁ hR₂ hR₃ using h_distinct_rectangles;
    have h_distinct_rectangles : Finset.card (Finset.image (fun p => R p.1 p.2) (Finset.attach F.pairs)) ≤ Finset.card rects := by
      exact Finset.card_le_card ( Finset.image_subset_iff.mpr fun p hp => hR₁ _ _ );
    rwa [ Finset.card_image_of_injOn, Finset.card_attach ] at h_distinct_rectangles;
    intro p hp q hq; specialize hR₃ p.1 p.2 q.1 q.2; aesop;

/-! ## The Rectangle Property -/

theorem rectangle_monochromatic {n : ℕ} (f : BoolFun n) (S : VarPartition n)
    (R : Rectangle n) (hR : R.IsOne f S)
    {αL : Assignment n} (hαL : αL ∈ R.rowSet)
    {αR : Assignment n} (hαR : αR ∈ R.colSet) :
    commMatrixEntry f S αL αR = true :=
  hR αL hαL αR hαR

theorem cover_covers_ones {n : ℕ} (f : BoolFun n) (S : VarPartition n)
    (C : RectangleCover n f S) (αL αR : Assignment n)
    (h : isOneEntry f S αL αR) :
    ∃ R ∈ C.rects, R.covers αL αR :=
  C.covers_all αL αR h

/-! ## Rank Lower Bound -/

noncomputable def boolRank {n : ℕ} (f : BoolFun n) (S : VarPartition n) : ℕ :=
  pp₁ f S

/-! ## Theorem 11: Lower Bound for Inner Product -/

/-- The **inner product function** on `2k` variables. -/
def innerProduct (k : ℕ) : BoolFun (2 * k) :=
  fun α =>
    let go : Fin k → Bool := fun i =>
      α ⟨i.val, by omega⟩ && α ⟨k + i.val, by omega⟩
    let vals : List Bool := (List.finRange k).map go
    vals.foldl xor false

/-- The natural partition for the inner product function. -/
def innerProductPartition (k : ℕ) : VarPartition (2 * k) where
  left := (Finset.univ : Finset (Fin k)).map ⟨fun i => ⟨i.val, by omega⟩, by
    intro a b hab; simp [Fin.ext_iff] at hab; exact Fin.ext hab⟩

/-! ### Fooling Set Construction for Inner Product -/

/-- The i-th left basis assignment: variable i is true, all others false. -/
def ipLeftAssign (k : ℕ) (i : Fin k) : Assignment (2 * k) :=
  fun j => j.val = i.val

/-- The i-th right basis assignment: variable k+i is true, all others false. -/
def ipRightAssign (k : ℕ) (i : Fin k) : Assignment (2 * k) :=
  fun j => j.val = k + i.val

/-
The inner product of two basis assignments (eᵢ, eⱼ) is 1 iff i = j.
    This is because combineAssign takes left vars from eᵢ and right vars from eⱼ,
    giving variable m = (m == i) for m < k and variable k+m = (m == j) for m < k.
    Then IP = XOR of ((m==i) AND (m==j)) = (i == j).
-/
theorem innerProduct_basis_diag (k : ℕ) (i : Fin k) :
    commMatrixEntry (innerProduct k) (innerProductPartition k)
      (ipLeftAssign k i) (ipRightAssign k i) = true := by
  unfold commMatrixEntry innerProduct innerProductPartition ipLeftAssign ipRightAssign;
  unfold combineAssign; simp +decide [ List.foldl_map ] ;
  induction' k with k ih <;> simp_all +decide [ List.finRange_succ ];
  · exact Fin.elim0 i;
  · cases i using Fin.inductionOn <;> simp_all +decide [ List.foldl_map ];
    · exact fun x => ne_of_lt ( Nat.lt_succ_of_le ( Fin.is_le x ) );
    · convert ih ‹_› using 1;
      grind +revert

/-
The inner product of two different basis assignments is 0.
-/
theorem innerProduct_basis_offdiag (k : ℕ) (i j : Fin k) (hij : i ≠ j) :
    commMatrixEntry (innerProduct k) (innerProductPartition k)
      (ipLeftAssign k i) (ipRightAssign k j) = false := by
  -- Apply the fact that the XOR of all false values is false.
  have h_xor_false : ∀ (l : List Bool), (∀ b ∈ l, b = false) → l.foldl xor false = false := by
    intro l hl; induction l <;> aesop;
  convert h_xor_false _ _;
  simp +decide [ combineAssign, ipLeftAssign, ipRightAssign ];
  grind

/-- The fooling set for the inner product: k pairs {(eᵢ, eᵢ)}_{i=0}^{k-1}. -/
noncomputable def innerProductFoolingSet (k : ℕ) :
    BoolFoolingSet (2 * k) (innerProduct k) (innerProductPartition k) where
  pairs := (Finset.univ : Finset (Fin k)).image
    (fun i => (ipLeftAssign k i, ipRightAssign k i))
  all_one := by
    intro p hp
    simp at hp
    obtain ⟨i, rfl⟩ := hp
    exact innerProduct_basis_diag k i
  fooling := by
    intro p hp q hq hpq
    simp at hp hq
    obtain ⟨i, rfl⟩ := hp
    obtain ⟨j, rfl⟩ := hq
    have hij : i ≠ j := by
      intro h; subst h; exact hpq rfl
    left
    exact innerProduct_basis_offdiag k i j hij

theorem innerProductFoolingSet_card (k : ℕ) :
    (innerProductFoolingSet k).pairs.card = k := by
  erw [ Finset.card_image_of_injective ];
  · simp +decide [ Finset.card_univ ];
  · intro i j h;
    simp_all +decide [ Fin.ext_iff, funext_iff ];
    have := h.1 ⟨ i, by linarith [ Fin.is_lt i ] ⟩ ; have := h.1 ⟨ j, by linarith [ Fin.is_lt j ] ⟩ ; simp_all +decide [ ipLeftAssign ] ;

/-- **Theorem 11**: The inner product function has protocol partition number
    `pp₁(IP_k, S) ≥ k`.
    This follows from the fooling set of size k consisting of standard basis pairs.
    The paper proves the stronger bound pp₁ ≥ 2^k - 1 using a counting argument
    (each 1-rectangle covers at most 2^(k-1) entries, and there are (2^k-1)·2^(k-1)
    total 1-entries), but the fooling set lower bound suffices for the key conclusion
    that the inner product has super-linear communication complexity. -/
theorem innerProduct_pp₁_lower_bound (k : ℕ) (_hk : 0 < k) :
    pp₁ (innerProduct k) (innerProductPartition k) ≥ k := by
  calc k = (innerProductFoolingSet k).pairs.card := (innerProductFoolingSet_card k).symm
    _ ≤ pp₁ (innerProduct k) (innerProductPartition k) :=
        foolingSet_le_pp₁ _ _ _

/-! ### Counting Argument for Strong Inner Product Lower Bound -/

/-- The set of all 1-entries for the inner product function. -/
def ipOneEntries (k : ℕ) : Finset (Assignment (2 * k) × Assignment (2 * k)) :=
  Finset.univ.filter (fun p => commMatrixEntry (innerProduct k) (innerProductPartition k) p.1 p.2 = true)

/-! #### Projection to relevant bits -/

/-- Extract the "left-relevant" bits from a full assignment: positions 0..k-1. -/
def leftBits (k : ℕ) (α : Assignment (2 * k)) : Fin k → Bool :=
  fun i => α ⟨i.val, by omega⟩

/-- Extract the "right-relevant" bits from a full assignment: positions k..2k-1. -/
def rightBits (k : ℕ) (α : Assignment (2 * k)) : Fin k → Bool :=
  fun i => α ⟨k + i.val, by omega⟩

/-
The communication matrix entry for the inner product depends only on the
    left-relevant bits of αL and the right-relevant bits of αR.
-/
theorem commMatrixEntry_ip_eq_gf2IP (k : ℕ) (αL αR : Assignment (2 * k)) :
    commMatrixEntry (innerProduct k) (innerProductPartition k) αL αR =
    gf2IP k (leftBits k αL) (rightBits k αR) := by
  unfold commMatrixEntry innerProduct innerProductPartition combineAssign leftBits rightBits gf2IP
  simp +decide [ Finset.mem_univ ]
  rw [ List.foldl_map ];
  grind

/-
Lemma 1 (Corrected): Total number of 1-entries.

    The original statement had `(2^k - 1) * 2^(k-1)` which counts only the
    "relevant" bit pairs. The full count includes a factor of `2^(2*k)` for
    the irrelevant bits (k bits from αL's right side and k bits from αR's left side).
-/
theorem ip_ones_count (k : ℕ) (hk : 0 < k) :
    (ipOneEntries k).card = (2^k - 1) * 2^(k - 1) * 2^(2 * k) := by
  rw [ ← gf2_ones_count k hk ];
  -- Let's rewrite the set we're interested in in terms of the GF(2) inner product.
  have h_set_eq : ipOneEntries k = Finset.image (fun (p : ((Fin k → Bool) × (Fin k → Bool)) × ((Fin k → Bool) × (Fin k → Bool))) => ((fun i : Fin (2 * k) => if h : i.val < k then p.1.1 ⟨i.val, h⟩ else p.2.1 ⟨i.val - k, by omega⟩), (fun i : Fin (2 * k) => if h : i.val < k then p.2.2 ⟨i.val, h⟩ else p.1.2 ⟨i.val - k, by omega⟩))) (gf2OnePairs k ×ˢ (Finset.univ : Finset ((Fin k → Bool) × (Fin k → Bool)))) := by
    ext ⟨αL, αR⟩;
    constructor;
    · intro h;
      simp_all +decide [ ipOneEntries ];
      refine' ⟨ leftBits k αL, rightBits k αR, _, _, _ ⟩;
      · exact Finset.mem_filter.mpr ⟨ Finset.mem_univ _, by simpa [ commMatrixEntry_ip_eq_gf2IP ] using h ⟩;
      · use fun i => αL ⟨k + i.val, by omega⟩;
        ext i; simp [leftBits];
        exact fun hi => by congr; simp +decide [ Nat.add_sub_of_le hi ] ;
      · use fun i => αR ⟨i.val, by omega⟩;
        ext i; simp [rightBits];
        exact fun hi => by congr; simp +decide [ add_tsub_cancel_of_le hi ] ;
    · simp +zetaDelta at *;
      rintro x y hxy z hx w hy; subst_vars; simp_all +decide [ ipOneEntries, gf2OnePairs ] ;
      convert hxy using 1;
      convert commMatrixEntry_ip_eq_gf2IP k _ _ using 2;
      · exact funext fun i => by unfold leftBits; aesop;
      · ext i; simp +decide [ rightBits ] ;
  rw [ h_set_eq, Finset.card_image_of_injOn, Finset.card_product, Finset.card_univ ];
  · norm_num [ pow_mul' ];
    exact Or.inl <| by ring;
  · intro p hp q hq h_eq; simp_all +decide [ funext_iff ] ;
    ext ⟨i, hi⟩; have := h_eq.1 ⟨i, by linarith⟩; have := h_eq.2 ⟨i, by linarith⟩; aesop;
    · have := h_eq.2 ⟨ k + i, by linarith ⟩ ; aesop;
    · simpa [ hi ] using h_eq.1 ⟨ i + k, by linarith ⟩;
    · simpa [ hi ] using h_eq.2 ⟨ i, by linarith ⟩

theorem ip_rect_max_size (k : ℕ) (hk : 0 < k) (R : Rectangle (2 * k))
    (hR : R.IsOne (innerProduct k) (innerProductPartition k)) :
    R.rowSet.card * R.colSet.card ≤ 2 ^ (k - 1) * 2 ^ (2 * k) := by
  -- By definition of $A$ and $B$, we know that $R.rowSet$ is a subset of ${αL | leftBits k αL ∈ A}$ which has exactly $|A| * 2^k$ elements.
  have hA : R.rowSet.card ≤ (Finset.image (leftBits k) R.rowSet).card * 2 ^ k := by
    have hA : ∀ a ∈ Finset.image (leftBits k) R.rowSet, (Finset.filter (fun αL => leftBits k αL = a) R.rowSet).card ≤ 2 ^ k := by
      intros a ha
      have hA_fiber : ∀ αL ∈ Finset.filter (fun αL => leftBits k αL = a) R.rowSet, ∀ αR ∈ Finset.filter (fun αL => leftBits k αL = a) R.rowSet, αL ≠ αR → rightBits k αL ≠ rightBits k αR := by
        intros αL hαL αR hαR hne h_eq_rightBits
        have h_eq_leftBits : leftBits k αL = leftBits k αR := by
          aesop
        have h_eq : αL = αR := by
          ext i; by_cases hi : i.val < k <;> simp_all +decide;
          · convert congr_fun h_eq_leftBits ⟨ i, hi ⟩ |> Eq.trans <| congr_fun hαL.2 ⟨ i, hi ⟩ |> Eq.symm using 1;
          · have := congr_fun h_eq_rightBits ⟨ i - k, by omega ⟩ ; simp_all +decide [ rightBits ] ;
        contradiction;
      have hA_fiber : Finset.card (Finset.image (rightBits k) (Finset.filter (fun αL => leftBits k αL = a) R.rowSet)) ≤ 2 ^ k := by
        exact le_trans ( Finset.card_le_univ _ ) ( by norm_num [ Finset.card_univ ] );
      rwa [ Finset.card_image_of_injOn fun x hx y hy hxy => Classical.not_not.1 fun h => ‹∀ αL ∈ { αL ∈ R.rowSet | leftBits k αL = a }, ∀ αR ∈ { αL ∈ R.rowSet | leftBits k αL = a }, αL ≠ αR → rightBits k αL ≠ rightBits k αR› x hx y hy h hxy ] at hA_fiber;
    convert Finset.sum_le_sum hA using 1;
    · rw [ Finset.card_eq_sum_ones, Finset.sum_image' ] ; aesop;
    · norm_num
  have hB : R.colSet.card ≤ (Finset.image (rightBits k) R.colSet).card * 2 ^ k := by
    have hB : ∀ b ∈ Finset.image (rightBits k) R.colSet, (Finset.filter (fun α => rightBits k α = b) R.colSet).card ≤ 2 ^ k := by
      intros b hb
      have h_irrelevant : Finset.card (Finset.filter (fun α => rightBits k α = b) Finset.univ) = 2 ^ k := by
        have hB : Finset.card (Finset.filter (fun α => rightBits k α = b) (Finset.univ : Finset (Fin (2 * k) → Bool))) = Finset.card (Finset.univ : Finset (Fin k → Bool)) := by
          fapply Finset.card_bij;
          use fun a ha => leftBits k a;
          · aesop;
          · simp +contextual [ funext_iff, leftBits, rightBits ];
            intro a₁ ha₁ a₂ ha₂ h x
            cases' lt_or_ge x.val k with hx hx
            · exact h ⟨ x.val, by linarith ⟩
            · have := ha₁ ⟨ x.val - k, by omega ⟩
              have := ha₂ ⟨ x.val - k, by omega ⟩
              simp_all +decide [ add_tsub_cancel_of_le (show k ≤ x.val from by linarith) ]
          · intro a ha
            use fun i => if h : i.val < k then a ⟨i.val, h⟩ else b ⟨i.val - k, by
              rw [ tsub_lt_iff_left ] <;> linarith [ Fin.is_lt i ]⟩
            generalize_proofs at *;
            unfold rightBits leftBits; aesop;
        simp_all +decide [ Finset.card_univ ];
      exact h_irrelevant ▸ Finset.card_le_card fun x hx => by aesop;
    have hB : R.colSet.card = Finset.sum (Finset.image (rightBits k) R.colSet) (fun b => (Finset.filter (fun α => rightBits k α = b) R.colSet).card) := by
      rw [ Finset.card_eq_sum_ones, Finset.sum_image' ] ; aesop;
    exact hB.symm ▸ le_trans ( Finset.sum_le_sum ‹_› ) ( by norm_num );
  -- By definition of $A$ and $B$, we know that $A.card * B.card ≤ 2^(k-1)$.
  have hAB : (Finset.image (leftBits k) R.rowSet).card * (Finset.image (rightBits k) R.colSet).card ≤ 2 ^ (k - 1) := by
    have hAB : ∀ a ∈ Finset.image (leftBits k) R.rowSet, ∀ b ∈ Finset.image (rightBits k) R.colSet, gf2IP k a b = true := by
      intro a ha b hb;
      obtain ⟨ αL, hαL, rfl ⟩ := Finset.mem_image.mp ha; obtain ⟨ αR, hαR, rfl ⟩ := Finset.mem_image.mp hb; exact commMatrixEntry_ip_eq_gf2IP k αL αR ▸ hR αL hαL αR hαR;
    by_cases hA : Finset.Nonempty ( Finset.image ( leftBits k ) R.rowSet ) <;> by_cases hB : Finset.Nonempty ( Finset.image ( rightBits k ) R.colSet ) <;> simp_all +decide [ gf2_rect_bound ];
  exact le_trans ( Nat.mul_le_mul hA hB ) ( by convert Nat.mul_le_mul_right ( 2 ^ ( 2 * k ) ) hAB using 1 ; ring )

open Classical in
theorem cover_size_counting {n : ℕ} (f : BoolFun n) (S : VarPartition n)
    (C : RectangleCover n f S) (total_ones max_rect : ℕ)
    (h_total : (Finset.univ.filter (fun p : Assignment n × Assignment n => isOneEntry f S p.1 p.2)).card = total_ones)
    (h_max : ∀ R ∈ C.rects, R.rowSet.card * R.colSet.card ≤ max_rect)
    (h_max_pos : 0 < max_rect) :
    C.rects.card ≥ total_ones / max_rect := by
  have _h_max_pos := h_max_pos
  have h_covered : (Finset.univ.filter (fun p : Assignment n × Assignment n => isOneEntry f S p.1 p.2)) ⊆ Finset.biUnion C.rects (fun R => R.rowSet ×ˢ R.colSet) := by
    intro p hp; have := C.covers_all p.1 p.2; aesop;
  refine' Nat.div_le_of_le_mul _;
  exact h_total ▸ le_trans ( Finset.card_le_card h_covered ) ( le_trans ( Finset.card_biUnion_le ) ( by simpa [ mul_comm ] using Finset.sum_le_sum fun x hx => h_max x hx ) )

/-
**Theorem 11 (Strong)**: The inner product function has protocol partition number
    `pp₁(IP_k, S) ≥ 2^k - 1`.
    This exactly matches the paper's exponential lower bound.
-/
theorem innerProduct_pp₁_lower_bound_strong (k : ℕ) (hk : 0 < k) :
    pp₁ (innerProduct k) (innerProductPartition k) ≥ 2^k - 1 := by
  refine' le_csInf _ _;
  · obtain ⟨C, hC⟩ : ∃ C : RectangleCover (2 * k) (innerProduct k) (innerProductPartition k), True := by
      obtain ⟨ F, hF ⟩ := exists_formula ( 2 * k ) ( innerProduct k );
      have := formulaToProtocol_valid F ( innerProduct k ) ( innerProductPartition k ) hF;
      exact ⟨ protocol_gives_rectangle_cover ( innerProduct k ) ( innerProductPartition k ) _ this |> Classical.choose, trivial ⟩;
    exact ⟨ _, ⟨ C, rfl ⟩ ⟩;
  · rintro _ ⟨ C, rfl ⟩;
    have := cover_size_counting ( innerProduct k ) ( innerProductPartition k ) C ( ( 2 ^ k - 1 ) * 2 ^ ( k - 1 ) * 2 ^ ( 2 * k ) ) ( 2 ^ ( k - 1 ) * 2 ^ ( 2 * k ) ) ?_ ?_ ?_;
    · exact le_trans ( by rw [ Nat.mul_assoc, Nat.mul_div_cancel _ ( by positivity ) ] ) this;
    · convert ip_ones_count k hk using 1;
      unfold ipOneEntries isOneEntry;
      convert rfl;
    · exact fun R hR => ip_rect_max_size k hk R ( C.all_one R hR );
    · positivity

end AUY83