import CubicPathwidth.CutSweep.Defs
import CubicPathwidth.C3.GaussianTools

/-!
# Two Gaussians on opposite sides of a threshold (`lem:crossing`(b))

With `U = ⟪a, x⟫`, `V = ⟪b, x⟫` and `ρ = ⟪a, b⟫`, the event
`min U V ≤ t < max U V` has probability at most `arccos ρ / π`.

The proof follows `lem:crossing` of the manuscript. Writing `P = ⟪a + b, x⟫` and
`Q = ⟪a - b, x⟫` (twice `S` and `D`), the event implies `(U - t)(V - t) ≤ 0`,
that is `|P - 2t| ≤ |Q|`. The coordinates of `x` along the orthonormal pair
`(a - b)/‖a - b‖`, `(a + b)/‖a + b‖` are independent standard normals, so the
probability is an integral over `Q` of the mass of an interval of fixed length
under the law of `P`; the centered interval has the largest mass
(`gaussianReal_abs_sub_le`). This bounds the probability by that of
`|P| ≤ |Q|`, i.e. of `UV ≤ 0`, which is `arccos ρ / π` by the sign
disagreement identity (part (a)).
-/

noncomputable section
open scoped RealInnerProductSpace
open MeasureTheory ProbabilityTheory Set

namespace CubicPathwidth
namespace CutSweep

/-- The standard normal density is monotone in `|x|`. -/
theorem gaussianPDFReal_anti {x y : ℝ} (h : |y| ≤ |x|) :
    gaussianPDFReal 0 1 x ≤ gaussianPDFReal 0 1 y := by
  rw [gaussianPDFReal_def]
  apply mul_le_mul_of_nonneg_left _ (by positivity)
  apply Real.exp_le_exp.mpr
  have hsq : y ^ 2 ≤ x ^ 2 := sq_le_sq.mpr h
  simp only [sub_zero, NNReal.coe_one, mul_one]
  linarith

theorem gaussianPDFReal_neg (x : ℝ) : gaussianPDFReal 0 1 (-x) = gaussianPDFReal 0 1 x :=
  le_antisymm (gaussianPDFReal_anti (by rw [abs_neg])) (gaussianPDFReal_anti (by rw [abs_neg]))

theorem gaussianReal_real_Icc {a b : ℝ} (hab : a ≤ b) :
    (gaussianReal 0 1).real (Icc a b) = ∫ x in a..b, gaussianPDFReal 0 1 x := by
  rw [Measure.real, gaussianReal_apply_eq_integral 0 one_ne_zero,
    ENNReal.toReal_ofReal (setIntegral_nonneg measurableSet_Icc
      (fun x _ => gaussianPDFReal_nonneg 0 1 x)),
    integral_Icc_eq_integral_Ioc, intervalIntegral.integral_of_le hab]

theorem intervalIntegrable_pdf (a b : ℝ) :
    IntervalIntegrable (gaussianPDFReal 0 1) volume a b :=
  (integrable_gaussianPDFReal 0 1).intervalIntegrable

/-- Shifting a symmetric window to a nonnegative center loses mass. -/
theorem window_le_of_nonneg {c w : ℝ} (hc : 0 ≤ c) (hw : 0 ≤ w) :
    (∫ x in (c - w)..(c + w), gaussianPDFReal 0 1 x) ≤
      ∫ x in (-w)..w, gaussianPDFReal 0 1 x := by
  set p := gaussianPDFReal 0 1
  have hi := intervalIntegrable_pdf
  have hsplit : (∫ x in (c - w)..(c + w), p x) =
      (∫ x in (c - w)..(-w), p x) + (∫ x in (-w)..w, p x) + ∫ x in w..(c + w), p x := by
    rw [intervalIntegral.integral_add_adjacent_intervals (hi _ _) (hi _ _),
      intervalIntegral.integral_add_adjacent_intervals (hi _ _) (hi _ _)]
  have hright : (∫ x in w..(c + w), p x) = ∫ u in (0 : ℝ)..c, p (u + w) := by
    rw [intervalIntegral.integral_comp_add_right]
    simp [add_comm]
  have hleft : (∫ x in (c - w)..(-w), p x) = -∫ u in (0 : ℝ)..c, p (u - w) := by
    rw [intervalIntegral.integral_comp_sub_right, intervalIntegral.integral_symm]
    simp
  have hmono : (∫ u in (0 : ℝ)..c, p (u + w)) ≤ ∫ u in (0 : ℝ)..c, p (u - w) := by
    apply intervalIntegral.integral_mono_on hc
    · simpa using (hi (0 + w) (c + w)).comp_add_right w
    · simpa using (hi (0 - w) (c - w)).comp_sub_right w
    · intro u hu
      apply gaussianPDFReal_anti
      rw [abs_of_nonneg (by linarith [hu.1] : 0 ≤ u + w), abs_le]
      constructor <;> linarith [hu.1]
  rw [hsplit, hright, hleft]
  linarith

