import CubicPathwidth.T3.GoodLabelsCount5A
import CubicPathwidth.T2.GoodLabelsCountB
import CubicPathwidth.T2.GoodLabelsCount

/-!
# Floors, neighbor enumeration and per-vertex star bound for maximum degree 5
-/

noncomputable section
open scoped RealInnerProductSpace
open MeasureTheory ProbabilityTheory Classical

namespace CubicPathwidth
namespace T3
namespace GLC5

open T2

/-! ## The floors lie in `(0, 1)` -/

theorem rFloor_pos (i j : ℕ) : 0 < rFloor i j := by
  unfold rFloor; exact one_div_pos.mpr (Real.cosh_pos _)

theorem rFloor_lt_one {i j : ℕ} (hi : 3 ≤ i) (hj : 3 ≤ j) : rFloor i j < 1 := by
  unfold rFloor
  have hi' : (3 : ℝ) ≤ i := by exact_mod_cast hi
  have hj' : (3 : ℝ) ≤ j := by exact_mod_cast hj
  have hprod : (1 : ℝ) < ((i : ℝ) - 1) * ((j : ℝ) - 1) := by nlinarith
  have h6 : 0 < Real.log (((i : ℝ) - 1) * ((j : ℝ) - 1)) := Real.log_pos hprod
  have hx : Real.log (((i : ℝ) - 1) * ((j : ℝ) - 1)) / 4 ≠ 0 := by positivity
  have h1 := Real.add_one_lt_exp hx
  have h2 := Real.add_one_lt_exp (neg_ne_zero.mpr hx)
  have h : 1 < Real.cosh (Real.log (((i : ℝ) - 1) * ((j : ℝ) - 1)) / 4) := by
    rw [Real.cosh_eq]
    linarith
  rw [div_lt_one (by linarith)]
  exact h

theorem typeFloors_pos (d k₃ k₄ : ℕ) (i : Fin d) : 0 < typeFloors d k₃ k₄ i := by
  unfold typeFloors; split_ifs <;> exact rFloor_pos _ _

theorem typeFloors_lt_one {d : ℕ} (hd : 3 ≤ d) (k₃ k₄ : ℕ) (i : Fin d) :
    typeFloors d k₃ k₄ i < 1 := by
  unfold typeFloors; split_ifs <;> exact rFloor_lt_one hd (by norm_num)

/-! ## Choice of `η` from the continuity claim -/

