APSP in $O(n^{2.99791})$: definitions, part 2 of 5

/-
Copyright (c) 2026 Anthropic, PBC. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
SPDX-License-Identifier: Apache-2.0
-/
-- Generated by skeleton subtraction from the ThreeSumApsp library at commit 27a2fa4bab32.

import Definitions.Def_APSP299791_Defs1

set_option autoImplicit false
set_option relaxedAutoImplicit false

section
-- Source module: ThreeSumApsp.Sec4.Theorem30
/-!
# Equation (7) and Theorem 30: the data structure (Section 4.3)

The running times of Theorem 30 on the word RAM are `wordRam_theorem_30` and
`wordRam_theorem_30_wanted`; the expressions inside their `O(·)` are `cost8`, `cost9` and
`costQuery`. This file has the mathematics of the paper's proof, in the paper's order.

* *The decay rate.* `ρ = 9m/(L-m+1)` is less than 1 (`sec4_rho_lt_one`), the ratio `β_d/β_{d-1}` is
  at most `ρ` (`eq_7_ratio`), and so `β_d ≤ ρ^d M` (`eq_7`). This is equation (7).
* *Preprocessing: the count behind (8).* Here `sqrtKN0 L m` is `√K N₀`. Padding at most doubles `N`
  (`Theorem30.padding`; nothing else rests on this lemma). The list of subsets takes `K L ≤ 10^L`
  operations (`Theorem30.subsets`), which the last term of (8) absorbs
  (`Theorem30.subsets_absorbed`). There are at most `4N/(√K N₀)` bands (`Theorem30.bands`), with `N`
  the given size, by a slightly finer count than `K₀ ≥ √K/2` gives (`Theorem30.numBands_mul_le`);
  the input array of a band has at most `K N₀ D ≤ 7^L` nonzero entries
  (`Theorem30.K_mul_N0_mul_D_le`), and `L · 7^L ≤ 2 · 10^L` (`Theorem30.form_array`). There are at
  most `4N²/M` tiles (`Theorem30.tiles`) with at most `(m+1) ∑_{d ≥ t} β_d` boxes each (Lemma 29),
  and `∑_{d ≥ t} β_d ≤ M ρ^t/(1-ρ)` by (7) (`Theorem30.sum_beta`), which bounds the number of all
  boxes (`Theorem30.boxes_total`). The boxes (`Theorem30.boxes_cost`), the bands
  (`Theorem30.bands_cost`) and the list add up to at most a constant times the expression in (8)
  (`Theorem30.cost8_assembly`).
* *Query.* The sum of Lemma 28 for the output string of the position `(I, J)` is `(XY)[I, J]` by
  Section 2.4.4 (`Theorem30.query`). Each box of that sum is a box of the tile (`Theorem30.lookup`),
  and the dynamic program of Lemma 29 has stored its value (`dpValue_card_starLevels`); so the query
  returns `(XY)[I, J]` (`Theorem30.correct`). Expression (9) is `|W|` queries on top of (8)
  (`Theorem30.cost9_eq`).
* *Word size.* `10^L ≤ N^{5/2}` (`Theorem30.ten_pow_le`).

Here `N` is the given size throughout, and the padded size is `padN L m N`.
-/

@[expose] public section

open Finset

namespace ThreeSumApsp

/-! ### The decay rate `ρ` and equation (7) -/

/-- Section 4.3: `ρ = 9m/(L-m+1)` "is less than 1 since L ≥ 10m". -/
theorem sec4_rho_lt_one (L m : ℕ) (hL : 10 * m ≤ L) : ThreeSumApsp.rho L m < 1 := by
  have hL' : (10 : ℝ) * m ≤ L := by exact_mod_cast hL
  have hm : (0 : ℝ) ≤ m := Nat.cast_nonneg m
  rw [ThreeSumApsp.rho, div_lt_one (by linarith)]
  linarith

/-- The decay rate `ρ` is not negative when `L ≥ 10m`. -/
theorem rho_nonneg {L m : ℕ} (hL : 10 * m ≤ L) : 0 ≤ ThreeSumApsp.rho L m := by
  have hL' : (10 : ℝ) * m ≤ L := by exact_mod_cast hL
  have hm : (0 : ℝ) ≤ m := Nat.cast_nonneg m
  exact div_nonneg (by positivity) (by linarith)

/-- Equation (7), first half: for `1 ≤ d ≤ m`,
`β_d / β_{d-1} = 9(m-d+1)/(L-m+d) ≤ 9m/(L-m+1) = ρ`. -/
theorem eq_7_ratio_b299791 (L m d : ℕ) (hL : 10 * m ≤ L) (hd1 : 1 ≤ d) (hdm : d ≤ m) :
    (beta L m d : ℝ) / (beta L m (d - 1) : ℝ)
        = 9 * ((m : ℝ) - (d : ℝ) + 1) / ((L : ℝ) - (m : ℝ) + (d : ℝ)) ∧
      9 * ((m : ℝ) - (d : ℝ) + 1) / ((L : ℝ) - (m : ℝ) + (d : ℝ)) ≤ ThreeSumApsp.rho L m := by
  have hL' : (10 : ℝ) * m ≤ L := by exact_mod_cast hL
  have hd1' : (1 : ℝ) ≤ d := by exact_mod_cast hd1
  have hdm' : (d : ℝ) ≤ m := by exact_mod_cast hdm
  constructor
  · -- the ratio was computed for (5), in the rational numbers
    have hratio := congrArg (fun q : ℚ => (q : ℝ))
      (eq_5_ratio (L := L) (m := m) (by omega) d hd1 hdm)
    push_cast at hratio
    exact hratio
  · -- the numerator is at most `9m` and the denominator at least `L - m + 1`
    exact div_le_div₀ (by positivity) (by linarith) (by linarith) (by linarith)

/-- Equation (7), second half: "so β_d ≤ ρ^d M" (for `0 ≤ d ≤ m`). -/
theorem eq_7 (L m d : ℕ) (hL : 10 * m ≤ L) (hdm : d ≤ m) :
    (beta L m d : ℝ) ≤ ThreeSumApsp.rho L m ^ d * (ThreeSumApsp.M L m : ℝ) := by
  induction d with
  | zero => simp [beta_zero_eq_M]
  | succ d ih =>
    obtain ⟨hratio, hle⟩ := eq_7_ratio_b299791 L m (d + 1) hL (by omega) hdm
    have hpos : (0 : ℝ) < (beta L m d : ℝ) := by
      exact_mod_cast beta_pos (L := L) (m := m) (by omega) d
    rw [← hratio, Nat.add_sub_cancel, div_le_iff₀ hpos] at hle
    calc (beta L m (d + 1) : ℝ) ≤ ThreeSumApsp.rho L m * (beta L m d : ℝ) := hle
      _ ≤ ThreeSumApsp.rho L m * (ThreeSumApsp.rho L m ^ d * (ThreeSumApsp.M L m : ℝ)) :=
          mul_le_mul_of_nonneg_left (ih (by omega)) (rho_nonneg hL)
      _ = ThreeSumApsp.rho L m ^ (d + 1) * (ThreeSumApsp.M L m : ℝ) := by ring

/-! ### Theorem 30, "Preprocessing": the count behind (8) -/

/-- `√K N₀` is positive when `m ≤ L`. -/
theorem sqrtKN0_pos {L m : ℕ} (hmL : m ≤ L) : 0 < sqrtKN0 L m := by
  have hK : (0 : ℝ) < (ThreeSumApsp.K L m : ℝ) := by exact_mod_cast Nat.choose_pos hmL
  have hN0 : (0 : ℝ) < (ThreeSumApsp.N0 L m : ℝ) := by exact_mod_cast ThreeSumApsp.N0_pos L m
  exact mul_pos (Real.sqrt_pos.mpr hK) hN0

/-- `(√K N₀)² = K N₀² = M`. -/
theorem sqrtKN0_sq (L m : ℕ) : sqrtKN0 L m ^ 2 = (ThreeSumApsp.M L m : ℝ) := by
  rw [sqrtKN0, mul_pow, Real.sq_sqrt (Nat.cast_nonneg _)]
  simp [ThreeSumApsp.M]

/-! ### Theorem 30, "Query": a query returns `(XY)[I, J]` -/

/-! ### Theorem 30, "Word size" -/

end ThreeSumApsp
end
end

section
-- Source module: ThreeSumApsp.Spec.Sec4.Theorem30.Retained
/-!
# The retained leaves: what the data structure reads of the encodings

The data structure reads the two encodings of a tile only through `leafProduct`, and only at leaves
with at most m digits 9 (the retained leaves):

* the value of a box without stars (`boxValue` at e = 0) is read at a member of
  `starBoxes L m t 0`, a string with at most m - t nines (`retained_of_mem_starBoxes_zero`); the
  values of the boxes with stars are looked up in the trie;
* a query reads the products at the members of `lowList`, strings with at most m nines
  (`retained_of_mem_lowList`).

So two pairs of arrays that agree at the retained leaves (`RetAgree`) give the same tries
(`fillTrie_congr`, `fillUpTo_congr`, `tileTrie_congr`, `allTries_congr`,
`allTries_tileList_congr`) and the same queries (`queryTerms_congr`, `trieQuery_congr`).  An encoding with the other leaves set
to 0 (`clipT`) agrees with it at the retained leaves; bounds at the retained leaves are bounds of the
clipped encoding everywhere (`abs_clipT_le`).
-/

@[expose] public section

namespace ThreeSumApsp.Spec

