import Mathlib

namespace RobustOutcomeTest

/-
Formalization of the logical core of Theorem 1 of
  Gaebler & Goel, "A simple, statistically robust test of discrimination" (PNAS 2025).

Setup (following SI Appendix, Section 1.A). Two groups g ∈ {0,1}. For a threshold t,
let `decRate g t = Pr(R ≥ t | G = g)` be the decision rate and
`outRate g t = E[R | R ≥ t, G = g]` the outcome rate (tail-conditional mean of risk).
Decisions use group-specific thresholds t₀, t₁, so the realized decision/outcome
rates are `decRate g (t g)` and `outRate g (t g)`.

The proof of Theorem 1 imports two consequences of the MLRP from Shaked & Shanthikumar
(Thms 1.C.1, 1.C.5), stated at a COMMON threshold, distinguishing the lower-base-rate
group `lwr` from the higher-base-rate group `upr`:

  • Eq. (3)  stochastic dominance:   decRate upr t ≥ decRate lwr t      (∀ t)
  • Eq. (4)  tail-mean ordering:      outRate lwr t ≤ outRate upr t      (∀ t)

together with two monotonicity facts about threshold rules:

  • decision rate is non-increasing in the threshold (lowering t raises the decision rate)
  • outcome rate (tail mean) is non-decreasing in the threshold (lowering t lowers it)

We take these four facts as hypotheses — exactly as the paper imports Eqs. (3),(4) from
S&S rather than re-deriving them — and prove the paper's contrapositive:
if t₁ ≤ t₀ then the benchmark test or the standard outcome test fails to fire.
Hence (Theorem 1) if BOTH tests fire, then t₀ < t₁.

This formalizes the COMBINATION step that is the actual content of Theorem 1.
It does NOT formalize the MLRP ⇒ stochastic-dominance / tail-mean theory (Eqs. 3,4
themselves), which the paper likewise imports; that measure-theoretic development is
not in Mathlib and is a separate undertaking.
-/

/-- The lower-base-rate group `lwr` and higher-base-rate group `upr`. -/
structure MLRPInputs (decRate outRate : Bool → ℝ → ℝ) (lwr upr : Bool) : Prop where
  /-- Eq. (3): stochastic dominance — at any common threshold the higher-base-rate
      group has the larger decision rate. -/
  dom : ∀ s : ℝ, decRate lwr s ≤ decRate upr s
  /-- Eq. (4): at any common threshold the lower-base-rate group has the smaller tail mean. -/
  tail : ∀ s : ℝ, outRate lwr s ≤ outRate upr s

/-- Threshold-rule monotonicity of the two rates. -/
structure ThresholdMono (decRate outRate : Bool → ℝ → ℝ) : Prop where
  /-- Lowering the threshold cannot decrease the decision rate. -/
  decAnti : ∀ (g : Bool) {a b : ℝ}, a ≤ b → decRate g b ≤ decRate g a
  /-- Lowering the threshold cannot increase the outcome rate (tail mean). -/
  outMono : ∀ (g : Bool) {a b : ℝ}, a ≤ b → outRate g a ≤ outRate g b

/-- **Core of Theorem 1 (contrapositive).**
If group `g0` is the lower-base-rate group and `t g1 ≤ t g0`, then the benchmark test
fails to fire: the decision rate of `g0` does not strictly exceed that of `g1`. -/
theorem benchmark_correct_when_g0_lower
    {decRate outRate : Bool → ℝ → ℝ} {t : Bool → ℝ} {g0 g1 : Bool}
    (hmono : ThresholdMono decRate outRate)
    (hmlrp : MLRPInputs decRate outRate g0 g1)   -- g0 = lwr, g1 = upr
    (hle : t g1 ≤ t g0) :
    decRate g0 (t g0) ≤ decRate g1 (t g1) := by
  -- stochastic dominance at the common threshold t g0: decRate g0 (t g0) ≤ decRate g1 (t g0)
  have hdom : decRate g0 (t g0) ≤ decRate g1 (t g0) := hmlrp.dom (t g0)
  -- lowering g1's threshold from t g0 to t g1 (since t g1 ≤ t g0) raises g1's decision rate
  have hlow : decRate g1 (t g0) ≤ decRate g1 (t g1) := hmono.decAnti g1 hle
  exact le_trans hdom hlow

/-- **Core of Theorem 1 (contrapositive), other case.**
If group `g0` is the higher-base-rate group and `t g1 ≤ t g0`, then the standard outcome
test fails to fire: the outcome rate of `g1` does not strictly exceed that of `g0`. -/
theorem outcome_correct_when_g0_higher
    {decRate outRate : Bool → ℝ → ℝ} {t : Bool → ℝ} {g0 g1 : Bool}
    (hmono : ThresholdMono decRate outRate)
    (hmlrp : MLRPInputs decRate outRate g1 g0)   -- g1 = lwr, g0 = upr
    (hle : t g1 ≤ t g0) :
    outRate g1 (t g1) ≤ outRate g0 (t g0) := by
  -- lowering g1's threshold from t g0 to t g1 lowers g1's tail mean
  have hlow : outRate g1 (t g1) ≤ outRate g1 (t g0) := hmono.outMono g1 hle
  -- tail-mean ordering at the common threshold t g0 (g1 = lwr ≤ g0 = upr)
  have htail : outRate g1 (t g0) ≤ outRate g0 (t g0) := hmlrp.tail (t g0)
  exact le_trans hlow htail

/-- **Theorem 1 (robust outcome test).**
Suppose one of the two groups is the lower-base-rate group (either `g0` or `g1`), with the
corresponding MLRP inputs, and threshold-rule monotonicity holds. If BOTH tests fire —
the benchmark test (`decRate g1 (t g1) < decRate g0 (t g0)`, group 1 has the strictly lower
decision rate) AND the standard outcome test (`outRate g0 (t g0) < outRate g1 (t g1)`,
group 1 has the strictly higher outcome rate) — then group 1 is held to a strictly higher
threshold: `t g0 < t g1`. -/
theorem robust_outcome_test
    {decRate outRate : Bool → ℝ → ℝ} {t : Bool → ℝ} {g0 g1 : Bool}
    (hmono : ThresholdMono decRate outRate)
    -- the MLRP holds with one group the lower-base-rate group; we don't know which
    (hcase : MLRPInputs decRate outRate g0 g1 ∨ MLRPInputs decRate outRate g1 g0)
    -- benchmark test fires: group 1's decision rate is strictly lower
    (hbench : decRate g1 (t g1) < decRate g0 (t g0))
    -- standard outcome test fires: group 1's outcome rate is strictly higher
    (houtcome : outRate g0 (t g0) < outRate g1 (t g1)) :
    t g0 < t g1 := by
  -- prove the contrapositive: ¬ (t g1 ≤ t g0)
  by_contra hcon
  push_neg at hcon          -- hcon : t g1 ≤ t g0
  rcases hcase with hlwr0 | hlwr1
  · -- g0 is the lower-base-rate group ⇒ benchmark test must be correct, contradicting hbench
    have : decRate g0 (t g0) ≤ decRate g1 (t g1) :=
      benchmark_correct_when_g0_lower hmono hlwr0 hcon
    exact (not_lt.mpr this) hbench
  · -- g0 is the higher-base-rate group ⇒ outcome test must be correct, contradicting houtcome
    have : outRate g1 (t g1) ≤ outRate g0 (t g0) :=
      outcome_correct_when_g0_higher hmono hlwr1 hcon
    exact (not_lt.mpr this) houtcome

end RobustOutcomeTest