import Mathlib

set_option autoImplicit false

namespace DoubleML

/-
Neyman orthogonality for the partially linear model — the algebraic core of
Double/Debiased Machine Learning (Chernozhukov et al., 2018).

Random variables are elements of a commutative ℝ-algebra `R` (think: functions of
the data, pointwise ops); expectation is an ℝ-linear functional `E : R →ₗ[ℝ] ℝ`.
The target θ enters as `algebraMap ℝ R θ` (the constant random variable θ).

Partially linear model:  Y = (θ₀)·D + g₀ + ε,   D = m₀ + V, with the moments below.

DML orthogonal score (partials out BOTH nuisances):  ψ(θ,g,m) = (Y - θ·D - g)(D - m).
Result: at the truth under nuisance perturbations g = g₀+a, m = m₀+b, E[ψ] = E[a·b]
— SECOND order, so first-order (Gâteaux) derivative vanishes (Neyman orthogonality).
The naive score (Y - θ·D - g)·D is FIRST order in a, hence biased by nuisance error.

SCOPE. This formalizes the orthogonality IDENTITY (why DML works) in a linear-functional
model of expectation, with the model's conditional-mean-zero moments as hypotheses. It
does NOT formalize the DML asymptotic guarantee (√n-consistency / asymptotic normality
under cross-fitting and o(n^{-1/4}) nuisance rates), which needs empirical-process/CLT
machinery not developed here.
-/

variable {R : Type*} [CommRing R] [Algebra ℝ R]

/-- A partially linear model with its orthogonality moment conditions. -/
structure PLR (E : R →ₗ[ℝ] ℝ) where
  θ₀ : ℝ
  D : R
  ε : R
  V : R
  g₀ : R
  m₀ : R
  Y : R
  hY : Y = (algebraMap ℝ R θ₀) * D + g₀ + ε
  hD : D = m₀ + V
  m_eps_V : E (ε * V) = 0
  m_eps : ∀ b : R, E (ε * b) = 0
  m_V : ∀ a : R, E (a * V) = 0

/-- The DML orthogonal (partialling-out) score at parameter `θ` and nuisances `g, m`. -/
def orthScore {E : R →ₗ[ℝ] ℝ} (M : PLR E) (θ : ℝ) (g m : R) : R :=
  (M.Y - (algebraMap ℝ R θ) * M.D - g) * (M.D - m)

/-- The naive score (partials out only the outcome nuisance `g`). -/
def naiveScore {E : R →ₗ[ℝ] ℝ} (M : PLR E) (θ : ℝ) (g : R) : R :=
  (M.Y - (algebraMap ℝ R θ) * M.D - g) * M.D

/-- **Substitution lemma.** At the true `θ₀` and perturbed nuisances `g₀+a, m₀+b`,
the orthogonal score reduces to `(ε - a)(V - b)`. -/
theorem orthScore_truth {E : R →ₗ[ℝ] ℝ} (M : PLR E) (a b : R) :
    orthScore M M.θ₀ (M.g₀ + a) (M.m₀ + b) = (M.ε - a) * (M.V - b) := by
  unfold orthScore
  rw [M.hY, M.hD]
  ring

/-- **Neyman orthogonality.** The expected orthogonal score at the truth, under nuisance
perturbations `(a, b)`, equals `E[a*b]` — second order in the perturbations. -/
theorem orthScore_expectation {E : R →ₗ[ℝ] ℝ} (M : PLR E) (a b : R) :
    E (orthScore M M.θ₀ (M.g₀ + a) (M.m₀ + b)) = E (a * b) := by
  rw [orthScore_truth]
  have hexp : (M.ε - a) * (M.V - b) = M.ε * M.V - M.ε * b - a * M.V + a * b := by ring
  rw [hexp, map_add, map_sub, map_sub, M.m_eps_V, M.m_eps b, M.m_V a]
  simp

/-- **First-order insensitivity.** Perturbing only `g` (b = 0) gives expected score 0. -/
theorem orthScore_firstOrder_g {E : R →ₗ[ℝ] ℝ} (M : PLR E) (a : R) :
    E (orthScore M M.θ₀ (M.g₀ + a) (M.m₀ + 0)) = 0 := by
  rw [orthScore_expectation]; simp

/-- **First-order insensitivity.** Perturbing only `m` (a = 0) gives expected score 0. -/
theorem orthScore_firstOrder_m {E : R →ₗ[ℝ] ℝ} (M : PLR E) (b : R) :
    E (orthScore M M.θ₀ (M.g₀ + 0) (M.m₀ + b)) = 0 := by
  rw [orthScore_expectation]; simp

/-- **Naive score is first-order sensitive.** At the truth with outcome nuisance perturbed
by `a`, the expected naive score equals `- E[a * m₀]` — first order in `a`, hence biased. -/
theorem naiveScore_expectation {E : R →ₗ[ℝ] ℝ} (M : PLR E) (a : R) :
    E (naiveScore M M.θ₀ (M.g₀ + a)) = - E (a * M.m₀) := by
  unfold naiveScore
  rw [M.hY]
  have hexp : ((algebraMap ℝ R M.θ₀) * M.D + M.g₀ + M.ε
        - (algebraMap ℝ R M.θ₀) * M.D - (M.g₀ + a)) * M.D
      = M.ε * M.D - a * M.D := by ring
  rw [hexp]
  have haD : a * M.D = a * M.m₀ + a * M.V := by rw [M.hD]; ring
  rw [map_sub, haD, map_add, M.m_V a]
  have heD : M.ε * M.D = M.ε * M.m₀ + M.ε * M.V := by rw [M.hD]; ring
  rw [heD, map_add, M.m_eps M.m₀, M.m_eps_V]
  simp

end DoubleML