theorem eta_exists5 (hC : StarBoundContClaim) {δ : ℝ} (hδ : 0 < δ) :
    ∃ η : ℝ, 0 < η ∧ η < 1 ∧
      ∀ d k₃ k₄ : ℕ, (d = 3 ∨ d = 4 ∨ d = 5) → k₃ ≤ 5 → k₄ ≤ 5 → ∀ r' : Fin d → ℝ,
        (∀ i, (1 - η) * typeFloors d k₃ k₄ i ≤ r' i ∧ r' i ≤ typeFloors d k₃ k₄ i) →
        ∀ τ ∈ Set.Icc (-3 : ℝ) 3,
          starBound r' τ ≤ starBound (typeFloors d k₃ k₄) τ + δ := by
  have h : ∀ p : Fin 6 × Fin 6 × Fin 6, ∃ η : ℝ, 0 < η ∧
      (((p.1 : ℕ) = 3 ∨ (p.1 : ℕ) = 4 ∨ (p.1 : ℕ) = 5) →
      ∀ r' : Fin (p.1 : ℕ) → ℝ,
        (∀ i, (1 - η) * typeFloors p.1 p.2.1 p.2.2 i ≤ r' i ∧
          r' i ≤ typeFloors p.1 p.2.1 p.2.2 i) →
        ∀ τ ∈ Set.Icc (-3 : ℝ) 3,
          starBound r' τ ≤ starBound (typeFloors p.1 p.2.1 p.2.2) τ + δ) := by
    intro p
    by_cases hd : ((p.1 : ℕ) = 3 ∨ (p.1 : ℕ) = 4 ∨ (p.1 : ℕ) = 5)
    · have h3 : 3 ≤ (p.1 : ℕ) := by omega
      obtain ⟨η, hη, hη'⟩ := hC (p.1 : ℕ) (typeFloors p.1 p.2.1 p.2.2)
        (fun i => ⟨typeFloors_pos _ _ _ i, typeFloors_lt_one h3 _ _ i⟩) δ hδ
      exact ⟨η, hη, fun _ => hη'⟩
    · exact ⟨1, one_pos, fun h' => absurd h' hd⟩
  choose e hep he using h
  obtain ⟨p0, -, hp0⟩ := Finset.exists_min_image Finset.univ e
    ⟨(0, 0, 0), Finset.mem_univ _⟩
  refine ⟨min (1 / 2) (e p0), lt_min (by norm_num) (hep p0), ?_, ?_⟩
  · exact lt_of_le_of_lt (min_le_left _ _) (by norm_num)
  · intro d k₃ k₄ hd hk3 hk4 r' hr τ hτ
    have hd6 : d < 6 := by omega
    have := he (⟨d, hd6⟩, ⟨k₃, by omega⟩, ⟨k₄, by omega⟩) hd r' ?_ τ hτ
    · exact this
    · intro i
      refine ⟨le_trans ?_ (hr i).1, (hr i).2⟩
      have hη : min (1 / 2) (e p0) ≤ e (⟨d, hd6⟩, ⟨k₃, by omega⟩, ⟨k₄, by omega⟩) :=
        (min_le_right _ _).trans (hp0 _ (Finset.mem_univ _))
      have hp := typeFloors_pos d k₃ k₄ i
      show (1 - e (⟨d, hd6⟩, ⟨k₃, by omega⟩, ⟨k₄, by omega⟩)) * typeFloors d k₃ k₄ i ≤
        (1 - min (1 / 2) (e p0)) * typeFloors d k₃ k₄ i
      nlinarith

/-! ## Enumeration of a finite set, grouped by a key in `{3,4,5}` -/

theorem exists_enum3 {α : Type*} (N : Finset α) (key : α → ℕ) (d : ℕ) (hd : N.card = d)
    (hk : ∀ a ∈ N, key a = 3 ∨ key a = 4 ∨ key a = 5) :
    ∃ f : Fin d → α, (∀ i, f i ∈ N) ∧ (∀ z ∈ N, ∃ i, f i = z) ∧
      ∀ i : Fin d,
        ((i : ℕ) < (N.filter fun a => key a = 3).card → key (f i) = 3) ∧
        ((N.filter fun a => key a = 3).card ≤ (i : ℕ) →
          (i : ℕ) < (N.filter fun a => key a = 3).card + (N.filter fun a => key a = 4).card →
          key (f i) = 4) ∧
        ((N.filter fun a => key a = 3).card + (N.filter fun a => key a = 4).card ≤ (i : ℕ) →
          key (f i) = 5) := by
  set A3 := N.filter fun a => key a = 3 with hA3
  set A4 := N.filter fun a => key a = 4 with hA4
  set A5 := N.filter fun a => key a = 5 with hA5
  set l : List α := A3.toList ++ (A4.toList ++ A5.toList) with hl
  have hcard : A3.card + A4.card + A5.card = d := by
    have hN : N.card = ∑ b ∈ ({3, 4, 5} : Finset ℕ), (N.filter fun a => key a = b).card :=
      Finset.card_eq_sum_card_fiberwise (f := key) (t := ({3, 4, 5} : Finset ℕ))
        (by intro x hx; simpa using hk x hx)
    rw [← hd, hN]
    simp [Finset.sum_insert, hA3, hA4, hA5, add_assoc]
  have hlen : l.length = d := by
    rw [hl, List.length_append, List.length_append, Finset.length_toList, Finset.length_toList,
      Finset.length_toList]
    omega
  have hmem : ∀ a, a ∈ l ↔ a ∈ N := by
    intro a
    rw [hl, List.mem_append, List.mem_append, Finset.mem_toList, Finset.mem_toList,
      Finset.mem_toList, hA3, hA4, hA5, Finset.mem_filter, Finset.mem_filter, Finset.mem_filter]
    constructor
    · rintro (h | h | h) <;> exact h.1
    · intro h
      rcases hk a h with h' | h' | h'
      · exact Or.inl ⟨h, h'⟩
      · exact Or.inr (Or.inl ⟨h, h'⟩)
      · exact Or.inr (Or.inr ⟨h, h'⟩)
  refine ⟨fun i => l[(i : ℕ)]'(by omega), ?_, ?_, ?_⟩
  · intro i
    exact (hmem _).mp (List.getElem_mem _)
  · intro z hz
    obtain ⟨j, hj, hjz⟩ := List.getElem_of_mem ((hmem z).mpr hz)
    exact ⟨⟨j, by omega⟩, hjz⟩
  · intro i
    refine ⟨fun hi => ?_, fun hi1 hi2 => ?_, fun hi => ?_⟩
    · have hlt : (i : ℕ) < A3.toList.length := by rwa [Finset.length_toList]
      have : l[(i : ℕ)]'(by omega) = A3.toList[(i : ℕ)]'hlt :=
        List.getElem_append_left hlt
      simp only [this]
      have hm : A3.toList[(i : ℕ)]'hlt ∈ A3 := Finset.mem_toList.mp (List.getElem_mem _)
      exact (Finset.mem_filter.mp hm).2
    · have hge : A3.toList.length ≤ (i : ℕ) := by rw [Finset.length_toList]; omega
      have hlt : (i : ℕ) - A3.toList.length < A4.toList.length := by
        rw [Finset.length_toList, Finset.length_toList]; omega
      have h1 : l[(i : ℕ)]'(by omega) =
          (A4.toList ++ A5.toList)[(i : ℕ) - A3.toList.length]'(by
            rw [List.length_append]; omega) :=
        List.getElem_append_right hge
      have h2 : (A4.toList ++ A5.toList)[(i : ℕ) - A3.toList.length]'(by
            rw [List.length_append]; omega) =
          A4.toList[(i : ℕ) - A3.toList.length]'hlt := List.getElem_append_left hlt
      simp only [h1, h2]
      have hm : A4.toList[(i : ℕ) - A3.toList.length]'hlt ∈ A4 :=
        Finset.mem_toList.mp (List.getElem_mem _)
      exact (Finset.mem_filter.mp hm).2
    · have hge : A3.toList.length ≤ (i : ℕ) := by rw [Finset.length_toList]; omega
      have hge2 : A4.toList.length ≤ (i : ℕ) - A3.toList.length := by
        rw [Finset.length_toList, Finset.length_toList]; omega
      have hlt : (i : ℕ) - A3.toList.length - A4.toList.length < A5.toList.length := by
        rw [Finset.length_toList, Finset.length_toList, Finset.length_toList]; omega
      have h1 : l[(i : ℕ)]'(by omega) =
          (A4.toList ++ A5.toList)[(i : ℕ) - A3.toList.length]'(by
            rw [List.length_append]; omega) :=
        List.getElem_append_right hge
      have h2 : (A4.toList ++ A5.toList)[(i : ℕ) - A3.toList.length]'(by
            rw [List.length_append]; omega) =
          A5.toList[(i : ℕ) - A3.toList.length - A4.toList.length]'hlt :=
        List.getElem_append_right hge2
      simp only [h1, h2]
      have hm : A5.toList[(i : ℕ) - A3.toList.length - A4.toList.length]'hlt ∈ A5 :=
        Finset.mem_toList.mp (List.getElem_mem _)
      exact (Finset.mem_filter.mp hm).2