/-- A centered window has the largest standard normal mass. -/
theorem window_le {c w : ℝ} (hw : 0 ≤ w) :
    (∫ x in (c - w)..(c + w), gaussianPDFReal 0 1 x) ≤
      ∫ x in (-w)..w, gaussianPDFReal 0 1 x := by
  rcases le_total 0 c with hc | hc
  · exact window_le_of_nonneg hc hw
  · have hsym : (∫ x in (c - w)..(c + w), gaussianPDFReal 0 1 x) =
        ∫ x in (-c - w)..(-c + w), gaussianPDFReal 0 1 x := by
      have h := intervalIntegral.integral_comp_neg (a := c - w) (b := c + w)
        (f := gaussianPDFReal 0 1)
      simp only [gaussianPDFReal_neg] at h
      rw [h]
      congr 1 <;> ring
    rw [hsym]
    exact window_le_of_nonneg (by linarith) hw

/-- One-dimensional step of `lem:crossing`(b): for `α > 0`, the event `|α g - t| ≤ h`
is at most as likely as `|α g| ≤ h`. -/
theorem gaussianReal_abs_sub_le {α : ℝ} (hα : 0 < α) (t : ℝ) {h : ℝ} (hh : 0 ≤ h) :
    gaussianReal 0 1 {y | |α * y - t| ≤ h} ≤ gaussianReal 0 1 {y | |α * y| ≤ h} := by
  have hset : ∀ s : ℝ, {y : ℝ | |α * y - s| ≤ h} = Icc (s / α - h / α) (s / α + h / α) := by
    intro s
    ext y
    rw [mem_ofPred_eq, abs_le, mem_Icc, ← sub_div, ← add_div, div_le_iff₀ hα, le_div_iff₀ hα]
    constructor <;> rintro ⟨h1, h2⟩ <;> constructor <;> nlinarith
  have h0 : {y : ℝ | |α * y| ≤ h} = {y : ℝ | |α * y - 0| ≤ h} := by simp
  rw [h0, hset, hset, ← ENNReal.toReal_le_toReal (measure_ne_top _ _) (measure_ne_top _ _)]
  have hw : 0 ≤ h / α := div_nonneg hh hα.le
  change (gaussianReal 0 1).real _ ≤ (gaussianReal 0 1).real _
  rw [gaussianReal_real_Icc (by linarith), gaussianReal_real_Icc (by linarith)]
  simpa using window_le (c := t / α) hw