/-- Two arrays on the leaves (indexed by the codes) agree at the leaves with at most r nines. -/
def RetAgree (L r : ℕ) (a a' : List ℤ) : Prop :=
  ∀ l : List ℕ, l.length = L → (∀ d ∈ l, d < 10) → l.count 9 ≤ r →
    a.getD (ofDigitList 10 l) 0 = a'.getD (ofDigitList 10 l) 0

/-! ## The strings that are read are retained -/

/-! ## Congruence -/

section

variable {L m : ℕ} {a a' b b' : List ℤ} (hA : RetAgree L m a a') (hB : RetAgree L m b b')

include hA hB

end

/-- Two tiles whose arrays agree at the retained leaves. -/
def TileEnc.Agree (L m : ℕ) (x y : TileEnc) : Prop :=
  Spec.RetAgree L m x.encA y.encA ∧ Spec.RetAgree L m x.encB y.encB

/-! ## The clipped encoding -/

/-- The encoding with the leaves of more than r nines set to 0. -/
def clipT {L : ℕ} (r : ℕ) (f : Leaf L → ℤ) : Leaf L → ℤ :=
  fun τ => if (digitsT τ).count 9 ≤ r then f τ else 0

end ThreeSumApsp.Spec
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.Theorem30.Contracts
/-!
# Theorem 30 in the light language: procedure numbers, time functions, specifications

This file contains no program, and its only proof is a short remark on the trie area when other
cells change. It states the specifications of the procedures 40 to 53 (`nineFirst` to `queryCore`).
The proof of a routine uses only the specifications of the routines that it calls.

* A routine takes its scalars and the base addresses of its arrays as arguments; no routine but
  `preCore` and `queryAt` knows the memory map. The routines for a tile assume the order in which
  their areas lie (`TileCtxP.layout`).
* The specification of a routine is a proposition XSpec lim P: "procedure number Proc.x of the
  program P, started on these arguments in a memory that satisfies this, ends within tX steps with
  this result in a memory that satisfies that". The file of x proves XSpec for every program P that
  holds the body of x at the number Proc.x and meets the specifications of the routines that x
  calls. A caller assumes XSpec and feeds it to the rule for calls. XSpec holds at every depth d of
  calls that leaves the levels which x needs below itself.
* Conventions of the programs: the result of a procedure is what it leaves in its local 0; locals
  that are not arguments start at 0. In the proofs, lines such as "have hw := C.std.space_le" or
  "obtain ⟨⟩ := areas p t b0" state the facts on the limits and on the map that the later steps use.
* The larger routines have two records: XArgs holds the arguments and the data behind them, and
  XPre lim μ x says what the routine assumes about the limits and the memory.
* The pure models are those of the specifications of Section 4 (strings are lists of digits, level 1
  first: P_ij is 3(i - 1) + (j - 1), P₀ is 9, the star is 10).
* Time functions are definitions. The time function of a caller is written in terms of the time
  functions of its callees. Their constants are upper bounds, some with room to spare; only the
  shape matters. The lemmas `within8_tAllTiles` and `exists_tQueryCore_le` compare `tAllTiles` and
  `tQueryCore` with the expressions (8) and L ∑ α_d of the paper.
* Tries use addresses relative to the base tr of the trie area: the cell tr + a of the memory is the
  cell a of the trie array T. The area has cap cells, of which the first T.length are in use;
  nothing is assumed about the others (a new vertex clears its eleven cells); the number T.length
  is kept in the cell fp outside the area.

The routines, the pure functions that model them, and the lemmas that say what the models compute:

| number | routine | model | what is proved about the model |
|---|---|---|---|
| 40 | `nineFirst` | `nineFirst` | `head?_nineStrs`: it is the first string of `nineStrs` |
| 41 | `nineNext` | `nineNext` | `nineNext_getElem`: from string i of `nineStrs` to string i + 1 |
| 42 | `scatter` | `scatter`, `starRunIf` | `lowList` and `boxesOf` are built from them |
| 43 | `starFirst` | `starFirst`, `lastStar` | `starBoxes` is `nineStrs` mapped by `starFirst` |
| 44 | `horner` | `ofDigitList 10` | `ofDigitList_digitsT`: the code of a leaf |
| 45 | `lookup` | `trieLookup` | `TrieRep.lookup`, `TrieRep.walkOK` |
| 46 | `insert` | `trieInsert` | `TrieRep.insert`, `TrieRep.insertOK` |
| 47 | `newRoot` | `trieNew` | `TrieRep.new` |
| 48 | `sumTen` | sum of `tenValues` | `exists_lastStar_of_mem_starBoxes`, `abs_dp_partial_sum_le` |
| 49 | `fillList` | `fillTrie` | `fillTrie_rep`, `length_fillTrie_le` |
| 50 | `tile` | `tileTrie` | `fillUpTo_rep`, `length_tileTrie_le` |
| 51 | `allTiles` | `allTries` | `allTries_rep`, `length_allTries_le`, `lookup_allTries` |
| 52 | `outDigits` | `OutDigitsArgs.digits` | `digitsO_outStrOfPos`: the string of a position |
| 53 | `queryCore` | `trieQuery`, `queryTerms` | `trieQuery_allTries`, `abs_query_partial_sum_le` |

The last line of the table continues with `Theorem30.correct`. The model `boxValue` is computed by a
part of `fillList`, which calls `sumTen` for a box with stars.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Sec4

open ThreeSumApsp.Spec

/-! ## Vocabulary -/

/-! ## Procedure numbers

The numbers 0 to 39 belong to Section 2, the numbers from 40 on to Section 4. This file fixes 40 to
53; the files of the other routines of Section 4 fix theirs. -/

namespace Proc

end Proc

/-! ## Time functions -/

/-! ## The enumeration of the strings with a bounded number of nines -/

/-! ## Strings -/

/-! ## Tries -/

/-! ## The dynamic program of Lemma 29 -/

/-- What the routines for a tile assume about their surroundings: the parameters, the two encodings
of the tile and where they are, the trie area, two scratch strings, and the limits. -/
structure TileCtxP (lim : Limits) (L m t : ℕ) (encA encB : Leaf L → ℤ)
    (aA aB tr cap fp cur box : ℕ) : Prop where
  std : Std lim
  ht : t ≤ m
  hmL : m ≤ L
  /-- Every value and every partial sum fits in a word: A and B bound the encodings at the leaves
  with at most m symbols P₀, the only ones that are read. -/
  value_le : ∃ A B : ℤ, (∀ τ, (digitsT τ).count 9 ≤ m → |encA τ| ≤ A) ∧
    (∀ τ, (digitsT τ).count 9 ≤ m → |encB τ| ≤ B) ∧ 10 ^ m * (A * B) ≤ lim.word
  pow_le : (10 : ℤ) ^ L ≤ lim.word
  /-- The two encodings, the two scratch strings, the cell fp and the trie area lie in this order,
  and the last cell of the memory lies behind them. -/
  layout : InOrder (lim.space - 1)
    [(aA, 10 ^ L), (aB, 10 ^ L), (cur, L), (box, L), (fp, 1), (tr, cap)]

namespace FillListArgs

variable (x : FillListArgs)

end FillListArgs

/-- What `fillList` assumes: the surroundings of a tile, a trie that holds the boxes with fewer
stars, the tries and the two encodings in the memory, and room for L vertices for each box. -/
structure FillListPreP (lim : Limits) (μ : ℕ → ℤ) (x : FillListArgs) : Prop where
  ctx : TileCtxP lim x.L x.m x.t x.encA x.encB x.aA x.aB x.tr x.cap x.fp x.cur x.box
  stars_le : x.e ≤ x.m - x.t
  root_ne : x.root ≠ 0
  rep : TrieRep x.L x.lo x.T x.roots (x.stored [])
  trie : TrieMem μ x.tr x.cap x.fp x.T
  segA : Seg μ x.aA (arrT x.encA)
  segB : Seg μ x.aB (arrT x.encB)
  room : x.T.length + 11 * x.L * x.boxes.length ≤ x.cap

/-- fillList(e, root, aA, aB, tr, fp, cur, box, L, mt), with mt = m - t, inserts the boxes with e
stars, with their values, into the trie of the tile, which has the given root and holds the boxes
with fewer stars; for e ≥ 1 the values are sums of ten values looked up in it. -/
def FillListSpecP (lim : Limits) (P : Program) : Prop :=
  ∀ (x : FillListArgs) (μ : ℕ → ℤ), FillListPreP lim μ x →
    ∀ d, d + 2 ≤ lim.depth →
    Meets lim P Proc.fillList d x.vals μ (tFillList x.L x.boxes.length) fun _ μ' =>
      TrieMem μ' x.tr x.cap x.fp (fillTrie (arrT x.encA) (arrT x.encB) x.root x.e x.boxes x.T) ∧
        SameOutsideTile μ μ' x.L x.tr x.cap x.fp x.cur x.box

/-- What `tile` assumes: the surroundings of a tile, the older tries and roots and the two encodings
in the memory, room for the new trie, and the place of the cells of the roots. -/
structure TilePreP (lim : Limits) (μ : ℕ → ℤ) (x : TileArgs) : Prop where
  ctx : TileCtxP lim x.L x.m x.t x.encA x.encB x.aA x.aB x.tr x.cap x.fp x.cur x.box
  rep : x.s.Holds x.L x.lo x.old
  trie : TrieMem μ x.tr x.cap x.fp x.s.cells
  roots : SegN μ x.aR x.s.roots
  segA : Seg μ x.aA (arrT x.encA)
  segB : Seg μ x.aB (arrT x.encB)
  room : x.s.cells.length + 11 * (1 + x.L * (boxes x.L x.m x.t).card) ≤ x.cap
  /-- The cells of the roots lie between the cell fp and the trie area. -/
  rootsPlace : x.fp < x.aR ∧ x.aR + x.cells ≤ x.tr

/-- `tile` builds the trie of one tile, with the values of all its boxes, on top of the older tries,
and writes its root into the cell behind the older roots. It changes only the trie area, the cell
fp, the two scratch strings and this cell. -/
def TileSpecP (lim : Limits) (P : Program) : Prop :=
  ∀ (x : TileArgs) (μ : ℕ → ℤ), TilePreP lim μ x →
    ∀ d, d + 3 ≤ lim.depth → Meets lim P Proc.tile d
      [x.aA, x.aB, (x.aR + x.s.roots.length : ℕ), x.tr, x.fp, x.cur, x.box, x.L, (x.m - x.t : ℕ)] μ
      (tTile x.L x.m x.t) fun _ μ' =>
        TrieMem μ' x.tr x.cap x.fp x.tries.cells ∧ SegN μ' x.aR x.tries.roots ∧ SameOn x.Kept μ μ'

/-- What `allTiles` assumes: the limits, room in the trie area for the tries of all tiles, the order
of the areas, the encodings in the memory, and the trie array [0]. -/
structure AllTilesPreP (lim : Limits) (μ : ℕ → ℤ) (x : AllTilesArgs) : Prop where
  std : Std lim
  t_le : x.t ≤ x.m
  m_le : x.m ≤ x.L
  /-- Every value and every partial sum fits in a word: A and B bound the encodings of the bands at
  the leaves with at most m symbols P₀, the only ones that are read. -/
  value_le : ∃ A B : ℤ, (∀ β < x.nB, ∀ τ, (digitsT τ).count 9 ≤ x.m → |x.encA β τ| ≤ A) ∧
    (∀ β < x.nB, ∀ τ, (digitsT τ).count 9 ≤ x.m → |x.encB β τ| ≤ B) ∧
    10 ^ x.m * (A * B) ≤ lim.word
  pow_le : (10 : ℤ) ^ x.L ≤ lim.word
  cap_ge : 1 + x.nB * x.nB * (11 * (1 + x.L * (boxes x.L x.m x.t).card)) ≤ x.cap
  /-- The areas lie in this order, without overlap, below the end of the memory. -/
  order : x.aENCA + x.nB * 10 ^ x.L ≤ x.aENCB ∧ x.aENCB + x.nB * 10 ^ x.L ≤ x.cur ∧
    x.cur + x.L ≤ x.box ∧ x.box + x.L ≤ x.fp ∧ x.fp < x.aR ∧
    x.aR + x.nB * x.nB ≤ x.tr ∧ x.tr + x.cap < lim.space
  segA : ∀ β < x.nB, Seg μ (x.aENCA + β * 10 ^ x.L) (arrT (x.encA β))
  segB : ∀ β < x.nB, Seg μ (x.aENCB + β * 10 ^ x.L) (arrT (x.encB β))
  trie : TrieMem μ x.tr x.cap x.fp [0]

/-- `allTiles` builds the tries of all tiles, starting from the trie array [0], and changes nothing
below cur or behind the trie area. -/
def AllTilesSpecP (lim : Limits) (P : Program) : Prop :=
  ∀ (x : AllTilesArgs) (μ : ℕ → ℤ), AllTilesPreP lim μ x →
    ∀ d, d + 4 ≤ lim.depth →
    Meets lim P Proc.allTiles d [x.aENCA, x.aENCB, x.nB, (10 ^ x.L : ℕ), x.aR, x.tr, x.fp,
      x.cur, x.box, x.L, (x.m - x.t : ℕ)] μ (tAllTiles x.L x.m x.t x.nB) fun _ μ' =>
        TrieMem μ' x.tr x.cap x.fp x.tries.cells ∧ SegN μ' x.aR x.tries.roots ∧ SameOn x.Kept μ μ'

/-! ## A query -/

namespace OutDigitsArgs

variable (x : OutDigitsArgs)

end OutDigitsArgs

/-- What `queryCore` assumes: the data stand in the memory, every box of the query is found in the
trie of the tile, all partial sums fit in a word, and the two scratch strings are apart from each
other and from the data. -/
structure QueryCorePreP (lim : Limits) (μ : ℕ → ℤ) (x : QueryCoreArgs) : Prop where
  std : Std lim
  t_le : x.t ≤ x.m
  card : (innerSetO x.η).card = x.m
  segA : Seg μ x.aA (arrT x.encA)
  segB : Seg μ x.aB (arrT x.encB)
  segT : Seg μ x.tr x.T
  segDigits : SegN μ x.wd (digitsO x.η)
  walk : ∀ l ∈ boxesOf x.m x.t (digitsO x.η), WalkOK x.T x.root l
  sums : ∀ j, |(x.terms.take j).sum| ≤ lim.word
  /-- The products at the leaves with at most m symbols P₀, the only ones that are read. -/
  prod : ∀ τ, (digitsT τ).count 9 ≤ x.m → |x.encA τ * x.encB τ| ≤ lim.word
  pow : (10 : ℤ) ^ x.L ≤ lim.word
  ss_wd : Apart x.ss x.m x.wd x.L
  box_wd : Apart x.box x.L x.wd x.L
  ss_box : Apart x.ss x.m x.box x.L
  ss_aA : Apart x.ss x.m x.aA (10 ^ x.L)
  ss_aB : Apart x.ss x.m x.aB (10 ^ x.L)
  box_aA : Apart x.box x.L x.aA (10 ^ x.L)
  box_aB : Apart x.box x.L x.aB (10 ^ x.L)
  ss_tr : Apart x.ss x.m x.tr x.T.length
  box_tr : Apart x.box x.L x.tr x.T.length
  aA_lt : x.aA + 10 ^ x.L < lim.space
  aB_lt : x.aB + 10 ^ x.L < lim.space
  tr_lt : x.tr + x.T.length < lim.space
  wd_lt : x.wd + x.L < lim.space
  ss_lt : x.ss + x.m < lim.space
  box_lt : x.box + x.L < lim.space

/-- `queryCore` returns the sum `trieQuery` for the output string whose digits are at wd: the
products at its leaves of order below t, read from the two encodings, and what the trie with the
given root holds for its boxes. This is the sum of Lemma 28 if the trie holds the values of the
boxes (`trieQuery_allTries`). The routine changes only the two scratch strings. -/
def QueryCoreSpecP (lim : Limits) (P : Program) : Prop :=
  ∀ (x : QueryCoreArgs) (μ : ℕ → ℤ), QueryCorePreP lim μ x →
    ∀ d, d + 1 ≤ lim.depth →
    Meets lim P Proc.queryCore d [x.aA, x.aB, x.tr, x.root, x.wd, x.ss, x.box, x.L, x.m, x.t] μ
      (tQueryCore x.L x.m x.t) fun r μ' =>
        r = trieQuery x.m x.t (arrT x.encA) (arrT x.encB) x.T x.root (digitsO x.η) ∧
          SameOutside2 μ μ' x.ss x.m x.box x.L

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec2.PrunedEnc.Contracts
/-!
# Pruned encodings: the contracts

The data structure of Section 4 reads the encoding of a band only at the leaves whose code has at
most m digits 9 (the boxes with no star and the strings of the low lists).  The pruned encoder
computes only these leaves: below the term number 9 (the term P₀) it recurses with one 9 fewer
allowed, and with none left it skips that term.  The other cells of the output keep what was there.

* `ninesIdx L i`: the number of digits 9 among the L decimal digits of i (most significant first).
* `retW L r`: the number of leaves with at most r nines, `∑_{j ≤ r} C(L, j) 9^{L-j}`, by its
  recursion; the pruned encoder takes `O(retW L r)` steps.
* `RetSeg L r μ a l`: the cells from a hold l at the positions with at most r nines.
* `EncodePLSpec`, `EncodePRSpec`: the entries of the pruned encoder (procedure `pEncodeP`, arguments
  `[n, src, dst, scr, tab, p7, p10, r]`).
* `SharedReadyP`, `sharedShapeP`, `SharedSpecP`: the shared stage with pruned encodings (procedure
  `pSharedP`, the same arguments as `pShared`, budget `m`).
-/

@[expose] public section

namespace Light.Sec2

open ThreeSumApsp

/-- The number of digits 9 among the `L` decimal digits of `i`, most significant first. -/
def ninesIdx : ℕ → ℕ → ℕ
  | 0, _ => 0
  | L + 1, i => (if i / 10 ^ L = 9 then 1 else 0) + ninesIdx L (i % 10 ^ L)

/-- The number of strings of `L` decimal digits with at most `r` nines. -/
def retW : ℕ → ℕ → ℕ
  | 0, _ => 1
  | L + 1, 0 => 9 * retW L 0
  | L + 1, r + 1 => 9 * retW L (r + 1) + retW L r

/-- The cells from `a` hold the list `l` at the positions whose `L` digits have at most `r`
nines. -/
def RetSeg (L r : ℕ) (μ : ℕ → ℤ) (a : ℕ) (l : List ℤ) : Prop :=
  ∀ i < l.length, ninesIdx L i ≤ r → μ (a + i) = l.getD i 0

/-! ## The numbers of the procedures (unused numbers of Section 2) -/

/-- The number of the pruned encoder. -/
def pEncodeP : ℕ := 29
/-- The number of the pruned encodeBands. -/
def pEncodeBandsP : ℕ := 30
/-- The number of the shared stage with pruned encodings. -/
def pSharedP : ℕ := 31

section Entries

variable (lim : Limits) (P : Program) (c : ℕ)

/-- encodeP(n, src, dst, scr, tab, p7, p10, r), with the coefficients φ: the encoding of the array a
at src, at the leaves with at most r nines, within c · retW n r steps. -/
def EncodePLSpec : Prop :=
  ∀ (n Lmax src dst scr tab p7 p10 r : ℕ) (μ : ℕ → ℤ) (a : LeftStr n → ℤ) (V : ℤ),
    EncodePlaces lim n Lmax src dst scr tab p7 p10 μ →
    Seg μ tab Spec.phiFlat → Seg μ src (Spec.arrL a) → (∀ u, |a u| ≤ V) →
    7 ^ (n + 1) * V ≤ lim.word →
    ∀ d, d + (n + 1) ≤ lim.depth →
    Meets lim P pEncodeP d [n, src, dst, scr, tab, p7, p10, r] μ (c * retW n r) fun _ μ' =>
      RetSeg n r μ' dst (Spec.arrT (encodingL a)) ∧ SameOutside2 μ μ' dst (10 ^ n) scr (7 ^ n)

/-- The same procedure with the coefficients ψ. -/
def EncodePRSpec : Prop :=
  ∀ (n Lmax src dst scr tab p7 p10 r : ℕ) (μ : ℕ → ℤ) (b : RightStr n → ℤ) (V : ℤ),
    EncodePlaces lim n Lmax src dst scr tab p7 p10 μ →
    Seg μ tab Spec.psiFlat → Seg μ src (Spec.arrR b) → (∀ v, |b v| ≤ V) →
    7 ^ (n + 1) * V ≤ lim.word →
    ∀ d, d + (n + 1) ≤ lim.depth →
    Meets lim P pEncodeP d [n, src, dst, scr, tab, p7, p10, r] μ (c * retW n r) fun _ μ' =>
      RetSeg n r μ' dst (Spec.arrT (encodingR b)) ∧ SameOutside2 μ μ' dst (10 ^ n) scr (7 ^ n)

end Entries

/-! ## The shared stage with pruned encodings -/

/-- What the shared stage with pruned encodings leaves behind: the tables, the directory, and the
encodings of all bands at the leaves with at most m nines. -/
structure SharedReadyP (p : Par) (hmL : p.m ≤ p.L) (aX aY b0 : ℕ)
    (X : Matrix (Fin p.N) (Fin (ThreeSumApsp.D p.m)) ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin p.N) ℤ) (μ : ℕ → ℤ) : Prop
  extends SharedTables p b0 μ where
  dir : SegN μ (p.aDIR b0) (dirList p aX aY b0)
  encA : ∀ β < p.nB, RetSeg p.L p.m μ (p.aENCA b0 + β * p.T)
    (Spec.arrT (encodingL (bandArrayL (Spec.stdLayout hmL) X β)))
  encB : ∀ β < p.nB, RetSeg p.L p.m μ (p.aENCB b0 + β * p.T)
    (Spec.arrT (encodingR (bandArrayR (Spec.stdLayout hmL) Y β)))