/-! ## Per-vertex bounds -/

variable {n : ℕ} {G : SimpleGraph (Fin n)} {K : ℕ} {vec : Fin n → GaussianField.Sample n}

theorem deg_cases5 (hmin : MinDegreeAtLeast 3 G) (hmax : MaxDegreeAtMost 5 G) (u : Fin n) :
    G.degree u = 3 ∨ G.degree u = 4 ∨ G.degree u = 5 := by
  have h1 : 3 ≤ G.degree u := hmin u
  have h2 : G.degree u ≤ 5 := hmax u
  omega

variable {η : ℝ}

theorem prob_outer_le_typ (hS : HetStarClaim) (hs : GLC.Sch G K vec)
    (hmin : MinDegreeAtLeast 3 G) (hmax : MaxDegreeAtMost 5 G) (hη0 : 0 < η) (hη1 : η < 1)
    (hedge : ∀ u v, G.Adj u v →
      (1 - η) * rFloor (G.degree u) (G.degree v) ≤ ⟪vec u, vec v⟫)
    (v : Fin n) (τ : ℝ) :
    (stdGaussian (GaussianField.Sample n)).real (GLC.outerEv G vec τ v) ≤
      starBound (fun i : Fin (G.degree v) =>
        (1 - η) * typeFloors (G.degree v) (nbrCount G v 3) (nbrCount G v 4) i) τ := by
  have hcard : (G.neighborFinset v).card = G.degree v := G.card_neighborFinset_eq_degree v
  have hk : ∀ a ∈ G.neighborFinset v, G.degree a = 3 ∨ G.degree a = 4 ∨ G.degree a = 5 :=
    fun a _ => deg_cases5 hmin hmax a
  obtain ⟨f, hfm, hfs, hfp⟩ := exists_enum3 (G.neighborFinset v) (fun a => G.degree a) _ hcard hk
  have hv3 : 3 ≤ G.degree v := hmin v
  refine GLC.prob_outer_le_aux hS hs v (G.degree v) _ ?_ f
    (fun z hz => hfs z ((G.mem_neighborFinset v z).mpr hz)) ?_ τ
  · intro i
    have := typeFloors_pos (G.degree v) (nbrCount G v 3) (nbrCount G v 4) i
    have := typeFloors_lt_one hv3 (nbrCount G v 3) (nbrCount G v 4) i
    constructor <;> nlinarith
  · intro i
    have hadj : G.Adj v (f i) := (G.mem_neighborFinset v _).mp (hfm i)
    have h := hedge v (f i) hadj
    refine le_trans (le_of_eq ?_) h
    congr 1
    obtain ⟨p3, p4, p5⟩ := hfp i
    have e3 : nbrCount G v 3 = ((G.neighborFinset v).filter fun a => G.degree a = 3).card :=
      GLC.nbrCount_eq v 3
    have e4 : nbrCount G v 4 = ((G.neighborFinset v).filter fun a => G.degree a = 4).card :=
      GLC.nbrCount_eq v 4
    rw [typeFloors, e3, e4]
    by_cases h1 : (i : ℕ) < ((G.neighborFinset v).filter fun a => G.degree a = 3).card
    · simp only [h1, ↓reduceIte, p3 h1]
    · simp only [h1, ↓reduceIte]
      by_cases h2 : (i : ℕ) <
          ((G.neighborFinset v).filter fun a => G.degree a = 3).card +
            ((G.neighborFinset v).filter fun a => G.degree a = 4).card
      · simp only [h2, ↓reduceIte, p4 (not_lt.mp h1) h2]
      · simp only [h2, ↓reduceIte, p5 (not_lt.mp h2)]

