-- 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 → ℕ

noncomputable def Run.netFlow {S : Type*} [Fintype S] [DecidableEq S] (R : Run S) (x : S) : ℤ :=
  (∑ y : S, (R (x, y) : ℤ)) - (∑ y : S, (R (y, x) : ℤ))

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: noncomputable def Run.netFlow {S : Type*} [Fintype S] [DecidableEq S] (R : Run S) (x : S) : ℤ :=
  (∑ y : S, (R (x, y) : ℤ)) - (∑ y : S, (R (y, x) : ℤ))

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 containsPath_has_positive_transition (R : Run S) (path : List S)
    (h_contains : R.containsPath path) (t : Transition S)
    (h_in : t ∈ path.zip path.tail) :
    R t > 0

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])

content: 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

content: lemma acyclic_has_leaf (R : Run S) (root : S)
    (h_acyclic : R.isAcyclic)
    (h_has_out : ∃ y, R (root, y) > 0) :
    ∃ leaf, R.isLeaf root leaf

content: lemma single_le_sum_of_nonneg {α : Type*} [Fintype α] (f : α → ℤ) (a : α)
    (h_nonneg : ∀ x, f x ≥ 0) :
    f a ≤ ∑ x : α, f x

content: lemma sum_pos_of_pos_element {α : Type*} [Fintype α] (f : α → ℤ) (a : α)
    (h_pos : f a > 0)
    (h_nonneg : ∀ x, f x ≥ 0) :
    ∑ x : α, f x > 0

content: lemma sum_nat_cast_pos {α : Type*} [Fintype α] (f : α → ℕ) (a : α)
    (h_pos : f a > 0) :
    ∑ x : α, (f x : ℤ) > 0

content: lemma leaf_has_negative_netFlow (R : Run S) (root leaf : S)
    (h_leaf : R.isLeaf root leaf)
    (h_in : ∃ y, R (y, leaf) > 0) :
    R.netFlow leaf < 0

content: lemma sum_zero_of_all_zero {α : Type*} [Fintype α] (f : α → ℕ) (h : ∀ x, f x = 0) :
    ∑ x : α, (f x : ℤ) = 0

content: lemma sum_nat_cast_nonneg {α : Type*} [Fintype α] (f : α → ℕ) :
    0 ≤ ∑ x : α, (f x : ℤ)

content: lemma sum_nat_cast_zero_of_not_pos {α : Type*} [Fintype α] (f : α → ℕ)
    (h : ∀ x, ¬(f x > 0)) :
    ∑ x : α, (f x : ℤ) = 0

content: lemma positive_netFlow_has_outgoing_edge (R : Run S) (s : S)
    (h_pos : R.netFlow s > 0) :
    ∃ y, R (s, y) > 0

content: lemma last_has_incoming_transition {α : Type*} (l : List α) (x : α)
    (h_len : l.length ≥ 2)
    (h_last : l.getLast? = some x) :
    ∃ y, (y, x) ∈ l.zip l.tail

content: lemma path_distinct_head_last_length_ge_two {α : Type*} (path : List α) (x y : α)
    (h_nonempty : path ≠ [])
    (h_head : path.head? = some x)
    (h_last : path.getLast? = some y)
    (h_ne : x ≠ y) :
    path.length ≥ 2

content: lemma reachable_leaf_has_incoming_edge (R : Run S) (root leaf : S)
    (h_leaf : R.isLeaf root leaf)
    (h_ne : root ≠ leaf) :
    ∃ y, R (y, leaf) > 0

content: lemma unique_negative_netFlow (R : Run S) (s d x : S)
    (h_source : R.netFlow s = 1)
    (h_others : ∀ y, y ≠ s → y ≠ d → R.netFlow y = 0)
    (h_x_neg : R.netFlow x < 0) :
    x = d

content: lemma leaf_has_incoming_and_negative_netFlow (R : Run S) (root leaf : S)
    (h_leaf : R.isLeaf root leaf)
    (h_root_out : ∃ y, R (root, y) > 0) :
    R.netFlow leaf < 0

-- Target theorem and proof:
lemma acyclic_run_has_path_from_source_to_sink (R : Run S) (s d : S)
    (h_acyclic : R.isAcyclic)
    (h_source : R.netFlow s = 1)
    (h_others : ∀ x, x ≠ s → x ≠ d → R.netFlow x = 0) :
    ∃ (path : List S), path.head? = some s ∧ path.getLast? = some d ∧
      path ≠ [] ∧ R.containsPath path ∧ path.Nodup
:= by
  -- s has positive net flow, so it has an outgoing edge
  have h_s_out : ∃ y, R (s, y) > 0 := by
    apply positive_netFlow_has_outgoing_edge
    rw [h_source]
    omega

  -- Find a leaf reachable from s
  obtain ⟨leaf, h_leaf⟩ := acyclic_has_leaf R s h_acyclic h_s_out

  -- The leaf has negative net flow
  have h_leaf_neg := leaf_has_incoming_and_negative_netFlow R s leaf h_leaf h_s_out

  -- Identify leaf = d (only state with negative net flow)
  have h_leaf_eq_d := unique_negative_netFlow R s d leaf h_source h_others h_leaf_neg

  -- Extract the path from s to d
  rw [h_leaf_eq_d] at h_leaf
  obtain ⟨h_reach, _⟩ := h_leaf
  obtain ⟨path, h_head, h_last, h_nonempty, h_contains⟩ := h_reach
  use path
  refine ⟨h_head, h_last, h_nonempty, h_contains, ?_⟩
  exact acyclic_containsPath_nodup R path h_acyclic h_contains