/-- **`lem:crossing`(b).** -/
theorem crossingClaim : CrossingClaim := by
  intro m a b ha hb t
  set μ := stdGaussian (GaussianField.Sample m)
  have haa : ⟪a, a⟫ = 1 := by rw [real_inner_self_eq_norm_sq, ha]; norm_num
  have hbb : ⟪b, b⟫ = 1 := by rw [real_inner_self_eq_norm_sq, hb]; norm_num
  have hnonneg : 0 ≤ Real.arccos ⟪a, b⟫ / Real.pi :=
    div_nonneg (Real.arccos_nonneg _) Real.pi_pos.le
  by_cases hq : a - b = 0
  · have hab : a = b := sub_eq_zero.mp hq
    subst hab
    have hempty : {x : GaussianField.Sample m |
        min ⟪a, x⟫ ⟪a, x⟫ ≤ t ∧ t < max ⟪a, x⟫ ⟪a, x⟫} = ∅ := by
      ext x
      simp only [min_self, max_self, mem_ofPred_eq, mem_empty_iff_false, iff_false, not_and,
        not_lt]
      exact fun h => h
    rw [hempty, measureReal_empty]
    exact hnonneg
  by_cases hp : a + b = 0
  · have hba : b = -a := eq_neg_of_add_eq_zero_right hp
    rw [hba, inner_neg_right, haa, Real.arccos_neg_one, div_self Real.pi_ne_zero]
    exact measureReal_le_one
  -- the orthonormal pair `f = q/‖q‖`, `e = p/‖p‖`
  set p := a + b with hpdef
  set q := a - b with hqdef
  have hqp : ⟪q, p⟫ = 0 := by
    simp only [p, q, inner_sub_left, inner_add_right, haa, hbb, real_inner_comm a b]
    ring
  set e := ‖p‖⁻¹ • p
  set f := ‖q‖⁻¹ • q
  have hee : ⟪e, e⟫ = 1 := GaussianAngle.normalized_inner_self hp
  have hff : ⟪f, f⟫ = 1 := GaussianAngle.normalized_inner_self hq
  have hfe : ⟪f, e⟫ = 0 := by
    simp only [e, f, real_inner_smul_left, real_inner_smul_right, hqp, mul_zero]
  have hmap := GaussianAngle.orthogonal_pair_map hff hee hfe
  have hpn : ‖p‖ ≠ 0 := norm_ne_zero_iff.mpr hp
  have hqn : ‖q‖ ≠ 0 := norm_ne_zero_iff.mpr hq
  have hP : ∀ x, ⟪p, x⟫ = ‖p‖ * ⟪e, x⟫ := by
    intro x
    simp only [e, real_inner_smul_left]
    field_simp
  have hQ : ∀ x, ⟪q, x⟫ = ‖q‖ * ⟪f, x⟫ := by
    intro x
    simp only [f, real_inner_smul_left]
    field_simp
  have hPx : ∀ x, ⟪p, x⟫ = ⟪a, x⟫ + ⟪b, x⟫ := fun x => by simp only [p, inner_add_left]
  have hQx : ∀ x, ⟪q, x⟫ = ⟪a, x⟫ - ⟪b, x⟫ := fun x => by simp only [q, inner_sub_left]
  set Φ : GaussianField.Sample m → ℝ × ℝ := fun x => (⟪f, x⟫, ⟪e, x⟫) with hΦdef
  have hΦ : Measurable Φ := by fun_prop
  set S : Set (ℝ × ℝ) := {z | |‖p‖ * z.2 - 2 * t| ≤ |‖q‖ * z.1|}
  set S' : Set (ℝ × ℝ) := {z | |‖p‖ * z.2| ≤ |‖q‖ * z.1|}
  have hS : MeasurableSet S := measurableSet_le (by fun_prop) (by fun_prop)
  have hS' : MeasurableSet S' := measurableSet_le (by fun_prop) (by fun_prop)
  -- the sign-disagreement event
  set D : Set (GaussianField.Sample m) := {x | ¬(0 < ⟪a, x⟫ ↔ 0 < ⟪b, x⟫)}
  have hD : MeasurableSet D :=
    ((measurableSet_lt measurable_const (by fun_prop)).iff
      (measurableSet_lt measurable_const (by fun_prop))).compl
  -- step 1: the event implies `(U - t)(V - t) ≤ 0`, i.e. `|P - 2t| ≤ |Q|`
  have h1 : {x | min ⟪a, x⟫ ⟪b, x⟫ ≤ t ∧ t < max ⟪a, x⟫ ⟪b, x⟫} ⊆ Φ ⁻¹' S := by
    intro x hx
    simp only [mem_ofPred_eq, min_le_iff, lt_max_iff] at hx
    have hprod : (⟪a, x⟫ - t) * (⟪b, x⟫ - t) ≤ 0 := by
      rcases hx with ⟨h | h, h' | h'⟩ <;> nlinarith
    show |‖p‖ * ⟪e, x⟫ - 2 * t| ≤ |‖q‖ * ⟪f, x⟫|
    rw [← hP, ← hQ, hPx, hQx, ← sq_le_sq]
    nlinarith
  -- step 2: conditionally on `Q`, compare with the centered window
  have h2 : GaussianAngle.normal2 S ≤ GaussianAngle.normal2 S' := by
    rw [GaussianAngle.normal2, Measure.prod_apply hS, Measure.prod_apply hS']
    apply lintegral_mono
    intro d
    exact gaussianReal_abs_sub_le (norm_pos_iff.mpr hp) (2 * t) (abs_nonneg (‖q‖ * d))
  -- step 3: `|P| ≤ |Q|` means `UV ≤ 0`, which forces sign disagreement a.s.
  have h3 : Φ ⁻¹' S' ⊆ D ∪ {x | ⟪a, x⟫ = 0} ∪ {x | ⟪b, x⟫ = 0} := by
    intro x hx
    change |‖p‖ * ⟪e, x⟫| ≤ |‖q‖ * ⟪f, x⟫| at hx
    rw [← hP, ← hQ, hPx, hQx, ← sq_le_sq] at hx
    have huv : ⟪a, x⟫ * ⟪b, x⟫ ≤ 0 := by nlinarith
    simp only [mem_union, mem_ofPred_eq, D]
    by_cases hu : ⟪a, x⟫ = 0
    · exact Or.inl (Or.inr hu)
    by_cases hv : ⟪b, x⟫ = 0
    · exact Or.inr hv
    refine Or.inl (Or.inl fun hiff => ?_)
    rcases lt_or_gt_of_ne hu with hu' | hu'
    · have hv' : ⟪b, x⟫ < 0 :=
        lt_of_le_of_ne (not_lt.mp fun h => absurd (hiff.mpr h) (not_lt.mpr hu'.le)) hv
      nlinarith [mul_pos_of_neg_of_neg hu' hv']
    · have hv' : 0 < ⟪b, x⟫ := hiff.mp hu'
      nlinarith [mul_pos hu' hv']
  have ha0 : a ≠ 0 := by rintro rfl; simp at ha
  have hb0 : b ≠ 0 := by rintro rfl; simp at hb
  have hmain : μ {x | min ⟪a, x⟫ ⟪b, x⟫ ≤ t ∧ t < max ⟪a, x⟫ ⟪b, x⟫} ≤ μ D := by
    calc μ {x | min ⟪a, x⟫ ⟪b, x⟫ ≤ t ∧ t < max ⟪a, x⟫ ⟪b, x⟫}
        ≤ μ (Φ ⁻¹' S) := measure_mono h1
      _ = GaussianAngle.normal2 S := by rw [← hmap, Measure.map_apply hΦ hS]
      _ ≤ GaussianAngle.normal2 S' := h2
      _ = μ (Φ ⁻¹' S') := by rw [← hmap, Measure.map_apply hΦ hS']
      _ ≤ μ (D ∪ {x | ⟪a, x⟫ = 0} ∪ {x | ⟪b, x⟫ = 0}) := measure_mono h3
      _ ≤ μ D + μ {x | ⟪a, x⟫ = 0} + μ {x | ⟪b, x⟫ = 0} :=
          (measure_union_le _ _).trans (by gcongr; exact measure_union_le _ _)
      _ = μ D := by rw [C3.measure_inner_eq_zero ha0, C3.measure_inner_eq_zero hb0]; simp
  have hDval : μ.real D = Real.arccos ⟪a, b⟫ / Real.pi := by
    rw [← GaussianAngle.unit_sign_disagreement haa hbb, ← integral_indicator_one hD]
    apply integral_congr_ae
    refine ae_of_all _ fun x => ?_
    by_cases hx : (0 < ⟪a, x⟫ ↔ 0 < ⟪b, x⟫)
    · have : x ∉ D := fun h => h hx
      simp [GaussianAngle.disagreement, hx, this]
    · have : x ∈ D := hx
      simp [GaussianAngle.disagreement, hx, this]
  rw [← hDval]
  exact ENNReal.toReal_mono (measure_ne_top _ _) hmain

end CutSweep
end CubicPathwidth

#print axioms CubicPathwidth.CutSweep.crossingClaim