import Mathlib

open MeasureTheory ProbabilityTheory
open scoped MeasureTheory ProbabilityTheory

namespace StatsFormalization

variable {Ω : Type*} {mΩ : MeasurableSpace Ω} {μ : Measure Ω}

/-
Bias–variance decomposition of the mean squared error.

For a square-integrable estimator `est : Ω → ℝ` of a fixed real parameter `θ`
on a probability space `(Ω, μ)`,

    MSE(est) = Var(est) + Bias(est)²,

where  MSE(est) = 𝔼[(est - θ)²]  and  Bias(est) = 𝔼[est] - θ.

This is the foundational identity of point estimation theory: it decomposes the
mean squared error into estimator variance plus squared bias. The unbiased case
(Bias = 0) yields MSE = Var, underpinning minimum-variance unbiased estimation.
-/

theorem mse_eq_variance_add_bias_sq
    [IsProbabilityMeasure μ] (θ : ℝ) {est : Ω → ℝ} (hest : MemLp est 2 μ) :
    μ[fun ω => (est ω - θ) ^ 2] = Var[est; μ] + (μ[est] - θ) ^ 2 := by
  -- The centered-at-θ variable Y = est - θ is also in L².
  have hYmem : MemLp (fun ω => est ω - θ) 2 μ := hest.sub (memLp_const θ)
  -- Computational formula for variance: Var[Y] = 𝔼[Y²] - 𝔼[Y]².
  have hcomp : Var[fun ω => est ω - θ; μ]
      = μ[fun ω => (est ω - θ) ^ 2] - (μ[fun ω => est ω - θ]) ^ 2 := by
    have h := variance_eq_sub (μ := μ) hYmem
    simpa using h
  -- Variance is invariant under subtracting the constant θ.
  have hshift : Var[fun ω => est ω - θ; μ] = Var[est; μ] :=
    variance_sub_const hest.aestronglyMeasurable θ
  -- Linearity of expectation: 𝔼[est - θ] = 𝔼[est] - θ.
  have hmean : μ[fun ω => est ω - θ] = μ[est] - θ := by
    have hint : Integrable est μ := hest.integrable (by norm_num)
    rw [integral_sub hint (integrable_const θ)]
    simp
  -- Rearrange the computational formula and substitute.
  have : μ[fun ω => (est ω - θ) ^ 2]
      = Var[fun ω => est ω - θ; μ] + (μ[fun ω => est ω - θ]) ^ 2 := by
    rw [hcomp]; ring
  rw [this, hshift, hmean]

end StatsFormalization