/-- The shape of the running time of the shared stage with pruned encodings. -/
def sharedShapeP (p : Par) : ℕ :=
  (p.L + 1) ^ 2 + (Nat.sqrt p.K + 1) + (p.KK + 1) * (p.L + 1) + (p.N + 1) * (p.Lo + 1)
    + (p.D + 1) * (p.m + 1)
    + p.nB * (retW p.L p.m + bandArrayShape p)

variable (lim : Limits) (P : Program) (c : ℕ) in
/-- sharedP(L, m, N, D, aX, aY, b0): fills the shared block, with pruned encodings. -/
def SharedSpecP : Prop :=
  ∀ (p : Par) (hmL : p.m ≤ p.L) (aX aY b0 : ℕ) (μ : ℕ → ℤ)
    (X : Matrix (Fin p.N) (Fin (ThreeSumApsp.D p.m)) ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin p.N) ℤ) (V : ℤ),
    SharedPre lim p aX aY b0 X Y V μ →
    ∀ d, d + (p.L + 3) ≤ lim.depth →
    Meets lim P pSharedP d [p.L, p.m, p.N, p.D, aX, aY, b0] μ (c * sharedShapeP p) fun _ μ' =>
      SharedReadyP p hmL aX aY b0 X Y μ' ∧ SameOutside μ μ' b0 (p.sharedEnd b0 - b0)

/-! ## Basic facts -/

end Light.Sec2
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.Theorem30.SharedP
/-!
# The shared stage with pruned encodings, as Section 4 uses it

The shared stage of Section 4 (`Sec2.SharedReadyP`) writes the encodings only at the leaves with at
most m nines, the retained leaves; the other cells of the encoding areas hold whatever was there.
The routines of Section 4 are handed what the memory holds (`memEnc`, which is a segment of the
memory by definition, `seg_memEnc`), and this agrees with the true encodings at the retained leaves
(`memEnc_eq`), which are the only ones that the data structure reads (`Spec.RetAgree`).

