-- Clean/Utils/SourceSinkPath.lean

import Mathlib.Data.Finset.Basic

import Mathlib.Data.Fintype.Basic

import Mathlib.Data.Fintype.Prod

import Mathlib.Data.List.Basic

import Mathlib.Algebra.BigOperators.Group.Finset.Basic

import Mathlib.Algebra.Order.BigOperators.Group.Finset

import Mathlib.Algebra.BigOperators.Group.Finset.Piecewise

import Mathlib.Algebra.BigOperators.Ring.Finset

namespace Utils.StateTransition

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

def Transition (S : Type*) := S × S

def Run (S : Type*) := Transition S → ℕ

def countTransitionInPath [DecidableEq S] (t : Transition S) (path : List S) : ℕ :=
  (path.zip path.tail).count t

def Run.containsPath [DecidableEq S] (R : Run S) (path : List S) : Prop :=
  ∀ t : Transition S, countTransitionInPath t path ≤ R t

def Run.hasCycle [DecidableEq S] (R : Run S) : Prop :=
  ∃ (cycle : List S), cycle.length ≥ 2 ∧
    cycle.head? = cycle.getLast? ∧
    R.containsPath cycle

def Run.isAcyclic [DecidableEq S] (R : Run S) : Prop :=
  ¬R.hasCycle

def Run.reachable [DecidableEq S] (R : Run S) (start finish : S) : Prop :=
  ∃ (path : List S), path.head? = some start ∧ path.getLast? = some finish ∧
    path ≠ [] ∧ R.containsPath path

def Run.isLeaf (R : Run S) (root leaf : S) : Prop :=
  R.reachable root leaf ∧ ∀ y, R (leaf, y) = 0

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: def Transition (S : Type*) := S × S

content: def Run (S : Type*) := Transition S → ℕ

content: def countTransitionInPath [DecidableEq S] (t : Transition S) (path : List S) : ℕ :=
  (path.zip path.tail).count t

content: def Run.containsPath [DecidableEq S] (R : Run S) (path : List S) : Prop :=
  ∀ t : Transition S, countTransitionInPath t path ≤ R t

content: def Run.hasCycle [DecidableEq S] (R : Run S) : Prop :=
  ∃ (cycle : List S), cycle.length ≥ 2 ∧
    cycle.head? = cycle.getLast? ∧
    R.containsPath cycle

content: def Run.isAcyclic [DecidableEq S] (R : Run S) : Prop :=
  ¬R.hasCycle

content: def Run.reachable [DecidableEq S] (R : Run S) (start finish : S) : Prop :=
  ∃ (path : List S), path.head? = some start ∧ path.getLast? = some finish ∧
    path ≠ [] ∧ R.containsPath path

content: def Run.isLeaf (R : Run S) (root leaf : S) : Prop :=
  R.reachable root leaf ∧ ∀ y, R (leaf, y) = 0

-- Local lemmas used:
content: lemma finset_ssubset_univ_of_not_mem {α : Type*} [Fintype α] (s : Finset α) (x : α)
    (h : x ∉ s) :
    s ⊂ Finset.univ

content: lemma zip_drop_sublist (l : List S) (n : ℕ) :
    ((l.drop n).zip (l.drop (n + 1))).Sublist (l.zip l.tail)

