Expectation and tail bounds for matrix sampling

import RNLAPhaseB.B3Base
import RNLAPhaseB.Complexity

open MeasureTheory
open scoped Matrix
open scoped Matrix.Norms.L2Operator
open scoped BigOperators
open scoped ENNReal

namespace RNLAPhaseB.MtxSampling

/-- Measurable-space structure on finite matrices over a measurable scalar type,
inherited from their representation as finitely indexed functions. -/
noncomputable instance instMeasurableSpaceMatrixFinScalar {α : Type*} [MeasurableSpace α]
    (p q : ℕ) : MeasurableSpace (Matrix (Fin p) (Fin q) α) := by
  unfold Matrix
  infer_instance

/-- Expectation `𝔼 M` of a matrix-valued random variable, taken entrywise as a
Bochner integral against the sample measure. -/
noncomputable def matrixExpectation {Ω 𝕜 : Type*} [MeasurableSpace Ω] [RCLike 𝕜]
    {p q : ℕ} (P : Measure Ω) (M : Ω → Matrix (Fin p) (Fin q) 𝕜) :
    Matrix (Fin p) (Fin q) 𝕜 :=
  Matrix.of fun i j => ∫ ω, M ω i j ∂P

/-- The per-sample second moment
`v(X) := max { ‖𝔼[X X^*]‖, ‖𝔼[X^* X]‖ }` (spectral norms). -/
noncomputable def perSampleSecondMoment {Ω 𝕜 : Type*} [MeasurableSpace Ω] [RCLike 𝕜]
    {m n : ℕ} (P : Measure Ω) (X : Ω → Matrix (Fin m) (Fin n) 𝕜) : ℝ :=
  max ‖matrixExpectation P (fun ω => X ω * (X ω)ᴴ)‖
    ‖matrixExpectation P (fun ω => (X ω)ᴴ * X ω)‖

/-- The matrix sampling estimator `X̄_k = (1/k) ∑_{i=1}^k X_i`. -/
noncomputable def matrixSamplingEstimator {Ω 𝕜 : Type*} [RCLike 𝕜] {m n : ℕ} (k : ℕ)
    (Xs : Fin k → Ω → Matrix (Fin m) (Fin n) 𝕜) : Ω → Matrix (Fin m) (Fin n) 𝕜 :=
  fun ω => ((k : 𝕜)⁻¹) • ∑ i, Xs i ω

/-- Matrix sampling (matrix Bernstein) bound, mt2020::thm:mtx-sampling.
The field `𝔽` is ℝ or ℂ (`RCLike 𝕜`); `‖·‖` is the spectral norm (scoped L2
operator norm). `X` is a random matrix with `𝔼 X = B` and `‖X‖ ≤ R` (almost surely);
`X_1, …, X_k` are iid copies of `X`.

Domain: the source's bounds involve `log(m+n)` and the factor `(m+n)`, which
presuppose `m + n ≥ 1`. The degenerate case `m = n = 0` is not meaningfully covered
by the source (`log 0` is undefined), so the statement is restricted to
`0 < m + n` (hypothesis `hdim`). No other dimension restriction is imposed. -/
theorem matrix_sampling_bound
    {Ω 𝕜 : Type*} [MeasurableSpace Ω] [RCLike 𝕜] [MeasurableSpace 𝕜] [BorelSpace 𝕜]
    (P : Measure Ω) [IsProbabilityMeasure P]
    {m n : ℕ} (hdim : 0 < m + n) (B : Matrix (Fin m) (Fin n) 𝕜) (R : ℝ)
    (X : Ω → Matrix (Fin m) (Fin n) 𝕜) (hXmeas : Measurable X)
    (hmean : matrixExpectation P X = B)
    (hbound : ∀ᵐ ω ∂P, ‖X ω‖ ≤ R)
    (k : ℕ) (hk : 0 < k)
    (Xs : Fin k → Ω → Matrix (Fin m) (Fin n) 𝕜)
    (hindep : ProbabilityTheory.iIndepFun Xs P)
    (hident : ∀ i, ProbabilityTheory.IdentDistrib (Xs i) X P P) :
    (∫ ω, ‖matrixSamplingEstimator k Xs ω - B‖ ∂P ≤
        Real.sqrt (2 * perSampleSecondMoment P X * Real.log ((m : ℝ) + n) / k)
          + 2 * R * Real.log ((m : ℝ) + n) / (3 * k)) ∧
    (∀ t : ℝ, 0 ≤ t →
      P {ω | t ≤ ‖matrixSamplingEstimator k Xs ω - B‖} ≤
        ENNReal.ofReal (((m : ℝ) + n) *
          Real.exp ((-((k : ℝ) * t ^ 2 / 2)) /
            (perSampleSecondMoment P X + 2 * R * t / 3)))) := by
  sorry

end RNLAPhaseB.MtxSampling