import Mathlib

/-!
# Rooting a tree: the parent function

Helpers shared by `MaxLeafFormal.CDS` and `MaxLeafFormal.Antinorm3` (moved here from `CDS.lean`
so that both files can use them; names unchanged).
-/

namespace MaxLeafFormal.CDS

open Finset SimpleGraph

variable {V : Type*} [Fintype V] [DecidableEq V]

noncomputable def comp (G : SimpleGraph V) (X : Finset V) : ℕ :=
  Nat.card (G.induce (X : Set V)).ConnectedComponent

section Parent

variable {F : SimpleGraph V} (hF : F.IsTree) (r : V)
include hF

theorem exists_parent {v : V} (hv : v ≠ r) :
    ∃ w, F.Adj v w ∧ F.dist r w + 1 = F.dist r v := by
  obtain ⟨p, hp⟩ := hF.connected.exists_walk_length_eq_dist v r
  cases p with
  | nil => exact absurd rfl hv
  | @cons _ w _ h q =>
    refine ⟨w, h, ?_⟩
    have hq : F.dist w r ≤ q.length := dist_le q
    simp only [Walk.length_cons] at hp
    rw [dist_comm] at hp
    rcases hF.dist_eq_dist_add_one_of_adj r h with h1 | h1
    · omega
    · rw [dist_comm] at hq; omega

theorem parent_unique {v w₁ w₂ : V} (h₁ : F.Adj v w₁) (h₂ : F.Adj v w₂)
    (hd₁ : F.dist r w₁ + 1 = F.dist r v) (hd₂ : F.dist r w₂ + 1 = F.dist r v) : w₁ = w₂ := by
  obtain ⟨q₁, hq₁⟩ := hF.connected.exists_walk_length_eq_dist r w₁
  obtain ⟨q₂, hq₂⟩ := hF.connected.exists_walk_length_eq_dist r w₂
  have hp₁ : (q₁.concat h₁.symm).IsPath :=
    Walk.isPath_of_length_eq_dist _ (by rw [Walk.length_concat]; omega)
  have hp₂ : (q₂.concat h₂.symm).IsPath :=
    Walk.isPath_of_length_eq_dist _ (by rw [Walk.length_concat]; omega)
  have := hF.isAcyclic.path_unique ⟨_, hp₁⟩ ⟨_, hp₂⟩
  have h := congrArg (fun P : F.Path r v => P.1.penultimate) this
  simpa [Walk.penultimate_concat] using h

open Classical in
/-- The parent of `v` towards the root `r` (and `r` itself at the root). -/
noncomputable def par (v : V) : V :=
  if hv : v = r then r else Classical.choose (exists_parent hF r hv)

theorem par_spec {v : V} (hv : v ≠ r) :
    F.Adj v (par hF r v) ∧ F.dist r (par hF r v) + 1 = F.dist r v := by
  unfold par
  rw [dif_neg hv]
  exact Classical.choose_spec (exists_parent hF r hv)

theorem adj_iff_par {v w : V} :
    F.Adj v w ↔ (v ≠ r ∧ par hF r v = w) ∨ (w ≠ r ∧ par hF r w = v) := by
  constructor
  · intro h
    rcases hF.dist_eq_dist_add_one_of_adj r h with h1 | h1
    · have hv : v ≠ r := by rintro rfl; simp at h1
      exact Or.inl ⟨hv, parent_unique hF r (par_spec hF r hv).1 h (par_spec hF r hv).2 h1.symm⟩
    · have hw : w ≠ r := by rintro rfl; simp at h1
      exact Or.inr ⟨hw, parent_unique hF r (par_spec hF r hw).1 h.symm (par_spec hF r hw).2
        h1.symm⟩
  · rintro (⟨hv, rfl⟩ | ⟨hw, rfl⟩)
    · exact (par_spec hF r hv).1
    · exact (par_spec hF r hw).1.symm

theorem degree_eq_children [DecidableRel F.Adj] (x : V) :
    F.degree x = (univ.filter (fun w => w ≠ r ∧ par hF r w = x)).card
      + (if x = r then 0 else 1) := by
  rw [← card_neighborFinset_eq_degree]
  have hsplit : F.neighborFinset x = (univ.filter (fun w => w ≠ r ∧ par hF r w = x)) ∪
      (if x = r then ∅ else {par hF r x}) := by
    ext w
    rw [mem_neighborFinset, adj_iff_par hF r, mem_union, mem_filter]
    split_ifs with hx
    · simp [hx]
    · simp only [mem_univ, true_and, mem_singleton]
      constructor
      · rintro (⟨_, h⟩ | h)
        · exact Or.inr h.symm
        · exact Or.inl h
      · rintro (h | h)
        · exact Or.inr h
        · exact Or.inl ⟨hx, h.symm⟩
  rw [hsplit, card_union_of_disjoint]
  · split_ifs <;> simp
  · split_ifs with hx
    · simp
    · rw [disjoint_singleton_right, mem_filter]
      rintro ⟨-, hpr, hpp⟩
      have h1 := (par_spec hF r hx).2
      have h2 := (par_spec hF r hpr).2
      rw [hpp] at h2
      omega

/-- Every vertex of `X` reaches, inside `F[X]`, a "top": a vertex of `X` whose parent is
outside `X`, or the root. -/
theorem reach_top (X : Finset V) :
    ∀ d, ∀ x (hx : x ∈ X), F.dist r x = d →
      ∃ t ∈ X.filter (fun t => ¬ (t ≠ r ∧ par hF r t ∈ X)), ∃ ht : t ∈ X,
        (F.induce (X : Set V)).Reachable ⟨x, hx⟩ ⟨t, ht⟩ := by
  intro d
  induction d using Nat.strong_induction_on with
  | _ d ih =>
    intro x hx hd
    by_cases htop : x ≠ r ∧ par hF r x ∈ X
    · obtain ⟨t, htT, ht, hreach⟩ := ih _ (by have := (par_spec hF r htop.1).2; omega)
        (par hF r x) htop.2 rfl
      refine ⟨t, htT, ht, Reachable.trans ?_ hreach⟩
      exact Adj.reachable (by simpa using (par_spec hF r htop.1).1)
    · exact ⟨x, mem_filter.2 ⟨hx, htop⟩, hx, Reachable.refl _⟩

theorem comp_le_tops (X : Finset V) :
    comp F X ≤ (X.filter (fun t => ¬ (t ≠ r ∧ par hF r t ∈ X))).card := by
  set T := X.filter (fun t => ¬ (t ≠ r ∧ par hF r t ∈ X))
  unfold comp
  rw [← Fintype.card_coe T, ← Nat.card_eq_fintype_card]
  apply Nat.card_le_card_of_surjective
    (fun t : T => (F.induce (X : Set V)).connectedComponentMk
      ⟨t.1, (mem_filter.1 t.2).1⟩)
  intro c
  induction c using ConnectedComponent.ind with
  | h x =>
    obtain ⟨t, htT, ht, hreach⟩ := reach_top hF r X _ x.1 x.2 rfl
    exact ⟨⟨t, htT⟩, (ConnectedComponent.eq.2 hreach).symm⟩

end Parent

end MaxLeafFormal.CDS