* The number of nines of the code of a leaf is the number of its digits 9 (`ninesIdx_ofDigitList`).
* `SharedReadyP` depends only on the cells of the shared block, without cell 31 of the directory
  (`SharedReadyP.congr`, `SharedReadyP.of_agree`), as `SharedReady` does.
-/

@[expose] public section

namespace Light.Sec2

open ThreeSumApsp ThreeSumApsp.Spec

/-! ## Nines of a code -/

/-! ## Pruned segments -/

/-! ## What the shared stage leaves behind -/

variable {p : Par} {hmL : p.m ≤ p.L} {aX aY b0 : ℕ}
  {X : Matrix (Fin p.N) (Fin (ThreeSumApsp.D p.m)) ℤ}
  {Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin p.N) ℤ} {μ μ' : ℕ → ℤ}

end Light.Sec2

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.Spec

/-! ## The encodings that the routines are handed -/

/-- The cells from a, read as an array on the leaves of depth L: what the memory holds where an
encoding should be. -/
def memEnc (μ : ℕ → ℤ) (a L : ℕ) : Leaf L → ℤ := fun τ => μ (a + codeT τ)

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.Theorem30.Memory
/-!
# Theorem 30 in the light language: the memory map and the invariant of the data structure

The data structure occupies one block of the memory, from a base address b0 on. First comes the
shared block of Section 2 (directory, tables, the encodings of all bands, written by the pruned
shared stage only at the leaves with at most m nines), and behind it the areas of Section 4:

    WD (L) | CUR (L) | BOX (L) | SS (m) | FP (1) | ROOTS (nB²) | TR (cap)

WD holds the digits of the output string of a query; CUR, BOX and SS are scratch strings; FP holds
the length of the trie array; ROOTS holds the roots of the tries, that of tile (β, β') in the cell
β nB + β'; TR is the trie area. Cell 31 of the directory holds t. `Areas` says where the tables that
the routines of Section 4 read and the areas of Section 4 lie, and `areas` proves it.

Two routines know this map: preCore(L, m, t, N, D, aX, aY, b0) builds the block from the matrices at
aX and aY, which lie below b0, and queryAt(I, J, b0) answers a query from it. They are relocatable,
so that Theorem 30, its offline form and Corollary 26 use the same two routines. The main procedures
of Theorem 30's programs only read the sizes from the input and compute the addresses. For Corollary
26 a routine of its own stands in between: it finds m and, from the threshold of the proof on, L and
t, writes padded copies of the matrices, and runs `preCore` on them.

Section 4.3: "Our data structure consists of three components: the list of the K₀² subsets Q
assigned to the block products of a tile […], the encodings of all row bands and all column bands,
and for every tile, the values of all its boxes". In the invariant DSReadyP the field shared holds
the first two, and the fields roots and trie hold the third: "for each tile we store its boxes, with
their values, in a standard trie on their strings of L symbols". The tries of all tiles lie in one
array (trie), and the table roots has the root of the trie of each tile. The invariant depends only
on the cells from b0 on, without the scratch strings (`DSReadyP.congr`).

The routines are handed the encodings as the memory holds them (`memRow`, `memCol`). These agree
with the true encodings at the leaves with at most m nines (`memRow_eq`, `memCol_eq`), the only
ones that the data structure reads, so they give the same tries (`allTries_mem_eq`).
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.Spec

/-! ## The map -/

/-! ## Where the areas lie -/

/-! ## The directory -/

namespace Dir

end Dir

/-! ## The invariant -/

/-- **The block from b0 on holds the data structure for X and Y.** The preprocessing establishes
this, and every query needs it and keeps it. -/
structure DSReadyP (p : Sec2.Par) (t : ℕ) (hmL : p.m ≤ p.L) (aX aY b0 : ℕ)
    (X : Matrix (Fin p.N) (Fin (ThreeSumApsp.D p.m)) ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin p.N) ℤ) (μ : ℕ → ℤ) : Prop where
  shared : Sec2.SharedReadyP p hmL aX aY b0 X Y μ
  cellT : μ (b0 + 31) = t
  roots : SegN μ (aROOTS p b0) (dsTries p t hmL X Y).roots
  trie : Seg μ (aTR p b0) (dsTries p t hmL X Y).cells

/-- What the memory holds where the encoding of the row band β should be. -/
def memRow (p : Sec2.Par) (b0 : ℕ) (μ : ℕ → ℤ) (β : ℕ) : Leaf p.L → ℤ :=
  memEnc μ (p.aENCA b0 + β * p.T) p.L

/-- What the memory holds where the encoding of the column band β should be. -/
def memCol (p : Sec2.Par) (b0 : ℕ) (μ : ℕ → ℤ) (β : ℕ) : Leaf p.L → ℤ :=
  memEnc μ (p.aENCB b0 + β * p.T) p.L

section

variable {p : Sec2.Par} {t : ℕ} {hmL : p.m ≤ p.L} {aX aY b0 : ℕ}
  {X : Matrix (Fin p.N) (Fin (D p.m)) ℤ} {Y : Matrix (Fin (D p.m)) (Fin p.N) ℤ} {μ μ' : ℕ → ℤ}

end

/-! ## The two routines that know the map -/

/-- The time of `preCore`; c is the constant of the shared stage (`Sec2.SharedSpecP`). -/
def tPreCoreP (c : ℕ) (p : Sec2.Par) (t : ℕ) : ℕ :=
  c * Sec2.sharedShapeP p + tAllTiles p.L p.m t p.nB + 300

/-- queryAt(I, J, b0) returns (XY)[I, J] from the block at b0, of which it changes only cells of the
scratch strings WD to SS. -/
def QueryAtSpecP (lim : Limits) (P : Program) : Prop :=
  ∀ (p : Sec2.Par) (t : ℕ) (hmL : p.m ≤ p.L) (aX aY b0 : ℕ) (X : Matrix (Fin p.N) (Fin (ThreeSumApsp.D p.m)) ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin p.N) ℤ) (U : ℤ) (μ : ℕ → ℤ) (I J : Fin p.N),
    t ≤ p.m → Lim30 lim p t b0 U → (∀ i j, |X i j| ≤ U) → (∀ i j, |Y i j| ≤ U) →
    DSReadyP p t hmL aX aY b0 X Y μ →
    ∀ d, d + 2 ≤ lim.depth →
    Meets lim P Proc.queryAt d [(I : ℕ), (J : ℕ), b0] μ (tQueryAt p.L p.m t) fun r μ' =>
      r = (X * Y) I J ∧ DSReadyP p t hmL aX aY b0 X Y μ' ∧
        SameOutside μ μ' (aWD p b0) (3 * p.L + p.m)

/-- preCore(L, m, t, N, D, aX, aY, b0) builds the data structure for the matrices at aX and aY in
the block from b0 on. Nothing is assumed about the content of the block. -/
def PreCoreSpecP (lim : Limits) (P : Program) (c : ℕ) : Prop :=
  ∀ (p : Sec2.Par) (t : ℕ) (hmL : p.m ≤ p.L) (aX aY b0 : ℕ) (X : Matrix (Fin p.N) (Fin (ThreeSumApsp.D p.m)) ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin p.N) ℤ) (U : ℤ) (μ : ℕ → ℤ),
    t ≤ p.m → Lim30 lim p t b0 U → (∀ i j, |X i j| ≤ U) → (∀ i j, |Y i j| ≤ U) →
    MatAt μ aX X → MatAt μ aY Y → aX + p.N * p.D ≤ b0 → aY + p.D * p.N ≤ b0 →
    ∀ d, d + (p.L + 6) ≤ lim.depth →
    Meets lim P Proc.preCore d [p.L, p.m, t, p.N, p.D, aX, aY, b0] μ
      (tPreCoreP c p t) fun _ μ' =>
      DSReadyP p t hmL aX aY b0 X Y μ' ∧ SameOutside μ μ' b0 (Light.Sec4.top p t b0 - b0)

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Sec4.ChoosingParameters
/-!
# 4.4 Choosing the parameters: the entropy function, the exponents `γ` and `q`, equation (11)