content: lemma containsPath_drop (R : Run S) (path : List S) (n : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath (path.drop n)

content: lemma zip_take_sublist (l1 l2 : List S) (n m : ℕ) :
    ((l1.take n).zip (l2.take m)).Sublist (l1.zip l2)

content: lemma tail_take {α : Type*} (l : List α) (n : ℕ) :
    (l.take n).tail = (l.tail).take (n - 1)

content: lemma containsPath_take (R : Run S) (path : List S) (n : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath (path.take n)

content: lemma drop_take_length_ge_two {α : Type*} (path : List α) (n m : Fin path.length)
    (h_n_lt_m : n < m) :
    ((path.drop n.val).take (m.val - n.val + 1)).length ≥ 2

content: lemma getLast_drop_take {α : Type*} (path : List α) (n k : ℕ)
    (h_n_lt : n < path.length)
    (h_bound : n + k ≤ path.length)
    (h_k_pos : k > 0) :
    ((path.drop n).take k).getLast? = path[n + k - 1]?

content: lemma drop_take_cycle_same_endpoints (path : List S) (x : S) (n m : Fin path.length)
    (h_n_lt_m : n < m)
    (h_x_at_n : path[n] = x)
    (h_x_at_m : path[m] = x) :
    ((path.drop n.val).take (m.val - n.val + 1)).head? =
    ((path.drop n.val).take (m.val - n.val + 1)).getLast?

content: lemma containsPath_drop_take (R : Run S) (path : List S) (n m : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath ((path.drop n).take m)

content: lemma acyclic_containsPath_nodup (R : Run S) (path : List S)
    (h_acyclic : R.isAcyclic)
    (h_contains : R.containsPath path) :
    path.Nodup

content: lemma countTransitionInPath_append_singleton (path : List S) (x y : S)
    (h_nonempty : path ≠ [])
    (h_last : path.getLast? = some x)
    (h_not_in : (x, y) ∉ path.zip path.tail) :
    countTransitionInPath (x, y) (path ++ [y]) = 1

content: lemma countTransitionInPath_append_singleton_other (path : List S) (x y : S) (t : Transition S)
    (h_nonempty : path ≠ [])
    (h_last : path.getLast? = some x)
    (h_ne : t ≠ (x, y)) :
    countTransitionInPath t (path ++ [y]) = countTransitionInPath t path

content: lemma acyclic_no_self_loop (R : Run S) (s : S) (h_acyclic : R.isAcyclic) (h_edge : R (s, s) > 0) : False

content: lemma getLast_mem {α : Type*} (l : List α) (x : α) (h_last : l.getLast? = some x) :
    x ∈ l

content: lemma last_not_in_zip_tail {α : Type*} [DecidableEq α] (l : List α) (x : α)
    (h_nodup : l.Nodup)
    (h_last : l.getLast? = some x) :
    ∀ y : α, (x, y) ∉ l.zip l.tail

content: lemma drop_of_lt_length_nonempty {α : Type*} (path : List α) (i : ℕ)
    (h_i_lt : i < path.length) :
    path.drop i ≠ []

content: lemma cycle_from_suffix_contains (R : Run S) (suffix : List S) (current y : S)
    (h_suffix_nodup : suffix.Nodup)
    (h_contains_suffix : R.containsPath suffix)
    (h_suffix_nonempty : suffix ≠ [])
    (h_suffix_last : suffix.getLast? = some current)
    (h_edge : R (current, y) > 0) :
    ∀ t : Transition S, countTransitionInPath t (suffix ++ [y]) ≤ R t

content: lemma path_with_back_edge_creates_cycle (R : Run S) (path : List S) (current y : S)
    (h_acyclic : R.isAcyclic)
    (h_end : path.getLast? = some current)
    (h_contains : R.containsPath path)
    (h_y_in_path : y ∈ path)
    (h_edge : R (current, y) > 0) :
    R.hasCycle

content: lemma acyclic_edge_not_in_path (R : Run S) (path : List S) (current y : S)
    (h_acyclic : R.isAcyclic)
    (h_end : path.getLast? = some current)
    (h_contains : R.containsPath path)
    (h_edge : R (current, y) > 0)
    (h_y_in_path : y ∈ path) :
    False

content: lemma not_mem_implies_transition_not_in_zip_tail {α : Type*} (path : List α) (x y : α)
    (h_y_not_in : y ∉ path) :
    (x, y) ∉ path.zip path.tail

content: lemma containsPath_append_singleton (R : Run S) (path : List S) (x y : S)
    (h_nonempty : path ≠ [])
    (h_last : path.getLast? = some x)
    (h_contains : R.containsPath path)
    (h_y_not_in_path : y ∉ path)
    (h_edge : R (x, y) > 0) :
    R.containsPath (path ++ [y])

-- Target theorem and proof:
lemma acyclic_has_leaf_aux (R : Run S) (root current : S)
    (path : List S)
    (h_acyclic : R.isAcyclic)
    (h_start : path.head? = some root)
    (h_end : path.getLast? = some current)
    (h_nonempty : path ≠ [])
    (h_contains : R.containsPath path)
    (h_has_out : ∃ y, y ∉ path ∧ R (current, y) > 0) :
    ∃ leaf, R.isLeaf root leaf
:= by
  -- Get a successor not in path
  obtain ⟨y, h_y_not_in_path, h_edge⟩ := h_has_out

  -- Check if y has any outgoing edges to states not in path ++ [y]
  by_cases h_y_has_out : ∃ z, z ∉ path ∧ z ≠ y ∧ R (y, z) > 0
  case neg =>
    -- y has no outgoing edges to states not in path ++ [y]
    -- We'll show y is actually a leaf (has NO outgoing edges at all)
    use y
    constructor
    · -- Show y is reachable from root
      -- Extend the path by adding y
      use path ++ [y]
      constructor
      · aesop
      constructor
      · simp
      constructor
      · aesop
      · exact containsPath_append_singleton R path current y h_nonempty h_end h_contains h_y_not_in_path h_edge
    · -- Show y has no outgoing edges
      intro z
      by_contra h_pos
      push_neg at h_y_has_out
      -- If R(y,z) > 0, then by h_y_has_out, either z ∈ path or z = y
      have h_z_pos : R (y, z) > 0 := by omega
      have h_z_in_path_or_y : z ∈ path ∨ z = y := by grind

      -- Derive contradiction from cycle
      cases h_z_in_path_or_y with
      | inl h_z_in_path =>
        -- z ∈ path, so we can construct a cycle
        -- path ends with current, current → y, y → z, and z ∈ path
        -- This creates a cycle: (suffix of path from z to current) → y → z
        have h_z_in_extended : z ∈ path ++ [y] := by aesop
        exact acyclic_edge_not_in_path R (path ++ [y]) y z h_acyclic (by simp) (containsPath_append_singleton R path current y h_nonempty h_end h_contains h_y_not_in_path h_edge) h_z_pos h_z_in_extended
      | inr h_z_eq_y =>
        -- z = y, so we have a self-loop y → y
        rw [h_z_eq_y] at h_z_pos
        exact acyclic_no_self_loop R y h_acyclic h_z_pos

  case pos =>
    -- y has an outgoing edge to some z not in path and z ≠ y
    -- Recurse with path ++ [y]
    obtain ⟨z, h_z_not_in_path, h_z_ne_y, h_y_z_edge⟩ := h_y_has_out

    -- Construct new path
    let new_path := path ++ [y]

    -- Show properties of new_path
    have h_new_start : new_path.head? = some root := by aesop
    have h_new_end : new_path.getLast? = some y := by aesop
    have h_new_nonempty : new_path ≠ [] := by aesop
    have h_new_contains : R.containsPath new_path :=
      containsPath_append_singleton R path current y h_nonempty h_end h_contains h_y_not_in_path h_edge

    have h_new_has_out : ∃ w, w ∉ new_path ∧ R (y, w) > 0 := by grind
    exact acyclic_has_leaf_aux R root y new_path h_acyclic h_new_start h_new_end h_new_nonempty h_new_contains h_new_has_out
termination_by Fintype.card S - path.toFinset.card
decreasing_by
  simp_wf
  -- new_path = path ++ [y], and y ∉ path, so toFinset increases by 1
  have h_y_not_mem : y ∉ path.toFinset := by
    simp
    exact h_y_not_in_path
  have h_card_increase : (insert y path.toFinset).card = path.toFinset.card + 1 := by
    rw [Finset.card_insert_of_notMem h_y_not_mem]
  -- path.toFinset is a strict subset of univ
  have h_path_subset := finset_ssubset_univ_of_not_mem path.toFinset y h_y_not_mem
  have h_card_bound : path.toFinset.card < Fintype.card S := by
    rw [← Finset.card_univ]
    exact Finset.card_lt_card h_path_subset
  rw [h_card_increase]
  omega