/-- The star bound of the type of `v`, at level `τ`. -/
def tb5 (G : SimpleGraph (Fin n)) (v : Fin n) (τ : ℝ) : ℝ :=
  starBound (typeFloors (G.degree v) (nbrCount G v 3) (nbrCount G v 4)) τ

theorem prob_outer_le_tb5 (hS : HetStarClaim) (hs : GLC.Sch G K vec)
    (hmin : MinDegreeAtLeast 3 G) (hmax : MaxDegreeAtMost 5 G) {δ : ℝ} (hη0 : 0 < η)
    (hη1 : η < 1)
    (he : ∀ d k₃ k₄ : ℕ, (d = 3 ∨ d = 4 ∨ d = 5) → k₃ ≤ 5 → k₄ ≤ 5 → ∀ r' : Fin d → ℝ,
        (∀ i, (1 - η) * typeFloors d k₃ k₄ i ≤ r' i ∧ r' i ≤ typeFloors d k₃ k₄ i) →
        ∀ τ ∈ Set.Icc (-3 : ℝ) 3,
          starBound r' τ ≤ starBound (typeFloors d k₃ k₄) τ + δ)
    (hedge : ∀ u v, G.Adj u v →
      (1 - η) * rFloor (G.degree u) (G.degree v) ≤ ⟪vec u, vec v⟫)
    (v : Fin n) {τ : ℝ} (hτ : τ ∈ Set.Icc (-3 : ℝ) 3) :
    (stdGaussian (GaussianField.Sample n)).real (GLC.outerEv G vec τ v) ≤ tb5 G v τ + δ := by
  have h1 := prob_outer_le_typ hS hs hmin hmax hη0 hη1 hedge v τ
  have hk3 : nbrCount G v 3 ≤ 5 := (GLC.nbrCount_le_degree v 3).trans (hmax v)
  have hk4 : nbrCount G v 4 ≤ 5 := (GLC.nbrCount_le_degree v 4).trans (hmax v)
  have h2 := he (G.degree v) (nbrCount G v 3) (nbrCount G v 4) (deg_cases5 hmin hmax v) hk3 hk4
    (fun i => (1 - η) * typeFloors (G.degree v) (nbrCount G v 3) (nbrCount G v 4) i)
    (fun i => ⟨le_refl _, by
      have := typeFloors_pos (G.degree v) (nbrCount G v 3) (nbrCount G v 4) i
      nlinarith⟩) τ hτ
  exact h1.trans h2

end GLC5
end T3
end CubicPathwidth