Section 4.4 chooses the parameters `L` and `t` of Theorem 30. This file has the facts that its
proofs use about `D = 4^m` and about the notions with which Corollary 31 is stated ("Other choices
of the parameters"). The notation of the paper: `entropy` is `H`, `rhoC c` is `ρ_c = 9/(c-1)`,
`gammaOf c θ` is `γ = θ ln(1/ρ_c)/ln 4`, `qOf θ` is `q = (H(θ) + θ ln 9)/ln 4`, `lnΛ c γ` is the
denominator `ln Λ` of (11), `Rc c γ` is `R_c(γ) = ln 4/ln Λ`, and `D m` is `D = 4^m`.

* The entropy function `H` is Mathlib's `Real.binEntropy` (`entropy_eq_binEntropy`); continuity,
  concavity and the derivative come from there.
* Powers and logarithms of `D = 4^m` (`cast_D_pos`, `one_le_cast_D`, `log_cast_D`,
  `cast_le_log_cast_D`, `cast_D_rpow`).
* The bound `binom(n, k) ≥ e^{n H(k/n)}/(n+1)` (`exp_entropy_div_le_choose`): it is the standard
  bound `binom(n, k) ≥ 1/(n+1) · n^n/(k^k (n-k)^{n-k})`, as `e^{n H(k/n)} = n^n/(k^k (n-k)^{n-k})`.
* One section for each of `ρ_c`, `γ`, `q`, and `ln Λ` with `R_c(γ)`: signs, monotonicity and
  continuity, and the three facts that tie them to the costs: `ρ_c^{θm} = D^{-γ}`
  (`rhoC_rpow_eq_D_rpow`), `e^{m(H(θ) + θ ln 9)} = D^q` (`exp_entropy_eq_D_rpow_qOf`), and
  `Λ < 4^{1/ε}` if and only if `ε < R_c(γ)` (`eq_11_iff`).
-/

public section

open Finset

namespace ThreeSumApsp

/-! ### The entropy function: the link with Mathlib's `Real.binEntropy` -/

/-! ### Powers and logarithms of `D = 4^m` -/

/-- `ln 4 > 0`. -/
theorem log_four_pos : 0 < Real.log 4 := Real.log_pos (by norm_num)

/-- `D = 4^m` as a real number. -/
theorem cast_D_eq (m : ℕ) : (ThreeSumApsp.D m : ℝ) = (4 : ℝ) ^ m := by
  simp [ThreeSumApsp.D]

/-- `D > 0`. -/
theorem cast_D_pos (m : ℕ) : (0 : ℝ) < (ThreeSumApsp.D m : ℝ) := by
  rw [cast_D_eq]
  positivity

/-- `D ≥ 1`. -/
theorem one_le_cast_D (m : ℕ) : (1 : ℝ) ≤ (ThreeSumApsp.D m : ℝ) := by
  rw [cast_D_eq]
  exact one_le_pow₀ (by norm_num)

/-- `ln D = m ln 4`. -/
theorem log_cast_D (m : ℕ) : Real.log (ThreeSumApsp.D m : ℝ) = (m : ℝ) * Real.log 4 := by
  rw [cast_D_eq, Real.log_pow]

/-- `m ≤ ln D`. -/
theorem cast_le_log_cast_D (m : ℕ) : (m : ℝ) ≤ Real.log (ThreeSumApsp.D m : ℝ) := by
  rw [log_cast_D]
  exact le_mul_of_one_le_right (Nat.cast_nonneg m) Real.one_le_log_four

/-- `D^x = e^{x m ln 4}`. -/
theorem cast_D_rpow (m : ℕ) (x : ℝ) :
    (ThreeSumApsp.D m : ℝ) ^ x = Real.exp (x * ((m : ℝ) * Real.log 4)) := by
  rw [Real.rpow_def_of_pos (cast_D_pos m), log_cast_D]
  ring_nf

/-! ### The lower bound on binomial coefficients by the entropy function -/

/-! ### The decay rate `ρ_c` -/

/-- `ρ_c = 9/(c-1) < 1` for `c > 10`. -/
theorem rhoC_lt_one {c : ℝ} (hc : 10 < c) : rhoC c < 1 :=
  (div_lt_one (by linarith)).2 (by linarith)

/-- `1/ρ_c = (c - 1)/9`. -/
theorem one_div_rhoC (c : ℝ) : 1 / rhoC c = (c - 1) / 9 := by
  rw [rhoC, one_div_div]

/-- `ln(1/ρ_c) = ln((c - 1)/9)` is positive for `c > 10`. -/
theorem log_inv_rhoC_pos {c : ℝ} (hc : 10 < c) : 0 < Real.log (1 / rhoC c) :=
  Real.log_pos (by rw [one_div_rhoC]; linarith)

/-! ### The exponent `γ` -/

/-- Corollary 31: "γ := θ ln(1/ρ_c)/ln 4 > 0". -/
theorem corollary_31_gamma_pos (c θ : ℝ) (hc : 10 < c) (hθ0 : 0 < θ) : 0 < gammaOf c θ :=
  div_pos (mul_pos hθ0 (log_inv_rhoC_pos hc)) log_four_pos

/-! ### The exponent `q` -/

/-- `e^{m(H(θ) + θ ln 9)} = D^q` with `D = 4^m` and `q = (H(θ) + θ ln 9)/ln 4`. -/
theorem exp_entropy_eq_D_rpow_qOf (θ : ℝ) (m : ℕ) :
    Real.exp ((m : ℝ) * (entropy θ + θ * Real.log 9)) = (ThreeSumApsp.D m : ℝ) ^ qOf θ := by
  rw [cast_D_rpow]
  congr 1
  have hfour := log_four_pos.ne'
  unfold qOf
  field_simp

/-! ### The denominator `ln Λ` of equation (11), and `R_c(γ)` -/

end ThreeSumApsp
end
end

section
-- Source module: ThreeSumApsp.Sec4.ParameterSteps
/-!
# Section 4.4: the steps of the proof of Corollary 26 that hold for all parameters

The proof of Corollary 31 begins: "We repeat the proof of Corollary 26 with L := ⌈cm⌉ and t :=
⌈θm⌉". This file has the parts of the proof of Corollary 26 that mention neither `L = 21m` nor
`t = ⌈m/9⌉`, under the headings of that proof. Both corollaries use them.

* Setting up. The inner dimension is padded to `D = 4^m` with `m = ⌈log_4 D⌉`
  (`Corollary26.setting_up`). This changes no entry of the product (`Corollary26.padding`), and it
  changes the bounds by a constant factor (`padded_le`). The switching order `t = ⌈θm⌉` is
  `switchOf θ m`, and `t ≤ m` (`switchOf_le`).
* Queries. `∑_{d ≤ t} α_d ≤ (9x)^t (1 + 1/x)^m` for every `x ≥ 1/9` (`sum_alpha_le`), and for
  `x = (1-θ)/θ` the right-hand side at `t = θm` is `D^q` (`rpow_mul_pow_eq_D_rpow_qOf`).
* Encodings. Inequality (10), whose left-hand side is `lhs10 L m γ`, bounds the last term of (8) and
  gives the hypothesis `N ≥ √K N₀` of Theorem 30 (`Equation10.last_term`, `Equation10.tile_fits`).
* Conclusion. The expression (8) from a bound on its first term and (10)
  (`dominated_cost8_of_eq_10`).

Bounds "up to a constant" are written with `Dominated`, in the one parameter `m` or in the record
`Sizes` of the given inner dimension, of `N` and of `m`.
-/

@[expose] public section

open Finset

namespace ThreeSumApsp

/-! ### Setting up -/

/-- Proof of Corollary 26, "Setting up": for "m := ⌈log_4 D⌉" and `D ≥ 2` the padded inner dimension
satisfies `D ≤ 4^m` and "4^m < 4D", and "the original D was larger than 4^{m-1}". -/
theorem Corollary26.setting_up (D₀ m : ℕ) (hD : 2 ≤ D₀) (hm : m = ⌈Real.logb 4 (D₀ : ℝ)⌉₊) :
    D₀ ≤ ThreeSumApsp.D m ∧ ThreeSumApsp.D m < 4 * D₀ ∧ 4 ^ (m - 1) < D₀ := by
  obtain rfl : m = Nat.clog 4 D₀ := by
    rw [hm, ← Real.natCeil_logb_natCast]
    norm_num
  have hpos : 1 ≤ Nat.clog 4 D₀ := Nat.clog_pos (by norm_num) hD
  have hlt : 4 ^ (Nat.clog 4 D₀ - 1) < D₀ := Nat.pow_pred_clog_lt_self (by norm_num) (by omega)
  have hsucc : 4 ^ Nat.clog 4 D₀ = 4 * 4 ^ (Nat.clog 4 D₀ - 1) := by
    rw [← pow_succ', Nat.sub_add_cancel hpos]
  refine ⟨Nat.le_pow_clog (by norm_num) _, ?_, hlt⟩
  unfold ThreeSumApsp.D
  omega

/-- The padded inner dimension is at least the given one. -/
theorem Sizes.SetUp.le {p : ThreeSumApsp.Sizes} (h : p.SetUp) : p.D₀ ≤ ThreeSumApsp.D p.m :=
  (Corollary26.setting_up p.D₀ p.m h.two_le h.m_eq).1

/-- "4^m < 4D". -/
theorem Sizes.SetUp.lt {p : ThreeSumApsp.Sizes} (h : p.SetUp) : ThreeSumApsp.D p.m < 4 * p.D₀ :=
  (Corollary26.setting_up p.D₀ p.m h.two_le h.m_eq).2.1

/-- "the original D was larger than 4^{m-1}". -/
theorem Sizes.SetUp.gt {p : ThreeSumApsp.Sizes} (h : p.SetUp) : 4 ^ (p.m - 1) < p.D₀ :=
  (Corollary26.setting_up p.D₀ p.m h.two_le h.m_eq).2.2

/-- `m ≥ 1`, because `2 ≤ D₀ ≤ 4^m`. -/
theorem Sizes.SetUp.m_pos {p : ThreeSumApsp.Sizes} (h : p.SetUp) : 1 ≤ p.m :=
  Nat.pos_of_ne_zero fun h0 => by
    have hle := h.le
    have htwo := h.two_le
    rw [h0, ThreeSumApsp.D, pow_zero] at hle
    omega

/-- Proof of Corollary 26, "Setting up": padding "changes D by a factor less than 4, which only
affects the constants". A bound `D^a log^k D` in the padded `D = 4^m` is, up to a constant, the same
bound in the original inner dimension `D₀`. (Corollary 26 has `k = 0`; the bounds of Corollary 31
have logarithms.) -/
theorem padded_le (a : ℝ) (k : ℕ) :
    Dominated Sizes.SetUp (fun p => (ThreeSumApsp.D p.m : ℝ) ^ a * Real.log (ThreeSumApsp.D p.m) ^ k)
      fun p => (p.D₀ : ℝ) ^ a * Real.log p.D₀ ^ k := by
  refine .of_le_const_mul (C := 4 ^ |a| * 3 ^ k) (by positivity) fun ⟨D₀, _, m⟩ h => ?_
  have htwo : (2 : ℝ) ≤ D₀ := by exact_mod_cast h.two_le
  have hD₀ : (0 : ℝ) < D₀ := by linarith
  have hle' : 1 ≤ (ThreeSumApsp.D m : ℝ) / D₀ := (one_le_div hD₀).2 (by exact_mod_cast h.le)
  have hle4 : (ThreeSumApsp.D m : ℝ) ≤ 4 * D₀ := by exact_mod_cast h.lt.le
  have hpow : (ThreeSumApsp.D m : ℝ) ^ a ≤ 4 ^ |a| * (D₀ : ℝ) ^ a :=
    calc (ThreeSumApsp.D m : ℝ) ^ a = ((ThreeSumApsp.D m : ℝ) / D₀) ^ a * (D₀ : ℝ) ^ a := by
          rw [← Real.mul_rpow (by positivity) hD₀.le, div_mul_cancel₀ _ hD₀.ne']
      _ ≤ 4 ^ |a| * (D₀ : ℝ) ^ a :=
          mul_le_mul_of_nonneg_right
            ((Real.rpow_le_rpow_of_exponent_le hle' (le_abs_self a)).trans
              (Real.rpow_le_rpow (by positivity) ((div_le_iff₀ hD₀).2 hle4) (abs_nonneg a)))
            (by positivity)
  have hlog : Real.log (ThreeSumApsp.D m) ≤ 3 * Real.log D₀ :=
    calc Real.log (ThreeSumApsp.D m) ≤ Real.log (4 * D₀) := Real.log_le_log (cast_D_pos m) hle4
      _ = 2 * Real.log 2 + Real.log D₀ := by
          rw [Real.log_mul (by norm_num) hD₀.ne', Real.log_four]
      _ ≤ 3 * Real.log D₀ := by linarith [Real.log_le_log two_pos htwo]
  calc (ThreeSumApsp.D m : ℝ) ^ a * Real.log (ThreeSumApsp.D m) ^ k
      ≤ 4 ^ |a| * (D₀ : ℝ) ^ a * (3 * Real.log D₀) ^ k :=
        mul_le_mul hpow (pow_le_pow_left₀ (by positivity) hlog k) (by positivity) (by positivity)
    _ = 4 ^ |a| * 3 ^ k * ((D₀ : ℝ) ^ a * Real.log D₀ ^ k) := by
        rw [mul_pow]
        ring

/-- The hypothesis `t ≤ m` of Theorem 30 holds for `t = ⌈θm⌉` with `θ ≤ 1`. -/
theorem switchOf_le {θ : ℝ} (hθ : θ ≤ 1) (m : ℕ) : switchOf θ m ≤ m :=
  Nat.ceil_le.2 (mul_le_of_le_one_left (Nat.cast_nonneg m) hθ)

/-! ### Queries -/

/-- Proof of Corollary 26, "Queries": "A query reads ∑_{d≤t} α_d numbers, and we bound this sum as
in the proof of Lemma 11. Since 9^d = 72^d 8^{-d} ≤ 72^t 8^{-d} for d ≤ t", the sum `∑_{d ≤ t} α_d`
is at most `72^t ∑_{d=0}^{m} binom(m, d) 8^{-d} = 72^t (9/8)^m`. Here any `x > 0` with `9x ≥ 1`
stands for 8, because the proof of Corollary 31 makes "the same calculation with x := (1-θ)/θ in
place of 8". -/
theorem sum_alpha_le_b299791 {x : ℝ} (hx : 0 < x) (h9x : 1 ≤ 9 * x) {m t : ℕ} (htm : t ≤ m) :
    ∑ d ∈ range (t + 1), (alpha m d : ℝ) ≤ (9 * x) ^ t * (1 + 1 / x) ^ m := by
  -- `9^d = (9x)^d x^{-d} ≤ (9x)^t x^{-d}` for `d ≤ t`
  have hterm : ∀ d ∈ range (t + 1),
      (alpha m d : ℝ) ≤ (9 * x) ^ t * ((1 / x) ^ d * 1 ^ (m - d) * (m.choose d : ℝ)) := by
    intro d hd
    have hpow : (9 * x) ^ d ≤ (9 * x) ^ t :=
      pow_le_pow_right₀ h9x (by rw [mem_range] at hd; omega)
    calc (alpha m d : ℝ) = (9 * x) ^ d * ((1 / x) ^ d * 1 ^ (m - d) * (m.choose d : ℝ)) := by
          rw [alpha, one_pow, mul_one, ← mul_assoc, ← mul_pow,
            show 9 * x * (1 / x) = 9 by field_simp]
          push_cast
          ring
      _ ≤ (9 * x) ^ t * ((1 / x) ^ d * 1 ^ (m - d) * (m.choose d : ℝ)) := by gcongr
  calc ∑ d ∈ range (t + 1), (alpha m d : ℝ)
      ≤ ∑ d ∈ range (t + 1), (9 * x) ^ t * ((1 / x) ^ d * 1 ^ (m - d) * (m.choose d : ℝ)) :=
        sum_le_sum hterm
    _ ≤ ∑ d ∈ range (m + 1), (9 * x) ^ t * ((1 / x) ^ d * 1 ^ (m - d) * (m.choose d : ℝ)) :=
        sum_le_sum_of_subset_of_nonneg (range_subset_range.2 (by omega))
          fun _ _ _ => by positivity
    _ = (9 * x) ^ t * (1 + 1 / x) ^ m := by rw [← mul_sum, ← add_pow, add_comm]

/-- `(9x)^{θm} (1+1/x)^m = D^q` for "x := (1-θ)/θ". -/
theorem rpow_mul_pow_eq_D_rpow_qOf {θ : ℝ} (hθ0 : 0 < θ) (hθ1 : θ < 1) (m : ℕ) :
    (9 * ((1 - θ) / θ)) ^ (θ * (m : ℝ)) * (1 + 1 / ((1 - θ) / θ)) ^ m = (ThreeSumApsp.D m : ℝ) ^ qOf θ := by
  have h1θ : 0 < 1 - θ := sub_pos.2 hθ1
  have hinv : 1 + 1 / ((1 - θ) / θ) = (1 - θ)⁻¹ := by
    field_simp
    ring
  rw [← exp_entropy_eq_D_rpow_qOf, hinv, Real.rpow_def_of_pos (by positivity),
    ← Real.rpow_natCast, Real.rpow_def_of_pos (by positivity), ← Real.exp_add, Real.log_inv,
    Real.log_mul (by norm_num) (by positivity), Real.log_div h1θ.ne' hθ0.ne', entropy]
  congr 1
  ring

/-! ### Encodings -/

/-! ### Conclusion -/

end ThreeSumApsp
end
end

section
-- Source module: ThreeSumApsp.Sec4.Corollary31
/-!
# Corollary 31: the costs of Theorem 30 as powers of `D`

"We repeat the proof of Corollary 26 with L := ⌈cm⌉ and t := ⌈θm⌉". Here `L` is `levelsOf c m` and
`t` is `switchOf θ m`, for constants `c > 10` and `0 < θ < 0.9`; the inner dimension satisfies
`2 ≤ D ≤ N^ε` with `ε < R_c(γ)`, and `m = ⌈log_4 D⌉`. The main result is `Corollary31.costs`: for
all `m` from some `m₀` on, the preprocessing cost (8) of Theorem 30 (`cost8`, the sum of a first
term for the boxes and a last term for the encodings) is `O(N² log² D / D^γ)`, and its query cost
`L ∑_{d ≤ t} α_d` (`costQuery`) is `O(D^q log D)`. What the paper says of the finitely many smaller
`m` is `Corollary31.small_m`; nothing else rests on that lemma.

The steps are those of the proof of Corollary 26. What does not depend on `c` and `θ` is proved
before, for both corollaries; the rest is what the paper says of each step.

* Setting up. As for Corollary 26, the inner dimension is padded to `D = 4^m`
  (`Corollary26.setting_up`, `Corollary26.padding`); this changes the bounds by a constant factor
  (`padded_le`). The hypothesis reads `N ≥ 4^{(m-1)/ε}` (`Corollary31.hypothesis`).
* Boxes. `ρ < ρ_c` and `ρ^t ≤ D^{-γ}` (`Corollary31.rho_lt`, `Corollary31.rho_pow_le`), which bounds
  the first term of (8) (`Corollary31.first_term`).
* Queries. The calculation of Corollary 26, `∑_{d ≤ t} α_d ≤ (9x)^t (1 + 1/x)^m` (`sum_alpha_le`),
  with `x = (1-θ)/θ` in place of 8 gives `O(D^q)` (`Corollary31.sum_alpha`) and so
  `Corollary31.query_cost`.
* Encodings. As for Corollary 26, inequality (10) (`eq_10`) bounds the last term of (8)
  (`Equation10.last_term`) and gives the hypothesis `N ≥ √K N₀` of Theorem 30
  (`Equation10.tile_fits`). Its left-hand side is `O(√m · Λ^m)` (`Corollary31.lhs10_le`), and the
  base `Λ` is smaller than the base `4^{1/ε}` (`eq_11_iff`). This `Λ` is the general form of the `Λ`
  of Corollary 26 (`exp_lnΛ`, `baseLambda_eq_exp_lnΛ`).

The last sentence of the corollary, on `R_c(0)`, is `sec4_Rc_strictAntiOn` and
`sec4_Rc_zero_tendsto`. Bounds "up to a constant" are written with `Dominated`, in the one parameter
`m` or in the record `Sizes` of the given `D`, of `N` and of `m`.
-/

@[expose] public section

open Finset

namespace ThreeSumApsp

/-! ### Setting up -/

/-- With "L := ⌈cm⌉" (proof of Corollary 31) the hypothesis `L ≥ 10m` of Theorem 30 holds. -/
theorem ten_mul_le_levelsOf {c : ℝ} (hc : 10 < c) (m : ℕ) : 10 * m ≤ levelsOf c m := by
  have h : ((10 * m : ℕ) : ℝ) ≤ c * m := by
    push_cast
    exact mul_le_mul_of_nonneg_right hc.le (Nat.cast_nonneg m)
  exact_mod_cast h.trans (Nat.le_ceil _)

/-- `L = ⌈cm⌉ ≤ (c + 1) m` for `m ≥ 1`. -/
lemma levelsOf_le {c : ℝ} (hc : 10 < c) {m : ℕ} (hm : 1 ≤ m) :
    (levelsOf c m : ℝ) ≤ (c + 1) * m := by
  have hm' : (1 : ℝ) ≤ m := by exact_mod_cast hm
  have hL : (levelsOf c m : ℝ) < c * m + 1 := Nat.ceil_lt_add_one (by positivity)
  linarith

/-- `m = O(log D)` for `D = 4^m`. -/
lemma dominated_cast_log_D :
    Dominated (fun m : ℕ => 1 ≤ m) (fun m => (m : ℝ)) fun m => Real.log (ThreeSumApsp.D m) :=
  .of_le fun m _ => cast_le_log_cast_D m

/-- `L = ⌈cm⌉ = O(log D)` for `D = 4^m`. -/
lemma dominated_levelsOf_log_D {c : ℝ} (hc : 10 < c) :
    Dominated (fun m : ℕ => 1 ≤ m) (fun m => (levelsOf c m : ℝ)) fun m => Real.log (ThreeSumApsp.D m) :=
  (Dominated.of_le_const_mul (C := c + 1) (by linarith) fun _ hm => levelsOf_le hc hm).trans
    dominated_cast_log_D

/-- Proof of Corollary 31: "After the padding, the assumption D ≤ N^ε now reads N ≥ 4^{(m-1)/ε}",
because `4^{m-1} < D ≤ N^ε` for the original `D`. -/
theorem Corollary31.hypothesis {m D₀ N : ℕ} {ε : ℝ} (hm : 1 ≤ m) (hε : 0 < ε)
    (hgt : 4 ^ (m - 1) < D₀) (hthin : (D₀ : ℝ) ≤ (N : ℝ) ^ ε) :
    (4 : ℝ) ^ (((m : ℝ) - 1) / ε) ≤ (N : ℝ) := by
  have hle : (4 : ℝ) ^ ((m : ℝ) - 1) ≤ (N : ℝ) ^ ε := by
    rw [← Nat.cast_one, ← Nat.cast_sub hm, Real.rpow_natCast]
    exact le_trans (by exact_mod_cast hgt.le) hthin
  calc (4 : ℝ) ^ (((m : ℝ) - 1) / ε) = ((4 : ℝ) ^ ((m : ℝ) - 1)) ^ (1 / ε) := by
        rw [← Real.rpow_mul (by norm_num), mul_one_div]
    _ ≤ ((N : ℝ) ^ ε) ^ (1 / ε) := Real.rpow_le_rpow (by positivity) hle (by positivity)
    _ = (N : ℝ) := by
        rw [← Real.rpow_mul N.cast_nonneg, mul_one_div_cancel hε.ne', Real.rpow_one]

/-! ### Boxes -/

/-- Proof of Corollary 31: "ρ < ρ_c", where `ρ = 9m/(L-m+1)` and `L = ⌈cm⌉ ≥ cm`. -/
theorem Corollary31.rho_lt {c : ℝ} (hc : 10 < c) (m : ℕ) : ThreeSumApsp.rho (levelsOf c m) m < rhoC c := by
  have hL : c * m ≤ levelsOf c m := Nat.le_ceil _
  have hm : (m : ℝ) ≤ c * m := le_mul_of_one_le_left (Nat.cast_nonneg m) (by linarith)
  unfold ThreeSumApsp.rho rhoC
  -- after clearing denominators: `9m (c - 1) < 9 (L - m + 1)`, from `cm ≤ L`
  rw [div_lt_div_iff₀ (by linarith) (by linarith)]
  linarith

/-! ### Queries -/

/-- Proof of Corollary 31: "the same calculation with x := (1-θ)/θ in place of 8, where 9x > 1
because θ < 0.9, gives ∑_{d≤t} α_d ≤ (9x)^t (1+1/x)^m = O(D^q)". The constant is `9x`, because
`t = ⌈θm⌉ < θm + 1`. -/
theorem Corollary31.sum_alpha {θ : ℝ} (hθ0 : 0 < θ) (hθ1 : θ < 0.9) :
    Dominated (fun m : ℕ => 1 ≤ m) (fun m => ∑ d ∈ range (switchOf θ m + 1), (alpha m d : ℝ))
      fun m => (ThreeSumApsp.D m : ℝ) ^ qOf θ := by
  have hx : 0 < (1 - θ) / θ := div_pos (by linarith) hθ0
  have h9x : 1 < 9 * ((1 - θ) / θ) := by
    rw [← mul_div_assoc, lt_div_iff₀ hθ0]
    linarith
  refine .of_le_const_mul (C := 9 * ((1 - θ) / θ)) (by positivity) fun m _ => ?_
  have ht : ((switchOf θ m : ℕ) : ℝ) ≤ θ * m + 1 := (Nat.ceil_lt_add_one (by positivity)).le
  calc ∑ d ∈ range (switchOf θ m + 1), (alpha m d : ℝ)
      ≤ (9 * ((1 - θ) / θ)) ^ switchOf θ m * (1 + 1 / ((1 - θ) / θ)) ^ m :=
        sum_alpha_le_b299791 hx h9x.le (switchOf_le (by linarith) m)
    _ ≤ (9 * ((1 - θ) / θ)) ^ (θ * (m : ℝ) + 1) * (1 + 1 / ((1 - θ) / θ)) ^ m := by
        gcongr
        rw [← Real.rpow_natCast]
        exact Real.rpow_le_rpow_of_exponent_le h9x.le ht
    _ = 9 * ((1 - θ) / θ) * (ThreeSumApsp.D m : ℝ) ^ qOf θ := by
        rw [← rpow_mul_pow_eq_D_rpow_qOf hθ0 (by linarith) m, Real.rpow_add (by positivity),
          Real.rpow_one]
        ring

/-- The query cost `L ∑_{d ≤ t} α_d` of Theorem 30 is `O(L D^q) = O(D^q log D)`, for `D = 4^m`. (In
the proof of Corollary 26: "Hence a query takes O(L D^q) = O(m D^q) time".) -/
theorem Corollary31.query_cost {c θ : ℝ} (hc : 10 < c) (hθ0 : 0 < θ) (hθ1 : θ < 0.9) :
    Dominated (fun m : ℕ => 1 ≤ m) (fun m => costQuery (levelsOf c m) m (switchOf θ m))
      fun m => (ThreeSumApsp.D m : ℝ) ^ qOf θ * Real.log (ThreeSumApsp.D m) :=
  ((dominated_levelsOf_log_D hc).mul (Corollary31.sum_alpha hθ0 hθ1) (fun _ _ => by positivity)
    fun _ _ => by positivity).congr (fun _ _ => rfl) fun m _ => by ring

/-! ### Encodings

As in the proof of Corollary 26, `m`-th powers are compared by comparing their bases. The paper's
sentences are one lemma each: the bound on the binomial coefficient (`Corollary31.le_K`, for which
`Corollary31.mul_entropy_monotoneOn` passes from `L` to `cm`), the bound `O(√m · Λ^m)` on the
left-hand side (`Corollary31.lhs10_le`), where `Λ` is `Real.exp (lnΛ c γ)` (`exp_lnΛ`), and
`Λ < 4^{1/ε}` (`eq_11_iff`); `eq_10` puts them together with `Corollary31.hypothesis`. Throughout,
`L = ⌈cm⌉`, and `D₀` is the inner dimension before padding. -/

/-! ### The corollary -/

end ThreeSumApsp
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.ChoosingParameters.Parameters
/-!
# Corollaries 26 and 31 in the light language: rational parameters in the program text

Proof of Corollary 26: "Setting up. Let m := ⌈log₄ D⌉, and pad the inner dimension to 4^m < 4D";
proof of Corollary 31: "with L := ⌈cm⌉ and t := ⌈θm⌉". A program is a finite text, so it can hold c
and θ only if they are rational: c = a/b and θ = p/q (`RatParams`). For real c and θ the statements
follow by approximation (`exists_ratParams`). Corollary 26 is the case c = 21, θ = 1/9, with the
threshold 60 (`ratParams26`).

This file has the parameters, what a structure at the free pointer fr consists of (`structEnd`,
`Ready31P`), what the routines ask of the limits (`Lim31`), the query with its specification and its
proof (`QuerySpec31P`, `query31_meets`), and the text and the specification of the preprocessing
(`pre31Body`, `PreSpec31P`).
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.Spec

/-! ## Parameters -/

variable {lim : Limits} {P : Program} {G : RatParams} {N D₀ : ℕ} {X : Matrix (Fin N) (Fin D₀) ℤ}
  {Y : Matrix (Fin D₀) (Fin N) ℤ} {aX aY fr : ℕ} {U : ℤ} {μ μ' : ℕ → ℤ}

/-- The data structure of Theorem 30 for the padded matrices has been built in the cells after fr.
-/
abbrev Built31P (G : RatParams) {N D₀ : ℕ} (X : Matrix (Fin N) (Fin D₀) ℤ)
    (Y : Matrix (Fin D₀) (Fin N) ℤ) (fr : ℕ) (μ : ℕ → ℤ) : Prop :=
  DSReadyP (parOf G N D₀) (switchOf31 G D₀) (logFour_le_levels G N D₀) (paddedXAt fr)
    (paddedYAt N D₀ fr) (blockAt N D₀ fr) (padInnerCols (ThreeSumApsp.D (logFour D₀)) X)
    (padInnerRows (ThreeSumApsp.D (logFour D₀)) Y) μ

/-- **The cells from fr on hold what a query needs.** -/
structure Ready31P (G : RatParams) {N D₀ : ℕ} (X : Matrix (Fin N) (Fin D₀) ℤ)
    (Y : Matrix (Fin D₀) (Fin N) ℤ) (aX aY fr : ℕ) (μ : ℕ → ℤ) : Prop where
  matX : MatAt μ aX X
  matY : MatAt μ aY Y
  small : logFour D₀ < G.m₀ → μ fr = 0
  large : G.m₀ ≤ logFour D₀ → μ fr = 1 ∧ μ (fr + 1) = blockAt N D₀ fr ∧ Built31P G X Y fr μ

/-! ## The query: the text is query31Body -/

/-- query31 returns the entry (XY)[I, J], keeps the structure ready, and changes only cells of the
structure behind its first three. -/
def QuerySpec31P (lim : Limits) (P : Program) (G : RatParams) : Prop :=
  ∀ (N D₀ aX aY fr : ℕ) (X : Matrix (Fin N) (Fin D₀) ℤ) (Y : Matrix (Fin D₀) (Fin N) ℤ) (U : ℤ)
      (μ : ℕ → ℤ) (I J : Fin N), Input31 lim G X Y aX aY fr U → Ready31P G X Y aX aY fr μ →
    ∀ d, d + 3 ≤ lim.depth → Meets lim P Proc.query31 d [(I : ℕ), (J : ℕ), N, D₀, aX, aY, fr] μ
        (tQuery31 G D₀) fun r μ' =>
      r = (X * Y) I J ∧ Ready31P G X Y aX aY fr μ' ∧
          SameOutside μ μ' (fr + 3) (structEnd G N D₀ fr - (fr + 3))

/-! ## The preprocessing: text and interface -/

namespace Pre31

end Pre31

/-- The time of the preprocessing; c is the constant of the shared stage. -/
def tPre31P (c : ℕ) (G : RatParams) (N D₀ : ℕ) : ℕ :=
  tLog4 (logFour D₀) + 30 +
    if logFour D₀ < G.m₀ then 0 else
      tCeilMul G.a (logFour D₀) + tCeilMul G.p (logFour D₀) + tPadX N (ThreeSumApsp.D (logFour D₀))
        + tCopy (D₀ * N) + tFill ((ThreeSumApsp.D (logFour D₀) - D₀) * N)
        + tPreCoreP c (parOf G N D₀) (switchOf31 G D₀) + 120

/-- pre31 prepares the cells from fr on, about which nothing is assumed, for queries to XY. -/
def PreSpec31P (lim : Limits) (P : Program) (c : ℕ) (G : RatParams) : Prop :=
  ∀ (N D₀ aX aY fr : ℕ) (X : Matrix (Fin N) (Fin D₀) ℤ) (Y : Matrix (Fin D₀) (Fin N) ℤ) (U : ℤ)
      (μ : ℕ → ℤ), Input31 lim G X Y aX aY fr U → MatAt μ aX X → MatAt μ aY Y →
    ∀ d, d + (G.L (logFour D₀) + 7) ≤ lim.depth → Meets lim P Proc.pre31 d [N, D₀, aX, aY, fr] μ
        (tPre31P c G N D₀) fun _ μ' =>
      Ready31P G X Y aX aY fr μ' ∧ SameOutside μ μ' fr (structEnd G N D₀ fr - fr)

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.Theorem30.WordSize
/-!
# Theorem 30 in the light language: limits that are polynomial in N

"Word size", proof of Theorem 30.  For every exponent c of the bound N^c on the entries there are
limits (`lim30`) that the two routines preCore and queryAt can live with (`lim30_ok`) and that are
at most 2^9 ((N + 1) (D + 1))^(20 + 2c) (`small_lim30`).  The exponents are generous: since (10^L)²
≤ N^5, every quantity that depends on L is at most B = (N + 1)^5 (`below`), the block has at most
222 B⁴ cells (`top_sub_le_pow`), and the word bound `word30` is a sum of seven terms, each at most a
small multiple of (N + 1)^(20 + 2c) (`word30_le`).  What the two routines ask of the limits is
`Lim30`, in their specifications.  The polynomial bound is for a block whose base address b0 is at
most 10 (N + 1)³; the end theorems, which put the block behind the input, show this of their b0
(`b0_le`, `wantedB0_le`, `blockAt_le`).
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.Spec Finset

/-- The word bound: the sum of everything that has to fit in a word. -/
def word30 (p : Sec2.Par) (t b0 c : ℕ) : ℕ :=
  Light.Sec4.top p t b0 + 100 + 10 ^ p.m * ((7 ^ p.L * p.N ^ c) * (7 ^ p.L * p.N ^ c))
    + 7 ^ (p.L + 1) * p.N ^ c + 10 ^ (p.L + 1) + p.N ^ c + p.L

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.ChoosingParameters.Limits
/-!
# Section 4.4: the limits of the programs with rational parameters

`lim31` gives the limits within which the programs run, in two regimes: for m < m₀ queries compute
inner products; for m ≥ m₀ the limits are those of Theorem 30 plus the numbers formed in the
computation of L and t. They are what the programs need (`lim31_ok`), they allow exactly the cells
of the structure (`lim31_space`), and they are polynomial in N (`small_lim31`): every summand is at
most a constant, which depends on the parameters only, times ((N + 1) (D₀ + 1))^(20 + 2c).
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.Spec ThreeSumApsp.WordRam

/-! ## The limits -/

/-- The numbers that the computation of L and t and the comparison with the threshold form. -/
def extraWord31 (G : RatParams) (D₀ : ℕ) : ℕ := G.a * logFour D₀ + G.b + G.p * logFour D₀ + G.q +
    G.m₀

/-- The word bound when m is below the threshold. -/
def wordSmall31 (G : RatParams) (N D₀ fr c : ℕ) : ℕ :=
  fr + 3 + 100 + 4 ^ logFour D₀ + extraWord31 G D₀ + D₀ * (N ^ c * N ^ c) + N ^ c

/-- The limits of the programs. -/
def lim31 (G : RatParams) (N D₀ fr c : ℕ) : Limits :=
  if logFour D₀ < G.m₀ then ⟨(wordSmall31 G N D₀ fr c : ℤ), fr + 3, G.L G.m₀ + 10⟩
  else ⟨((word30 (parOf G N D₀) (switchOf31 G D₀) (blockAt N D₀ fr) c + extraWord31 G D₀ : ℕ) : ℤ),
    G.blockEnd N D₀ fr, G.L (logFour D₀) + 10⟩

/-! ## The limits are polynomial in N

All limits are at most `slopeConst31 G` times the polynomial B = ((N + 1) (D₀ + 1))^(20 + 2c): the
summands that grow with N and D₀ are at most a numeral times B, and each summand that depends on the
parameters only occurs in `slopeConst31 G`. -/

section poly

variable {G : RatParams} {N D₀ fr c : ℕ}

end poly

/-! ## Cells and natural numbers -/

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.ChoosingParameters.Offline
/-!
# Section 4.4, the offline form: preprocess, then one query for each wanted position

Corollary 26: "Hence, for every set W of positions of an N × N matrix, the entries (XY)[I, J], (I,
J) ∈ W, can be computed deterministically in O(|W| D^{0.437} + N²/D^{0.063}) time". Its proof: "The
bound for a set W follows by asking |W| queries." Proof of Theorem 25: "ask one query for each
position of W"; for Corollary 32 the paper uses "the same argument", "With Corollary 31 in place of
Theorem 24".

offline32(N, D, w, U, x, y, wi, wj, out, fr) preprocesses the matrices at x and y (pre31, which uses
the cells from the free pointer fr on) and then asks one query for each of the w positions, whose
rows are at wi and whose columns are at wj; the answers go to out. The arguments are those of the
task `thinTask`. The text `offline32Body`, with the round `offlineAsk32` of its loop, serves all
rational parameters c and θ: they are in the body of the procedure pre31 (`pre31Body`).

A round adds one answer (`offlineAsk32_ends`), with the invariant `Answered` that the offline form
of Theorem 30 uses too, and the routine meets its specification `OfflineSpec32P` for all parameters G
(`offline32_meets`).
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.Spec

variable {lim : Limits} {P : Program} {G : RatParams} {N D₀ : ℕ} {X : Matrix (Fin N) (Fin D₀) ℤ}
  {Y : Matrix (Fin D₀) (Fin N) ℤ} {aX aY fr : ℕ} {U : ℤ} {μ μ' : ℕ → ℤ}

/-! ## The text -/

namespace Offline32

end Offline32

/-! ## What it does -/

/-- The time of offline32: the preprocessing, and for each wanted position a query, the reading of
the position and the storing of the answer. -/
def tOffline32P (c : ℕ) (G : RatParams) (N D₀ w : ℕ) : ℕ :=
  tPre31P c G N D₀ + w * (tQuery31 G D₀ + 30) + 20

/-- offline32 writes the wanted entries of XY to out. Below fr it changes only the output; nothing
is assumed about the cells from fr on. -/
def OfflineSpec32P (lim : Limits) (P : Program) (c : ℕ) (G : RatParams) : Prop :=
  ∀ (N D₀ aX aY aI aJ out fr : ℕ) (X : Matrix (Fin N) (Fin D₀) ℤ) (Y : Matrix (Fin D₀) (Fin N) ℤ)
    (U : ℤ) (u : ℤ) (μ : ℕ → ℤ) (WI WJ : List ℕ),
    Input31 lim G X Y aX aY fr U → OfflineInput32 X Y aX aY aI aJ out fr WI WJ μ →
    ∀ d, d + (G.L (logFour D₀) + 8) ≤ lim.depth → Meets lim P Proc.offline32 d
      [N, D₀, WI.length, u, aX, aY, aI, aJ, out, fr] μ (tOffline32P c G N D₀ WI.length)
      fun _ μ' => Seg μ' out ((WI.zip WJ).map fun q => entryN X Y q.1 q.2) ∧
        ∀ a < fr, (a < out ∨ out + WI.length ≤ a) → μ' a = μ a

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.ChoosingParameters.Preprocessing
/-!
# Corollaries 26 and 31 in the light language: the preprocessing, with rational parameters

Proof of Corollary 26: "Setting up. Let m := ⌈log₄ D⌉, and pad the inner dimension to 4^m < 4D with
zero columns of X and zero rows of Y." Proof of Corollary 31: "We repeat the proof of Corollary 26
with L := ⌈cm⌉ and t := ⌈θm⌉". Then the preprocessing of Theorem 30 on the padded matrices; "for
smaller m the corollary again holds trivially".

The routine meets its specification (`pre31_meets`). It computes m and branches. Below the threshold
it writes the flag 0 (`PreInput31.ready_small`). From the threshold on it computes L, t and the
addresses, pads X, and pads Y by a copy and a fill with zeros (`prePad31_ends`,
`matAt_padInnerRows`); then it calls the preprocessing of Theorem 30 and writes the base address and
the flag 1 (`preBuild31_ends`, `PreInput31.ready_large`).
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.Spec

variable {lim : Limits} {P : Program} {G : RatParams} {N D₀ : ℕ} {X : Matrix (Fin N) (Fin D₀) ℤ}
  {Y : Matrix (Fin D₀) (Fin N) ℤ} {aX aY fr : ℕ} {U : ℤ} {μ μ' : ℕ → ℤ}

/-! ## What the routine is given, and what it leaves -/

/-- The routines that the preprocessing calls meet their specifications. -/
structure PreCalls31P (lim : Limits) (P : Program) (c : ℕ) (G : RatParams) : Prop where
  log4 : Log4Spec lim P
  levels : CeilMulSpec lim P Proc.levels31 G.a G.b
  switch : CeilMulSpec lim P Proc.switch31 G.p G.q
  padX : PadXSpec lim P
  copy : CopySpec lim P
  fill : FillSpec lim P
  preCore : PreCoreSpecP lim P c

/-! ## From the threshold on -/

/-! ## The routine -/

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Lang.Lib.CountSort
/-!
# Counting sort of a list of items by a key that is read through the item

countSort(n, src, dst, key, stride, off, nb, cnt) sorts the n items at src stably by their keys and
writes them to dst.  The key of the item x is the number in the cell key + x · stride + off, a
natural number below nb.  cnt is a scratch area of nb + 1 cells, of arbitrary content on entry; on
exit cnt[t] is the number of items with a key below t, which is where bucket t starts.

The routine has five phases, each of them one loop:
* `csZero` sets the cells of cnt to 0 (`csZero_ends`);
* `csCount` counts the items with key t in cnt[t + 1] (`csCount_ends`);
* `csPrefix` adds up, so that cnt[t] is the start of bucket t (`csPrefix_ends`);
* `csPlace` writes each item to the next free place of its bucket, which it reads from cnt and moves
  on by one (`csPlace_ends`);
* `csShift` moves the contents of cnt up by one cell, so that cnt[t] is again the start of bucket t
  (`csShift_ends`).

Each phase has an invariant, a structure about the memory, with two lemmas: `start` says that it
holds before the first round, and `succ` says what a round does to the memory.  Three numbers
describe the memory, where kf j is the key of item number j: `cntLt kf t n` is the number of items
with a key below t, the start of bucket t; `cntEq kf t i` is the number of items before item number
i with key t; and `sortPos kf n j`, their sum at t = kf j and i = j, is the place of item number j.
What is assumed is `CountSort.Pre`, what is achieved is `CountSort.Post`, and the result is
`countSort_meets`.
-/

@[expose] public section

open ThreeSumApsp

namespace Light

variable {lim : Limits} {P : Program} {d : ℕ}

/-! ## The program -/

namespace CountSort

end CountSort

@[simp] theorem csKey_cost : csKey.cost = 11 := rfl

/-! ## What is assumed, and the state between the phases -/

namespace CountSort

end CountSort

open CountSort

variable {μ μ' : ℕ → ℤ} {A : Args} {it kf : ℕ → ℕ} {σ : State}

/-! ## The first phase: the counters are set to 0 -/

/-! ## The second phase: the sizes of the buckets -/

/-! ## The third phase: the starts of the buckets -/

/-! ## The fourth phase: the items go to their places -/

/-! ## The fifth phase: the contents of cnt move up by one cell -/

/-! ## The routine -/

end Light
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.Theorem30.Preprocessing
/-!
# The preprocessing of Theorem 30, at a given place of the memory

Proof of Theorem 30, "Preprocessing". The routine first runs the stage that it shares with
Theorem 5: the list of the subsets, the tables for the rows and columns, and the encodings of all
row bands and column bands, at the leaves with at most m nines (the pruned shared stage). It notes t
in the directory, computes the addresses of the areas of Section 4 (preCoreAddr), starts the trie
area with the array [0], and computes the values of all the boxes of all tiles, by Lemma 29
(preCoreFinish). The routine allTiles is handed the encodings as the memory holds them; at the
leaves that it reads they are the true ones (`allTries_mem_eq`). There is one lemma for each of the
two named parts; preCore_spec puts them behind the shared stage.
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.Spec

/-! ## The program -/

namespace PreCore

end PreCore

open PreCore in
/-- preCore(L, m, t, N, D, aX, aY, b0) builds the data structure for the matrices at aX and aY in
the block from b0 on. -/
def preCoreBodyP : Stmt :=
  .call Sec2.pSharedP [Light.v Light.Sec4.PreCore.Levels, Light.v Light.Sec4.PreCore.Inner, Light.v Light.Sec4.PreCore.Size, Light.v Light.Sec4.PreCore.Width, Light.v Light.Sec4.PreCore.MatX, Light.v Light.Sec4.PreCore.MatY, Light.v Light.Sec4.PreCore.Base] Light.Sec4.PreCore.Void ;;
  .store (Light.v Light.Sec4.PreCore.Base +' Light.k Dir.switch) (Light.v Light.Sec4.PreCore.Switch) ;;
  preCoreAddr ;;
  preCoreFinish

variable {lim : Limits} {P : Program} {d : ℕ} {p : Sec2.Par} {t : ℕ} {hmL : p.m ≤ p.L}
  {aX aY b0 : ℕ} {X : Matrix (Fin p.N) (Fin (D p.m)) ℤ} {Y : Matrix (Fin (D p.m)) (Fin p.N) ℤ}
  {U : ℤ} {μ : ℕ → ℤ}

/-! ## The addresses -/

/-! ## The tries -/

/-- The arguments of the call of allTiles, and the data behind them: the encodings are what the
memory μ holds. -/
def allArgsP (p : Sec2.Par) (t b0 : ℕ) (μ : ℕ → ℤ) : AllTilesArgs where
  L := p.L
  m := p.m
  t := t
  nB := p.nB
  encA := memRow p b0 μ
  encB := memCol p b0 μ
  aENCA := p.aENCA b0
  aENCB := p.aENCB b0
  aR := aROOTS p b0
  tr := aTR p b0
  cap := trieCap p t
  fp := aFP p b0
  cur := aCUR p b0
  box := aBOX p b0

/-! ## The routine -/

end Light.Sec4
end
end