import CubicPathwidth.Sw4.Defs
import CubicPathwidth.Sw4.ZeroVertexNum
import CubicPathwidth.C3.Switch
import CubicPathwidth.C3.GaussianTools
import CubicPathwidth.T3.Constants

/-!
# Q3 (`lem:switch4-zero` per vertex): the level-zero cost plus the transfer `s`

`zeroVertex_of : PsiLowerClaim → SwitchMixedClaim → ZeroVertexClaim`.

For each of the nine `{3,4}`-types `(d; k₃, k₄)` we find a precision `η` and take the
minimum.

* Types `(3; 2, 1)`, `(3; 1, 2)` and `(4; j, 4 - j)`: `PsiLowerClaim` with the bound
  `V` of H6 (parts (a), (b), (c)); then `V + transfer ≤ a_d` (for `(3; 1, 2)` this is
  `c₃ ≤ 0.1401846`, `c3_le`).
* Types `(3; 3, 0)` and `(3; 0, 3)`: all floors are at most `r*`, so the proved
  `C3.switchClaim` bounds both probabilities by `κ(r)` for `r = (1 - η) · floor`;
  continuity of `κ` and `κ(r₃₃) = c₃`, `κ(r₃₄) ≤ c₃ + 3 s` (`kappa_r34_le`) finish.
-/

noncomputable section
open scoped RealInnerProductSpace
open MeasureTheory ProbabilityTheory

namespace CubicPathwidth
namespace Sw4

open T2 T3

/-- The per-type conclusion of `ZeroVertexClaim` at precision `η`. -/
def ZGood (δ η : ℝ) (d k₃ k₄ : ℕ) : Prop :=
  ∀ (m : ℕ) (a : GaussianField.Sample m) (b : Fin d → GaussianField.Sample m),
    ‖a‖ = 1 → (∀ i, ‖b i‖ = 1) →
    (∀ i, (1 - η) * typeFloors d k₃ k₄ i ≤ ⟪a, b i⟫) →
      zeroProb a b + transferTerm d k₃ k₄ 0 sM 0 0 ≤ coeff C3.c3 aM 0 d + δ ∧
      zeroProb' a b + transferTerm d k₃ k₄ 0 sM 0 0 ≤ coeff C3.c3 aM 0 d + δ

theorem typeFloors_mem {d k₃ k₄ : ℕ} (hd : 3 ≤ d) (hk : k₃ + k₄ = d) (i : Fin d) :
    0 < typeFloors d k₃ k₄ i ∧ typeFloors d k₃ k₄ i < 1 := by
  have hi := i.isLt
  unfold typeFloors
  split_ifs with h1 h2
  · exact ⟨rFloor_pos hd le_rfl, rFloor_lt_one hd le_rfl⟩
  · exact ⟨rFloor_pos hd (by norm_num), rFloor_lt_one hd (by norm_num)⟩
  · omega

