import Mathlib

structure CapacityGraph where
  adj : Array (Array (Nat × Int))

def CapacityGraph.size (g : CapacityGraph) : Nat :=
  g.adj.size

def CapacityGraph.has_capacity (g : CapacityGraph) (u v : Nat) (c : Int) : Prop :=
  u < g.size ∧
  ∃ pair, pair ∈ g.adj.getD u #[] ∧ pair.1 = v ∧ pair.2 = c

def CapacityGraph.well_formed (g : CapacityGraph) : Prop :=
  ∀ u, u < g.size →
    (∀ pair, pair ∈ g.adj.getD u #[] → pair.1 < g.size) ∧
    -- Unique targets
    (∀ p1 p2, p1 ∈ g.adj.getD u #[] → p2 ∈ g.adj.getD u #[] → p1.1 = p2.1 → p1 = p2)

-- A flow map assigns a flow value to each edge (u, v)
-- We can model this as a function or list of entries.
-- Since the result is just the max flow value, we can keep FlowMap abstract in the spec
-- or define it as `Nat -> Nat -> Int`.
-- Dijkstra and others used abstract/implicit paths. here we need explicit flow values.
def FlowMap := Nat → Nat → Int

def FlowMap.get (f : FlowMap) (u v : Nat) : Int := f u v

def CapacityGraph.respects_capacity (g : CapacityGraph) (f : FlowMap) : Prop :=
  ∀ u v,
    (f.get u v > 0 →
      ∃ c, g.has_capacity u v c ∧ f.get u v ≤ c) ∧
    (f.get u v >= 0) -- Non-negative flow

def flow_out (g : CapacityGraph) (f : FlowMap) (u : Nat) : Int :=
  -- Sum of f(u, v) for all v. Hard to define as sum over infinite/large domain without Finset.
  -- But we can sum over the adjacency list of u.
  -- Since well_formed ensures unique neighbors, we can map adj to flow and sum.
  let neighbors := g.adj.getD u #[]
  neighbors.foldl (λ sum pair => sum + f.get u pair.1) 0

def flow_in (g : CapacityGraph) (f : FlowMap) (u : Nat) : Int :=
  -- Sum of f(v, u). This requires iterating over all nodes v that have edge to u.
  -- We can iterate over all nodes v from 0 to size-1.
  -- This is a bit computationally heavy for definition but fine for spec.
  (List.range g.size).foldl (λ sum v => sum + f.get v u) 0

def is_conserved (g : CapacityGraph) (f : FlowMap) (s t : Nat) : Prop :=
  ∀ u, u < g.size → u ≠ s → u ≠ t →
    flow_in g f u = flow_out g f u

def is_valid_flow (g : CapacityGraph) (f : FlowMap) (s t : Nat) : Prop :=
  g.respects_capacity f ∧ is_conserved g f s t

def flow_val (g : CapacityGraph) (f : FlowMap) (s : Nat) : Int :=
  flow_out g f s - flow_in g f s

def is_max_flow (g : CapacityGraph) (f : FlowMap) (s t : Nat) : Prop :=
  is_valid_flow g f s t ∧
  ∀ other, is_valid_flow g other s t → flow_val g f s ≥ flow_val g other s

def CapacityGraph.capacities_bounded (g : CapacityGraph) : Prop :=
  g.size ≤ 1000 ∧
  ∀ u v c, g.has_capacity u v c → c ≥ 0 ∧ c ≤ 100000

-- Precondition definitions
@[reducible, simp]
def max_flow_value_precond (graph : CapacityGraph) (s t : Nat) : Prop :=
  -- !benchmark @start precond
  graph.well_formed ∧
  graph.capacities_bounded ∧
  s < graph.size ∧
  t < graph.size ∧
  s ≠ t
  -- !benchmark @end precond

-- !benchmark @start auxcode
-- !benchmark @end auxcode

-- Main function definition
def max_flow_value (graph : CapacityGraph) (s t : Nat)
    (h_precond : max_flow_value_precond graph s t) : Int :=
  -- !benchmark @start code
  sorry
  -- !benchmark @end code

-- Postcondition definitions
@[reducible, simp]
def max_flow_value_postcond (graph : CapacityGraph) (s t : Nat) (result : Int)
    (_ : max_flow_value_precond graph s t) : Prop :=
  -- !benchmark @start postcond
  ∃ f, is_max_flow graph f s t ∧ flow_val graph f s = result
  -- !benchmark @end postcond

-- !benchmark @start lemma
-- !benchmark @end lemma

-- Proof content
theorem max_flow_value_postcond_satisfied (graph : CapacityGraph) (s t : Nat)
    (h_precond : max_flow_value_precond graph s t) :
    max_flow_value_postcond graph s t (max_flow_value graph s t h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof