import Mathlib

open MeasureTheory intervalIntegral
open scoped Real

namespace UniformMaxUMVUE

/-
Moments and UMVUE variance for the maximum of `n = m+1` i.i.d. Uniform(0, θ) draws.

The maximum `M = X_(n)` has density  f_M(t) = n · t^(n-1) / θ^n  on [0, θ].
Writing n = m+1 (so n ≥ 1), this is  (m+1) · t^m / θ^(m+1).

We verify, as real-analysis facts:
  • mean of M:           ∫₀^θ t · (m+1) t^m / θ^(m+1) dt = (m+1)/(m+2) · θ
  • second moment of M:  ∫₀^θ t² · (m+1) t^m / θ^(m+1) dt = (m+1)/(m+3) · θ²
and then derive algebraically:
  • Var(M) = (m+1) θ² / ((m+2)² (m+3))
  • Var(UMVUE) = Var( (n+1)/n · M ) = θ² / (n (n+2))   with n = m+1,
    i.e. = θ² / ((m+1)(m+3)).
These are the computational engine of the Lehmann–Scheffé UMVUE solution.
-/

variable (m : ℕ) (θ : ℝ)

/-- Mean of the maximum of `m+1` i.i.d. Uniform(0,θ). -/
theorem mean_max (hθ : 0 < θ) :
    ∫ t in (0:ℝ)..θ, t * ((m + 1) * t ^ m / θ ^ (m + 1))
      = (m + 1) / (m + 2) * θ := by
  have hθn : θ ^ (m + 1) ≠ 0 := pow_ne_zero _ (ne_of_gt hθ)
  -- rewrite integrand as a constant times t^(m+1)
  have hrw : ∀ t : ℝ, t * ((m + 1) * t ^ m / θ ^ (m + 1))
      = ((m + 1) / θ ^ (m + 1)) * t ^ (m + 1) := by
    intro t; rw [pow_succ]; ring
  simp_rw [hrw]
  rw [intervalIntegral.integral_const_mul, integral_pow]
  have hz : (0:ℝ) ^ (m + 1 + 1) = 0 := by simp
  rw [hz, pow_succ θ (m + 1)]
  push_cast
  field_simp
  ring

/-- Second moment of the maximum of `m+1` i.i.d. Uniform(0,θ). -/
theorem sq_moment_max (hθ : 0 < θ) :
    ∫ t in (0:ℝ)..θ, t ^ 2 * ((m + 1) * t ^ m / θ ^ (m + 1))
      = (m + 1) / (m + 3) * θ ^ 2 := by
  have hθn : θ ^ (m + 1) ≠ 0 := pow_ne_zero _ (ne_of_gt hθ)
  have hrw : ∀ t : ℝ, t ^ 2 * ((m + 1) * t ^ m / θ ^ (m + 1))
      = ((m + 1) / θ ^ (m + 1)) * t ^ (m + 2) := by
    intro t
    have : t ^ 2 * t ^ m = t ^ (m + 2) := by ring
    calc t ^ 2 * ((m + 1) * t ^ m / θ ^ (m + 1))
        = ((m + 1) / θ ^ (m + 1)) * (t ^ 2 * t ^ m) := by ring
      _ = ((m + 1) / θ ^ (m + 1)) * t ^ (m + 2) := by rw [this]
  simp_rw [hrw]
  rw [intervalIntegral.integral_const_mul, integral_pow]
  have hz : (0:ℝ) ^ (m + 2 + 1) = 0 := by simp
  have hpow : θ ^ (m + 2 + 1) = θ ^ (m + 1) * θ ^ 2 := by
    ring
  rw [hz, hpow]
  push_cast
  field_simp
  ring

/-- Variance of the maximum, from the two moments above (pure algebra). -/
theorem var_max (hθ : 0 < θ) :
    ((m + 1) / (m + 3) * θ ^ 2) - ((m + 1) / (m + 2) * θ) ^ 2
      = (m + 1) * θ ^ 2 / ((m + 2) ^ 2 * (m + 3)) := by
  have h2 : (m : ℝ) + 2 ≠ 0 := by positivity
  have h3 : (m : ℝ) + 3 ≠ 0 := by positivity
  field_simp
  ring

/-- Variance of the UMVUE  ((n+1)/n)·M  with n = m+1, i.e. scale factor (m+2)/(m+1).
    Equals θ² / (n(n+2)) = θ² / ((m+1)(m+3)). -/
theorem var_umvue (hθ : 0 < θ) :
    ((m + 2) / (m + 1)) ^ 2 * ((m + 1) * θ ^ 2 / ((m + 2) ^ 2 * (m + 3)))
      = θ ^ 2 / ((m + 1) * (m + 3)) := by
  have h1 : (m : ℝ) + 1 ≠ 0 := by positivity
  have h2 : (m : ℝ) + 2 ≠ 0 := by positivity
  have h3 : (m : ℝ) + 3 ≠ 0 := by positivity
  field_simp

end UniformMaxUMVUE