namespace Batteries

/-- Union-find node type -/
structure UFNode where
  /-- Parent of node -/
  parent : Nat

namespace UnionFind

/-- Parent of a union-find node, defaults to self when the node is a root -/
def parentD (arr : Array UFNode) (i : Nat) : Nat :=
  if h : i < arr.size then arr[i].parent else i

/-- Rank of a union-find node, defaults to 0 when the node is a root -/
def rankD (arr : Array UFNode) (i : Nat) : Nat := 0

theorem parentD_of_not_lt : ¬i < arr.size → parentD arr i = i := sorry

theorem parentD_set {arr : Array UFNode} {x h v i} :
    parentD (arr.set x v h) i = if x = i then v.parent else parentD arr i := by
  rw [parentD]
  sorry

end UnionFind

open UnionFind

structure UnionFind where
  arr : Array UFNode

namespace UnionFind

/-- Size of union-find structure. -/
@[inline] abbrev size (self : UnionFind) := self.arr.size

/-- Parent of union-find node -/
abbrev parent (self : UnionFind) (i : Nat) : Nat := parentD self.arr i

theorem parent_lt (self : UnionFind) (i : Nat) : self.parent i < self.size ↔ i < self.size :=
  sorry

/-- Rank of union-find node -/
abbrev rank (self : UnionFind) (i : Nat) : Nat := rankD self.arr i

/-- Maximum rank of nodes in a union-find structure -/
noncomputable def rankMax (self : UnionFind) := 0

/-- Root of a union-find node. -/
def root (self : UnionFind) (x : Fin self.size) : Fin self.size :=
  let y := self.arr[x.1].parent
  if h : y = x then
    x
  else
    have : self.rankMax - self.rank self.arr[x.1].parent < self.rankMax - self.rank x :=
      sorry
    self.root ⟨y, sorry⟩
termination_by self.rankMax - self.rank x

/-- Root of a union-find node. Returns input if index is out of bounds. -/
def rootD (self : UnionFind) (x : Nat) : Nat :=
  if h : x < self.size then self.root ⟨x, h⟩ else x

theorem rootD_parent (self : UnionFind) (x : Nat) : self.rootD (self.parent x) = self.rootD x := by
  simp only [rootD, Array.length_toList, parent_lt]
  split
  · simp only [parentD, ↓reduceDIte, *]
    conv => rhs; rw [root]
    split
    · rw [root, dif_pos] <;> simp_all
    · simp
  · simp only [not_false_eq_true, parentD_of_not_lt, *]