import CubicPathwidth.Sw4.Defs
import CubicPathwidth.T2.Bridge
import CubicPathwidth.C3.GaussianTools

/-!
# Lowering the floors of `Ψ` (Q2, the part of `lem:switch-cont` used)

`psiLower : PsiLowerClaim`. Route:
* `psi_eq`: by the Gaussian bridge, `Ψ(Gram(a, b))` is
  `Pr(⟪a,x⟫ ≤ 0, N ≥ 1) + Pr(⟪a,x⟫ ≤ 0, N ≥ 3)` with `N = #{i : ⟪b_i,x⟫ > 0}`,
  since `⌈N/2⌉ = 1{N ≥ 1} + 1{N ≥ 3}` for `N ≤ 4`.
* `zeroProb_le`, `zeroProb'_le`: by the symmetry `x ↦ -x` and since every
  `{⟪u,x⟫ = 0}` with `u ≠ 0` is null, `zeroProb` and `zeroProb'` are at most the same sum.
* `exists_raise`: a unit `b` with `⟪a,b⟫ ≥ (1-η) r` is moved to a unit `b'` with
  `⟪a,b'⟫ ≥ r` and `‖b - b'‖² ≤ 2η`.
* `real_diffEv_le`: the signs of `⟪b,x⟫` and `⟪b',x⟫` differ with probability at most
  `3t/2 + ‖b - b'‖²/t²` (a small band around zero, plus Chebyshev).
* `real_ev_le`: each event changes by at most the sum of the sign-difference probabilities.
-/

noncomputable section
open scoped RealInnerProductSpace
open MeasureTheory ProbabilityTheory

namespace CubicPathwidth
namespace Sw4
namespace PsiLower

open GaussianField (Sample)

variable {m : ℕ}

/-! ## Events -/

/-- `N(x) = #{i : ⟪b_i, x⟫ > 0}`. -/
def cnt {d : ℕ} (b : Fin d → Sample m) (x : Sample m) : ℕ := by
  classical
  exact (Finset.univ.filter fun i => 0 < ⟪b i, x⟫).card

/-- `{⟪a, x⟫ ≤ 0, N(x) ≥ k}`. -/
def ev {d : ℕ} (k : ℕ) (a : Sample m) (b : Fin d → Sample m) : Set (Sample m) :=
  {x | ⟪a, x⟫ ≤ 0 ∧ k ≤ cnt b x}

