import Mathlib

structure UnionFind where
  parent : List Nat
  rank   : List Nat -- Added to match Dafny/Verus
deriving Inhabited

namespace UnionFind

def len (uf : UnionFind) : Nat := uf.parent.length

def is_valid (uf : UnionFind) : Prop :=
  uf.parent.length = uf.rank.length ∧
  ∀ i, i < uf.len →
    let p := uf.parent.getD i 0
    p < uf.len ∧
    uf.rank.getD i 0 < uf.len ∧
    -- CRITICAL: Rank monotonicity enforces acyclicity (DAG structure)
    (p ≠ i → uf.rank.getD i 0 < uf.rank.getD p 0)

-- Measure for termination proof, matching Dafny
def measure (uf : UnionFind) (i : Nat) : Nat :=
  if h : i < uf.len then
    uf.len - uf.rank.getD i 0
  else 0

-- The Abstract Model ("Ghost Spec")
def spec_find (uf : UnionFind) (i : Nat) : Nat :=
  if h : i < uf.len ∧ uf.rank.getD i 0 < uf.len then
    let p := uf.parent.getD i 0
    -- If not root and parent has higher rank (valid edge up)
    if p ≠ i ∧ p < uf.len ∧ uf.rank.getD i 0 < uf.rank.getD p 0 then
      spec_find uf p
    else
      i
  else
    i
termination_by measure uf i
decreasing_by sorry -- Dafny/Verus solve this automatically; Lean needs 'sorry' or manual proof

-- Precondition definitions
@[reducible, simp]
def find_precond (uf : UnionFind) (i : Nat) : Prop :=
  -- !benchmark @start precond
  uf.is_valid ∧ i < uf.len
  -- !benchmark @end precond

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

-- Main function definition
def find (uf : UnionFind) (i : Nat)
    (h_precond : find_precond uf i) : Nat × UnionFind :=
  -- !benchmark @start code
  sorry
  -- !benchmark @end code

-- Postcondition definitions
@[reducible, simp]
def find_postcond (uf : UnionFind) (i : Nat) (result : Nat × UnionFind)
    (pre : find_precond uf i) : Prop :=
  -- !benchmark @start postcond
  let (root, new_uf) := result
  -- 1. Validity preservation
  new_uf.is_valid ∧
  -- 2. Functional Correctness: Returns the specific representative
  root = spec_find uf i ∧
  -- 3. Logical Stability: Path compression doesn't change sets
  (∀ j, j < uf.len → spec_find new_uf j = spec_find uf j) ∧
  -- 4. Structure: Length and Ranks preserved
  new_uf.len = uf.len ∧
  new_uf.rank = uf.rank
  -- !benchmark @end postcond

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

-- Proof content
theorem find_postcond_satisfied (uf : UnionFind) (i : Nat)
    (h_precond : find_precond uf i) :
    find_postcond uf i (find uf i h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof

end UnionFind