theorem ZGood.mono {δ η η' : ℝ} {d k₃ k₄ : ℕ} (hd : 3 ≤ d) (hk : k₃ + k₄ = d)
    (h : ZGood δ η d k₃ k₄) (hη' : η' ≤ η) : ZGood δ η' d k₃ k₄ := by
  intro m a b ha hb hab
  refine h m a b ha hb fun i => le_trans ?_ (hab i)
  have := (typeFloors_mem hd hk i).1
  nlinarith

/-! ## Types obtained from `PsiLowerClaim` -/

theorem zgood_of_psi (hlow : PsiLowerClaim) {d k₃ k₄ : ℕ} (hd : 3 ≤ d) (hd4 : d ≤ 4)
    (hk : k₃ + k₄ = d) {V : ℝ}
    (hV : ∀ S : Matrix (Fin (d + 1)) (Fin (d + 1)) ℝ,
      SwitchAdmissible (typeFloors d k₃ k₄) S → psi S ≤ V)
    (hVt : V + transferTerm d k₃ k₄ 0 sM 0 0 ≤ coeff C3.c3 aM 0 d) {δ : ℝ} (hδ : 0 < δ) :
    ∃ η : ℝ, 0 < η ∧ ZGood δ η d k₃ k₄ := by
  obtain ⟨η, hη, h⟩ := hlow d hd4 (typeFloors d k₃ k₄) (typeFloors_mem hd hk) V hV δ hδ
  refine ⟨η, hη, fun m a b ha hb hab => ?_⟩
  have := h m a b ha hb hab
  constructor <;> linarith [this.1, this.2]

theorem zgood_321 (hlow : PsiLowerClaim) (h6 : SwitchMixedClaim) {δ : ℝ} (hδ : 0 < δ) :
    ∃ η : ℝ, 0 < η ∧ ZGood δ η 3 2 1 := by
  refine zgood_of_psi hlow (d := 3) (k₃ := 2) (k₄ := 1) (V := 1529003 / 10000000)
    (by norm_num) (by norm_num) rfl ?_ ?_ hδ
  · exact fun S hS => h6.1 (Equiv.refl _) S hS
  · simp only [transferTerm, coeff, sM]; norm_num

theorem zgood_312 (hlow : PsiLowerClaim) (h6 : SwitchMixedClaim) {δ : ℝ} (hδ : 0 < δ) :
    ∃ η : ℝ, 0 < η ∧ ZGood δ η 3 1 2 := by
  refine zgood_of_psi hlow (d := 3) (k₃ := 1) (k₄ := 2) (V := 165616 / 1000000)
    (by norm_num) (by norm_num) rfl ?_ ?_ hδ
  · exact fun S hS => h6.2.1 (Equiv.refl _) S hS
  · have := c3_le
    simp only [transferTerm, coeff, sM]; norm_num; linarith

theorem zgood_four (hlow : PsiLowerClaim) (h6 : SwitchMixedClaim) {δ : ℝ} (hδ : 0 < δ)
    {j : ℕ} (hj : j ≤ 4) : ∃ η : ℝ, 0 < η ∧ ZGood δ η 4 j (4 - j) := by
  refine zgood_of_psi hlow (d := 4) (k₃ := j) (k₄ := 4 - j) (V := aM - (j : ℝ) * sM)
    (by norm_num) (by norm_num) (by omega) ?_ ?_ hδ
  · exact fun S hS => h6.2.2 j hj (Equiv.refl _) S hS
  · simp only [transferTerm, coeff]; norm_num

/-! ## The two pure types via the proved switch lemma -/

theorem card_eq_three {p : Fin 3 → Prop} [DecidablePred p]
    (h : 3 ≤ (Finset.univ.filter p).card) : ∀ i, p i := by
  intro i
  by_contra hi
  have hsub : Finset.univ.filter p ⊆ Finset.univ.erase i := by
    intro j hj
    simp only [Finset.mem_filter, Finset.mem_univ, true_and] at hj
    simp only [Finset.mem_erase, Finset.mem_univ, and_true]
    rintro rfl
    exact hi hj
  have := Finset.card_le_card hsub
  rw [Finset.card_erase_of_mem (Finset.mem_univ i)] at this
  simp at this
  omega

section Cubic

variable {m : ℕ}

private theorem meas_A (a : GaussianField.Sample m) (b : Fin 3 → GaussianField.Sample m) :
    MeasurableSet {x : GaussianField.Sample m | ⟪a, x⟫ ≤ 0 ∧ ∃ i, 0 ≤ ⟪b i, x⟫} := by
  have h1 : MeasurableSet {x : GaussianField.Sample m | ⟪a, x⟫ ≤ 0} :=
    measurableSet_le (by fun_prop) measurable_const
  have h2 : MeasurableSet {x : GaussianField.Sample m | ∃ i, 0 ≤ ⟪b i, x⟫} := by
    have : {x : GaussianField.Sample m | ∃ i, 0 ≤ ⟪b i, x⟫} =
        ⋃ i, {x | 0 ≤ ⟪b i, x⟫} := by ext x; simp
    rw [this]
    exact MeasurableSet.iUnion fun i => measurableSet_le measurable_const (by fun_prop)
  exact h1.inter h2

private theorem meas_B (a : GaussianField.Sample m) (b : Fin 3 → GaussianField.Sample m) :
    MeasurableSet {x : GaussianField.Sample m | 0 ≤ ⟪a, x⟫ ∧ ∀ i, ⟪b i, x⟫ ≤ 0} := by
  have h1 : MeasurableSet {x : GaussianField.Sample m | 0 ≤ ⟪a, x⟫} :=
    measurableSet_le measurable_const (by fun_prop)
  have h2 : MeasurableSet {x : GaussianField.Sample m | ∀ i, ⟪b i, x⟫ ≤ 0} := by
    have : {x : GaussianField.Sample m | ∀ i, ⟪b i, x⟫ ≤ 0} =
        ⋂ i, {x | ⟪b i, x⟫ ≤ 0} := by ext x; simp
    rw [this]
    exact MeasurableSet.iInter fun i => measurableSet_le (by fun_prop) measurable_const
  exact h1.inter h2

/-- For `d = 3` and `r ≤ r*`, both level-zero probabilities are at most `κ(r)`. -/
theorem zeroProb_le_kappa {r : ℝ} (hr0 : 0 < r) (hr : r ≤ C3.rStar)
    (a : GaussianField.Sample m) (b : Fin 3 → GaussianField.Sample m)
    (ha : ‖a‖ = 1) (hb : ∀ i, ‖b i‖ = 1) (hab : ∀ i, r ≤ ⟪a, b i⟫) :
    zeroProb a b ≤ C3.kappa r ∧ zeroProb' a b ≤ C3.kappa r := by
  classical
  have hs := C3.switchClaim m r hr0 hr a 3 le_rfl b ha hb hab
  simp only [ite_true] at hs
  set A : Set (GaussianField.Sample m) := {x | ⟪a, x⟫ ≤ 0 ∧ ∃ i, 0 ≤ ⟪b i, x⟫} with hAdef
  set B : Set (GaussianField.Sample m) := {x | 0 ≤ ⟪a, x⟫ ∧ ∀ i, ⟪b i, x⟫ ≤ 0} with hBdef
  have hA := meas_A a b
  have hB := meas_B a b
  have hnA := C3.real_neg_preimage hA
  have hnB := C3.real_neg_preimage hB
  constructor
  · unfold zeroProb
    have h1 : (μ m).real {x | ⟪a, x⟫ ≤ 0 ∧ ∃ i, 0 < ⟪b i, x⟫} ≤ (μ m).real A := by
      refine measureReal_mono ?_
      rintro x ⟨hx1, i, hi⟩
      exact ⟨hx1, i, hi.le⟩
    have h2 : (μ m).real {x | 0 < ⟪a, x⟫ ∧
        3 ≤ (Finset.univ.filter fun i => ⟪b i, x⟫ ≤ 0).card} ≤ (μ m).real B := by
      refine measureReal_mono ?_
      rintro x ⟨hx1, hx2⟩
      exact ⟨hx1.le, card_eq_three hx2⟩
    simp only [μ] at h1 h2 ⊢
    linarith
  · unfold zeroProb'
    have h1 : (μ m).real {x | 0 < ⟪a, x⟫ ∧ ∃ i, ⟪b i, x⟫ ≤ 0} ≤
        (μ m).real (Neg.neg ⁻¹' A) := by
      refine measureReal_mono ?_
      rintro x ⟨hx1, i, hi⟩
      simp only [Set.mem_preimage, hAdef, Set.mem_ofPred_eq, inner_neg_right]
      exact ⟨by linarith, i, by linarith⟩
    have h2 : (μ m).real {x | ⟪a, x⟫ ≤ 0 ∧
        3 ≤ (Finset.univ.filter fun i => 0 < ⟪b i, x⟫).card} ≤
        (μ m).real (Neg.neg ⁻¹' B) := by
      refine measureReal_mono ?_
      rintro x ⟨hx1, hx2⟩
      have := card_eq_three hx2
      simp only [Set.mem_preimage, hBdef, Set.mem_ofPred_eq, inner_neg_right]
      exact ⟨by linarith, fun i => by linarith [this i]⟩
    simp only [μ] at h1 h2 hnA hnB ⊢
    linarith

/-- Continuity of `κ((1 - η) F)` gives a precision `η` for a pure cubic type. -/
theorem cubic_case {F T L δ : ℝ} (hF0 : 0 < F) (hFs : F ≤ C3.rStar) (hδ : 0 < δ)
    (hT : C3.kappa F + T ≤ L) :
    ∃ η : ℝ, 0 < η ∧ η < 1 ∧
      ∀ (m : ℕ) (a : GaussianField.Sample m) (b : Fin 3 → GaussianField.Sample m),
        ‖a‖ = 1 → (∀ i, ‖b i‖ = 1) → (∀ i, (1 - η) * F ≤ ⟪a, b i⟫) →
          zeroProb a b + T ≤ L + δ ∧ zeroProb' a b + T ≤ L + δ := by
  have hc : Continuous fun t : ℝ => C3.kappa ((1 - t) * F) + T := by
    unfold C3.kappa C3.pairCorr
    fun_prop
  obtain ⟨ε, hε, hεδ⟩ := Metric.continuousAt_iff.mp hc.continuousAt δ hδ
  set η : ℝ := min (ε / 2) (1 / 2) with hη
  have hη0 : 0 < η := lt_min (by linarith) (by norm_num)
  have hη1 : η < 1 := lt_of_le_of_lt (min_le_right _ _) (by norm_num)
  have hηε : η < ε := lt_of_le_of_lt (min_le_left _ _) (by linarith)
  have hd := hεδ (x := η) (by rw [Real.dist_eq, sub_zero, abs_of_pos hη0]; exact hηε)
  rw [Real.dist_eq] at hd
  have hd' := (abs_lt.mp hd).2
  simp only [sub_zero, one_mul] at hd'
  refine ⟨η, hη0, hη1, fun m a b ha hb hab => ?_⟩
  have hr0 : 0 < (1 - η) * F := mul_pos (by linarith) hF0
  have hr : (1 - η) * F ≤ C3.rStar := by nlinarith
  have := zeroProb_le_kappa hr0 hr a b ha hb hab
  constructor <;> linarith [this.1, this.2]

end Cubic

theorem typeFloors_330 (i : Fin 3) : typeFloors 3 3 0 i = C3.rStar := by
  have hi := i.isLt
  unfold typeFloors
  simp only [show (i : ℕ) < 3 by omega, ↓reduceIte]
  rw [rFloor_33]
  rfl

theorem typeFloors_303 (i : Fin 3) : typeFloors 3 0 3 i = T2.r34 := by
  have hi := i.isLt
  unfold typeFloors
  simp only [show ¬ (i : ℕ) < 0 by omega, show (i : ℕ) < 0 + 3 by omega, ↓reduceIte]
  rw [rFloor_34]

theorem zgood_330 {δ : ℝ} (hδ : 0 < δ) : ∃ η : ℝ, 0 < η ∧ ZGood δ η 3 3 0 := by
  obtain ⟨η, hη, -, h⟩ := cubic_case (F := C3.rStar) (T := 0) (L := C3.c3) (δ := δ)
    T2.r33_pos (le_refl _) hδ (by rw [C3.kappa_rStar]; norm_num)
  refine ⟨η, hη, fun m a b ha hb hab => ?_⟩
  have := h m a b ha hb (fun i => by simpa [typeFloors_330] using hab i)
  simp only [transferTerm, coeff]
  norm_num
  constructor <;> linarith [this.1, this.2]

theorem zgood_303 {δ : ℝ} (hδ : 0 < δ) : ∃ η : ℝ, 0 < η ∧ ZGood δ η 3 0 3 := by
  obtain ⟨η, hη, -, h⟩ := cubic_case (F := T2.r34) (T := -(3 : ℝ) * sM) (L := C3.c3) (δ := δ)
    r34_pos' r34_le_rStar hδ (by have := kappa_r34_le; linarith)
  refine ⟨η, hη, fun m a b ha hb hab => ?_⟩
  have := h m a b ha hb (fun i => by simpa [typeFloors_303] using hab i)
  simp only [transferTerm, coeff]
  norm_num
  constructor <;> linarith [this.1, this.2]

/-! ## Assembly -/

theorem exists_good (hlow : PsiLowerClaim) (h6 : SwitchMixedClaim) {δ : ℝ} (hδ : 0 < δ)
    {d k₃ k₄ : ℕ} (hd : d = 3 ∨ d = 4) (hk : k₃ + k₄ = d) :
    ∃ η : ℝ, 0 < η ∧ ZGood δ η d k₃ k₄ := by
  rcases hd with rfl | rfl
  · have : k₄ ≤ 3 := by omega
    interval_cases k₄
    · obtain rfl : k₃ = 3 := by omega
      exact zgood_330 hδ
    · obtain rfl : k₃ = 2 := by omega
      exact zgood_321 hlow h6 hδ
    · obtain rfl : k₃ = 1 := by omega
      exact zgood_312 hlow h6 hδ
    · obtain rfl : k₃ = 0 := by omega
      exact zgood_303 hδ
  · obtain rfl : k₄ = 4 - k₃ := by omega
    exact zgood_four hlow h6 hδ (by omega)

theorem exists_min {ι : Type*} (S : Finset ι) (Q : ι → ℝ → Prop)
    (h : ∀ t ∈ S, ∃ η : ℝ, 0 < η ∧ ∀ η', 0 < η' → η' ≤ η → Q t η') :
    ∃ η : ℝ, 0 < η ∧ ∀ t ∈ S, ∀ η', 0 < η' → η' ≤ η → Q t η' := by
  classical
  induction S using Finset.induction_on with
  | empty => exact ⟨1, one_pos, by simp⟩
  | insert t S ht ih =>
    obtain ⟨η₁, h₁, H₁⟩ := ih fun s hs => h s (Finset.mem_insert_of_mem hs)
    obtain ⟨η₂, h₂, H₂⟩ := h t (Finset.mem_insert_self _ _)
    refine ⟨min η₁ η₂, lt_min h₁ h₂, ?_⟩
    intro s hs η' h0 hle
    rcases Finset.mem_insert.mp hs with rfl | hs
    · exact H₂ _ h0 (hle.trans (min_le_right _ _))
    · exact H₁ s hs η' h0 (hle.trans (min_le_left _ _))

/-- (Q3, `lem:switch4-zero` per vertex). -/
theorem zeroVertex_of (hlow : PsiLowerClaim) (h6 : SwitchMixedClaim) : ZeroVertexClaim := by
  intro δ hδ
  let S : Finset (ℕ × ℕ × ℕ) :=
    (Finset.range 5 ×ˢ Finset.range 5 ×ˢ Finset.range 5).filter
      fun t => (t.1 = 3 ∨ t.1 = 4) ∧ t.2.1 + t.2.2 = t.1
  obtain ⟨η, hη, H⟩ := exists_min S (fun t η => ZGood δ η t.1 t.2.1 t.2.2) (by
    rintro ⟨d, k₃, k₄⟩ ht
    simp only [S, Finset.mem_filter] at ht
    obtain ⟨η, hη, hg⟩ := exists_good hlow h6 hδ ht.2.1 ht.2.2
    refine ⟨η, hη, fun η' h0 hle => ?_⟩
    exact hg.mono (by rcases ht.2.1 with h | h <;> omega) ht.2.2 hle)
  refine ⟨η, hη, fun m d k₃ k₄ a b hd hk ha hb hab => ?_⟩
  have hmem : (d, k₃, k₄) ∈ S := by
    simp only [S, Finset.mem_filter, Finset.mem_product, Finset.mem_range]
    omega
  exact H _ hmem η hη le_rfl m a b ha hb hab

end Sw4
end CubicPathwidth

#print axioms CubicPathwidth.Sw4.zeroVertex_of