/-- The signs of `⟪b, x⟫` and `⟪b', x⟫` (positive or not) differ. -/
def diffEv (b b' : Sample m) : Set (Sample m) :=
  {x | ¬ (0 < ⟪b, x⟫ ↔ 0 < ⟪b', x⟫)}

/-- Some inner product vanishes (a null set for nonzero vectors). -/
def zeroEv {d : ℕ} (a : Sample m) (b : Fin d → Sample m) : Set (Sample m) :=
  {x | ⟪a, x⟫ = 0} ∪ ⋃ i, {x | ⟪b i, x⟫ = 0}

/-! ## Measure-theoretic helpers -/

theorem measurableSet_card_ge {X : Type*} [MeasurableSpace X] {k : ℕ} (p : Fin k → X → Prop)
    [∀ i x, Decidable (p i x)] (hp : ∀ i, MeasurableSet {x | p i x}) (n : ℕ) :
    MeasurableSet {x | n ≤ (Finset.univ.filter fun i => p i x).card} := by
  have hf : Measurable fun x => ∑ i, (if p i x then 1 else 0 : ℕ) :=
    Finset.measurable_sum _ fun i _ => Measurable.ite (hp i) measurable_const measurable_const
  have : {x | n ≤ (Finset.univ.filter fun i => p i x).card} =
      (fun x => ∑ i, (if p i x then 1 else 0 : ℕ)) ⁻¹' Set.Ici n := by
    ext x
    simp only [Set.mem_ofPred_eq, Set.mem_preimage, Set.mem_Ici, Finset.card_filter]
  rw [this]
  exact hf MeasurableSet.of_discrete

/-- Symmetry of the standard Gaussian, for every set. -/
theorem real_neg (s : Set (Sample m)) : (μ m).real (Neg.neg ⁻¹' s) = (μ m).real s := by
  have h : (μ m).map (MeasurableEquiv.neg (Sample m)) = μ m :=
    stdGaussian_map (E := Sample m) (LinearIsometryEquiv.neg ℝ (E := Sample m))
  conv_rhs => rw [← h]
  rw [measureReal_def, measureReal_def, MeasurableEquiv.map_apply]
  rfl

theorem real_zeroEv {d : ℕ} {a : Sample m} {b : Fin d → Sample m} (ha : ‖a‖ = 1)
    (hb : ∀ i, ‖b i‖ = 1) : (μ m).real (zeroEv a b) = 0 := by
  have hne : ∀ v : Sample m, ‖v‖ = 1 → v ≠ 0 := fun v hv h => by simp [h] at hv
  rw [measureReal_def, zeroEv, measure_union_null (C3.measure_inner_eq_zero (hne a ha))
    (measure_iUnion_null fun i => C3.measure_inner_eq_zero (hne _ (hb i)))]
  rfl

theorem real_le_of_subset_union {s t z : Set (Sample m)} (h : s ⊆ t ∪ z)
    (hz : (μ m).real z = 0) : (μ m).real s ≤ (μ m).real t :=
  (measureReal_mono h).trans ((measureReal_union_le t z).trans (by rw [hz, add_zero]))

/-! ## `zeroProb` and `zeroProb'` against the two events -/

theorem ev_one_eq {d : ℕ} (a : Sample m) (b : Fin d → Sample m) :
    {x | ⟪a, x⟫ ≤ 0 ∧ ∃ i, 0 < ⟪b i, x⟫} = ev 1 a b := by
  ext x
  simp [ev, cnt, Nat.one_le_iff_ne_zero]

theorem zeroProb_le {d : ℕ} {a : Sample m} {b : Fin d → Sample m} (ha : ‖a‖ = 1)
    (hb : ∀ i, ‖b i‖ = 1) :
    zeroProb a b ≤ (μ m).real (ev 1 a b) + (μ m).real (ev 3 a b) := by
  rw [zeroProb, ev_one_eq]
  refine add_le_add le_rfl ?_
  rw [← real_neg]
  refine real_le_of_subset_union (z := zeroEv a b) ?_ (real_zeroEv ha hb)
  intro x hx
  by_cases hz : x ∈ zeroEv a b
  · exact Or.inr hz
  left
  simp only [zeroEv, Set.mem_union, Set.mem_ofPred_eq, Set.mem_iUnion, not_or, not_exists] at hz
  simp only [Set.mem_preimage, Set.mem_ofPred_eq, inner_neg_right] at hx
  refine ⟨by linarith [hx.1], hx.2.trans (Finset.card_le_card ?_)⟩
  intro i
  simp only [Finset.mem_filter, Finset.mem_univ, true_and]
  intro hi
  exact lt_of_le_of_ne (by linarith) (Ne.symm (hz.2 i))

theorem zeroProb'_le {d : ℕ} {a : Sample m} {b : Fin d → Sample m} (ha : ‖a‖ = 1)
    (hb : ∀ i, ‖b i‖ = 1) :
    zeroProb' a b ≤ (μ m).real (ev 1 a b) + (μ m).real (ev 3 a b) := by
  rw [zeroProb']
  refine add_le_add ?_ le_rfl
  rw [← real_neg]
  refine real_le_of_subset_union (z := zeroEv a b) ?_ (real_zeroEv ha hb)
  intro x hx
  by_cases hz : x ∈ zeroEv a b
  · exact Or.inr hz
  left
  simp only [zeroEv, Set.mem_union, Set.mem_ofPred_eq, Set.mem_iUnion, not_or, not_exists] at hz
  simp only [Set.mem_preimage, Set.mem_ofPred_eq, inner_neg_right] at hx
  obtain ⟨i, hi⟩ := hx.2
  refine ⟨by linarith [hx.1], ?_⟩
  simp only [cnt]
  exact Finset.card_pos.mpr ⟨i, by
    simp only [Finset.mem_filter, Finset.mem_univ, true_and]
    exact lt_of_le_of_ne (by linarith) (Ne.symm (hz.2 i))⟩

/-! ## The bridge: `Ψ` as the sum of the two events -/

theorem psi_eq {d : ℕ} (hd : d ≤ 4) (a : Sample m) (b : Fin d → Sample m) :
    psi (T2.gram (Fin.cons a b : Fin (d + 1) → Sample m)) =
      (μ m).real (ev 1 a b) + (μ m).real (ev 3 a b) := by
  classical
  set u : Fin (d + 1) → Sample m := Fin.cons a b
  set G : ℕ → Set (EuclideanSpace ℝ (Fin (d + 1))) := fun k => {y | y 0 ≤ 0 ∧ k ≤ posCount y}
  have hc : ∀ j : Fin (d + 1), Measurable fun y : EuclideanSpace ℝ (Fin (d + 1)) => y j :=
    fun j => by fun_prop
  have hG : ∀ k, MeasurableSet (G k) := fun k => by
    simp only [G, Set.ofPred_and]
    unfold posCount
    exact (measurableSet_le (hc 0) measurable_const).inter
      (measurableSet_card_ge _ (fun i => measurableSet_lt measurable_const (hc i.succ)) k)
  have hI : ∀ k, Integrable ((G k).indicator (1 : EuclideanSpace ℝ (Fin (d + 1)) → ℝ)) (multivariateGaussian 0 (T2.gram u)) :=
    fun k => (integrable_const (1 : ℝ)).indicator (hG k)
  have hfun : (fun y : EuclideanSpace ℝ (Fin (d + 1)) =>
      if y 0 ≤ 0 then (((posCount y + 1) / 2 : ℕ) : ℝ) else 0) =
      fun y => (G 1).indicator 1 y + (G 3).indicator 1 y := by
    funext y
    have hN : posCount y ≤ 4 := by
      unfold posCount
      exact (Finset.card_filter_le _ _).trans (by simp [hd])
    simp only [G, Set.indicator, Set.mem_ofPred_eq, Pi.one_apply]
    generalize posCount y = N at hN ⊢
    split_ifs <;> first | (exfalso; tauto) | (interval_cases N <;> simp_all <;> norm_num)
  have hpre : ∀ k, T2.innerVec u ⁻¹' G k = ev k a b := by
    intro k
    ext x
    simp [G, ev, cnt, posCount, u, T2.innerVec]
  unfold psi
  rw [hfun, integral_add (hI 1) (hI 3),
    integral_indicator_one (hG 1), integral_indicator_one (hG 3), ← T2.map_innerVec_eq,
    map_measureReal_apply (T2.measurable_innerVec u) (hG 1),
    map_measureReal_apply (T2.measurable_innerVec u) (hG 3), hpre, hpre]

theorem switchAdmissible_gram {d : ℕ} {r : Fin d → ℝ} {a : Sample m} {b : Fin d → Sample m}
    (ha : ‖a‖ = 1) (hb : ∀ i, ‖b i‖ = 1) (hr : ∀ i, r i ≤ ⟪a, b i⟫) :
    SwitchAdmissible r (T2.gram (Fin.cons a b : Fin (d + 1) → Sample m)) := by
  refine ⟨⟨T2.gram_posSemidef _, fun i => ?_⟩, fun i => ?_⟩
  · refine Fin.cases ?_ (fun j => ?_) i
    · simp [T2.gram, ha]
    · simp [T2.gram, hb j]
  · simpa [T2.gram] using hr i

/-! ## Raising the floors -/

theorem exists_raise {a b : Sample m} (ha : ‖a‖ = 1) (hb : ‖b‖ = 1) {r η : ℝ} (hr0 : 0 < r)
    (hr1 : r < 1) (hη0 : 0 ≤ η) (hη1 : η < 1) (h : (1 - η) * r ≤ ⟪a, b⟫) :
    ∃ b' : Sample m, ‖b'‖ = 1 ∧ r ≤ ⟪a, b'⟫ ∧ ‖b - b'‖ ^ 2 ≤ 2 * η := by
  by_cases hc : r ≤ ⟪a, b⟫
  · exact ⟨b, hb, hc, by simp; positivity⟩
  push Not at hc
  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 hba : ⟪b, a⟫ = ⟪a, b⟫ := real_inner_comm a b
  set c := ⟪a, b⟫ with hc_def
  have hc0 : 0 < c := lt_of_lt_of_le (mul_pos (by linarith) hr0) h
  set w := b - c • a with hw_def
  have haw : ⟪a, w⟫ = 0 := by
    simp only [w, inner_sub_right, real_inner_smul_right, haa]; ring
  have hwa : ⟪w, a⟫ = 0 := by rw [real_inner_comm]; exact haw
  have hbw : ⟪b, w⟫ = 1 - c ^ 2 := by
    simp only [w, inner_sub_right, real_inner_smul_right, hbb, hba]; ring
  have hww : ⟪w, w⟫ = 1 - c ^ 2 := by
    simp only [w, inner_sub_left, real_inner_smul_left, hbw]
    rw [haw]; ring
  have hnw : ‖w‖ ^ 2 = 1 - c ^ 2 := by rw [← real_inner_self_eq_norm_sq, hww]
  have hwpos : 0 < ‖w‖ := by
    have h2 : 0 < ‖w‖ ^ 2 := by rw [hnw]; nlinarith
    exact lt_of_le_of_ne (norm_nonneg w) (fun h0 => by rw [← h0] at h2; simp at h2)
  set s := Real.sqrt (1 - r ^ 2)
  have hs : s ^ 2 = 1 - r ^ 2 := Real.sq_sqrt (by nlinarith)
  have hs0 : 0 ≤ s := Real.sqrt_nonneg _
  have hsw : s ≤ ‖w‖ := by
    rw [← Real.sqrt_sq (norm_nonneg w), hnw]
    exact Real.sqrt_le_sqrt (by nlinarith)
  set k := s / ‖w‖
  have hk : k * ‖w‖ = s := div_mul_cancel₀ _ hwpos.ne'
  set b' := r • a + k • w
  have hab' : ⟪a, b'⟫ = r := by
    simp only [b', inner_add_right, real_inner_smul_right, haa, haw]; ring
  have hbb' : ⟪b, b'⟫ = r * c + s * ‖w‖ := by
    simp only [b', inner_add_right, real_inner_smul_right, hba, hbw, ← hnw]
    rw [← hk]; ring
  have hb'b' : ⟪b', b'⟫ = 1 := by
    simp only [b', inner_add_left, inner_add_right, real_inner_smul_left, real_inner_smul_right,
      haa, haw, hwa, hww, ← hnw]
    have : k ^ 2 * ‖w‖ ^ 2 = s ^ 2 := by rw [← hk]; ring
    nlinarith
  have hnb' : ‖b'‖ ^ 2 = 1 := by rw [← real_inner_self_eq_norm_sq, hb'b']
  refine ⟨b', ?_, hab'.ge, ?_⟩
  · exact (pow_eq_one_iff_of_nonneg (norm_nonneg b') two_ne_zero).mp hnb'
  · rw [norm_sub_sq_real, hb, hnb', hbb']
    have h1 : s * s ≤ s * ‖w‖ := mul_le_mul_of_nonneg_left hsw hs0
    have h2 : r - c ≤ η * r := by nlinarith
    nlinarith

/-! ## The sign coupling -/

theorem real_diffEv_le {b b' : Sample m} (hb : ‖b‖ = 1) {t : ℝ} (ht : 0 < t) :
    (μ m).real (diffEv b b') ≤ 3 * t / 2 + ‖b - b'‖ ^ 2 / t ^ 2 := by
  have hsub : diffEv b b' ⊆ {x | ⟪b, x⟫ ∈ Set.Icc (-t) t} ∪ {x | t < |⟪b - b', x⟫|} := by
    intro x hx
    simp only [diffEv, Set.mem_ofPred_eq] at hx
    simp only [Set.mem_union, Set.mem_ofPred_eq, Set.mem_Icc, inner_sub_left]
    by_cases h : |⟪b, x⟫| ≤ t
    · exact Or.inl (abs_le.mp h)
    right
    push Not at h
    rw [lt_abs]
    rcases lt_or_ge 0 ⟪b, x⟫ with hp | hp
    · have hq : ⟪b', x⟫ ≤ 0 := by
        by_contra hq
        exact hx ⟨fun _ => lt_of_not_ge hq, fun _ => hp⟩
      rw [abs_of_pos hp] at h
      left; linarith
    · have hq : 0 < ⟪b', x⟫ := by
        by_contra hq
        exact hx ⟨fun h' => absurd h' (not_lt.mpr hp), fun h' => absurd h' hq⟩
      rw [abs_of_nonpos hp] at h
      right; linarith
  refine (measureReal_mono hsub).trans ((measureReal_union_le _ _).trans (add_le_add ?_ ?_))
  · rw [C3.real_unit_projection hb measurableSet_Icc]
    calc (gaussianReal 0 1).real (Set.Icc (-t) t)
        ≤ (gaussianReal 0 1).real (Set.Ioc (-t - t) t) :=
          measureReal_mono (fun z hz => ⟨by linarith [hz.1], hz.2⟩)
      _ ≤ (t - (-t - t)) / 2 := C3.gaussianReal_Ioc_le (by linarith)
      _ = 3 * t / 2 := by ring
  · by_cases hv : b - b' = 0
    · have : {x : Sample m | t < |⟪b - b', x⟫|} = ∅ := by
        ext x
        simp only [hv, inner_zero_left, abs_zero, Set.mem_ofPred_eq, Set.mem_empty_iff_false,
          iff_false, not_lt]
        exact ht.le
      rw [this, measureReal_empty]
      positivity
    · set v := b - b'
      have hnv : 0 < ‖v‖ := norm_pos_iff.mpr hv
      set u := ‖v‖⁻¹ • v
      have hu : ‖u‖ = 1 := by
        rw [norm_smul, norm_inv, norm_norm, inv_mul_cancel₀ hnv.ne']
      have hvu : ∀ x, ⟪v, x⟫ = ‖v‖ * ⟪u, x⟫ := fun x => by
        rw [real_inner_smul_left, ← mul_assoc, mul_inv_cancel₀ hnv.ne', one_mul]
      have hsub2 : {x : Sample m | t < |⟪v, x⟫|} ⊆ {x | ⟪u, x⟫ ∈ {z : ℝ | t / ‖v‖ ≤ |z|}} := by
        intro x hx
        simp only [Set.mem_ofPred_eq, hvu, abs_mul, abs_of_pos hnv] at hx ⊢
        rw [div_le_iff₀ hnv]
        linarith [mul_comm ‖v‖ |⟪u, x⟫|]
      refine (measureReal_mono hsub2).trans ?_
      rw [C3.real_unit_projection hu (measurableSet_le measurable_const continuous_abs.measurable)]
      refine (C3.gaussianReal_abs_ge_le (by positivity)).trans (le_of_eq ?_)
      field_simp

theorem real_ev_le {d k : ℕ} (a : Sample m) (b b' : Fin d → Sample m) :
    (μ m).real (ev k a b) ≤
      (μ m).real (ev k a b') + ∑ i, (μ m).real (diffEv (b i) (b' i)) := by
  have hsub : ev k a b ⊆ ev k a b' ∪ ⋃ i, diffEv (b i) (b' i) := by
    intro x hx
    by_cases h : x ∈ ⋃ i, diffEv (b i) (b' i)
    · exact Or.inr h
    left
    simp only [Set.mem_iUnion, diffEv, Set.mem_ofPred_eq, not_exists, not_not] at h
    refine ⟨hx.1, ?_⟩
    have : cnt b x = cnt b' x := by
      unfold cnt
      congr 1
      exact Finset.filter_congr fun i _ => h i
    rw [← this]
    exact hx.2
  calc (μ m).real (ev k a b)
      ≤ (μ m).real (ev k a b' ∪ ⋃ i, diffEv (b i) (b' i)) := measureReal_mono hsub
    _ ≤ (μ m).real (ev k a b') + (μ m).real (⋃ i, diffEv (b i) (b' i)) :=
        measureReal_union_le _ _
    _ ≤ _ := by gcongr; exact measureReal_iUnion_fintype_le _

end PsiLower

open PsiLower in
/-- (Q2) `PsiLowerClaim`. -/
theorem psiLower : PsiLowerClaim := by
  intro d hd r hr V hV δ hδ
  set t : ℝ := δ / 24 with ht_def
  have ht : 0 < t := by positivity
  refine ⟨min (1 / 2) (δ * t ^ 2 / 32), lt_min (by norm_num) (by positivity), ?_⟩
  intro m a b ha hb hab
  set η := min (1 / 2) (δ * t ^ 2 / 32)
  have hη0 : 0 ≤ η := le_min (by norm_num) (by positivity)
  have hη1 : η < 1 := (min_le_left _ _).trans_lt (by norm_num)
  have hηδ : η ≤ δ * t ^ 2 / 32 := min_le_right _ _
  choose b' hb'1 hb'r hb'd using
    fun i => exists_raise ha (hb i) (hr i).1 (hr i).2 hη0 hη1 (hab i)
  have hD : ∀ i, (μ m).real (diffEv (b i) (b' i)) ≤ δ / 8 := fun i => by
    refine (real_diffEv_le (hb i) ht).trans ?_
    have h1 : ‖b i - b' i‖ ^ 2 / t ^ 2 ≤ 2 * η / t ^ 2 :=
      div_le_div_of_nonneg_right (hb'd i) (by positivity)
    have h2 : 2 * η / t ^ 2 ≤ δ / 16 := by
      rw [div_le_iff₀ (by positivity)]
      nlinarith
    have h3 : 3 * t / 2 = δ / 16 := by rw [ht_def]; ring
    linarith
  have hsum : ∑ i, (μ m).real (diffEv (b i) (b' i)) ≤ δ / 2 := by
    calc ∑ i, (μ m).real (diffEv (b i) (b' i)) ≤ ∑ _i : Fin d, δ / 8 :=
          Finset.sum_le_sum fun i _ => hD i
      _ = d * (δ / 8) := by simp
      _ ≤ 4 * (δ / 8) := by gcongr; exact_mod_cast hd
      _ = δ / 2 := by ring
  have hpsi := hV _ (switchAdmissible_gram ha hb'1 hb'r)
  rw [psi_eq hd] at hpsi
  have h1 := real_ev_le (k := 1) a b b'
  have h3 := real_ev_le (k := 3) a b b'
  exact ⟨(zeroProb_le ha hb).trans (by linarith), (zeroProb'_le ha hb).trans (by linarith)⟩

end Sw4
end CubicPathwidth

#print axioms CubicPathwidth.Sw4.psiLower