import Mathlib

variable {V : Type*} [DecidableEq V] (G : SimpleGraph V)

namespace SimpleGraph.Walk

-- Helper lemma 1: The prefix of a shortest path is also a shortest path
lemma dist_eq_length_takeUntil {u v x : V} (p : G.Walk u v) (hp : p.IsPath)
    (hp_len : p.length = G.dist u v) (hx : x ∈ p.support) :
    (p.takeUntil x hx).length = G.dist u x := by
  -- Strategy: Prove takeUntil yields a path, and length ≤ dist
  -- Since dist ≤ length of any walk, they are equal
  let p_to_x := p.takeUntil x hx

  -- p_to_x is a path
  have hp_to_x : p_to_x.IsPath := by
    have : (p.takeUntil x hx).append (p.dropUntil x hx) = p := p.take_spec hx
    exact hp.takeUntil _
  -- p_to_x.length ≤ p.length
  have h_le : p_to_x.length ≤ p.length := SimpleGraph.Walk.length_takeUntil_le p hx
  have h_triangle : G.dist u v ≤ G.dist u x + G.dist x v := by
    by_cases h : G.Reachable u x ∧ G.Reachable x v;
    · obtain ⟨q, hq⟩ : ∃ q : G.Walk u x, q.length = G.dist u x := by
        have := h.1;
        exact Reachable.exists_walk_length_eq_dist this
      obtain ⟨r, hr⟩ : ∃ r : G.Walk x v, r.length = G.dist x v := by
        obtain ⟨r, hr⟩ : ∃ r : G.Walk x v, r.length = G.dist x v := by
          have h_reachable : G.Reachable x v := h.right
          exact Reachable.exists_walk_length_eq_dist h_reachable
        use r;
      exact SimpleGraph.dist_le ( q.append r ) |> le_trans <| by simp +decide [ hq, hr ] ;
    · cases hp_to_x ; simp_all +decide [ SimpleGraph.Reachable ];
      contrapose! h;
      exact ⟨ p_to_x, ⟨ p.dropUntil x hx ⟩ ⟩
  have h_dist_xv_le : G.dist x v ≤ (p.dropUntil x hx).length := SimpleGraph.dist_le (p.dropUntil x hx)

  have h_eq_parts : p.length = p_to_x.length + (p.dropUntil x hx).length := by
    have := congr_arg Walk.length (p.take_spec hx)
    -- The length of the appended walk is the sum of the lengths of the two parts.
    rw [← this];
    -- The length of the appended walk is the sum of the lengths of the two parts by definition.
    apply SimpleGraph.Walk.length_append
  have h_upper : p_to_x.length ≤ G.dist u x := by
    rw [hp_len] at h_eq_parts
    have := calc G.dist u v
      _ ≤ G.dist u x + G.dist x v := h_triangle
      _ ≤ G.dist u x + (p.dropUntil x hx).length := Nat.add_le_add_left h_dist_xv_le _
    omega
  refine' le_antisymm h_upper _;
  exact dist_le (p.takeUntil x hx)

-- Helper lemma 2: Length of the suffix of a shortest path
lemma length_dropUntil_eq_dist_sub {u v x : V} (p : G.Walk u v) (hp : p.IsPath)
    (hp_len : p.length = G.dist u v) (hx : x ∈ p.support) :
    (p.dropUntil x hx).length = G.dist u v - G.dist u x := by
  have h1 := congr_arg Walk.length (p.take_spec hx)
  have h4 : (p.takeUntil x hx).length = G.dist u x := by
    rw [ SimpleGraph.Walk.dist_eq_length_takeUntil ];
    · assumption;
    · exact hp_len;
  rw [ ← h4, ← hp_len, ← h1, SimpleGraph.Walk.length_append ] ; aesop