-- from the Lean documentation
--https://lean-lang.org/functional_programming_in_lean/Programming___-Proving___-and-Performance/Insertion-Sort-and-Array-Mutation/#insertion-sort-mutation

def insertSorted [Ord α] (arr : Array α) (i : Fin arr.size) : Array α :=
  match i with
  | ⟨0, _⟩ => arr
  | ⟨i' + 1, _⟩ =>
    match Ord.compare arr[i'] arr[i] with
    | .lt | .eq => arr
    | .gt =>
      insertSorted (arr.swap i' i) ⟨i', by simp; omega⟩

theorem insert_sorted_size_eq [Ord α] (len : Nat) (i : Nat) : (arr : Array α) → (isLt : i < arr.size) →
    (arr.size = len) →                                -- look here
    (insertSorted arr ⟨i, isLt⟩).size = len := by     -- look here
  induction i with
  | zero =>
    intro arr isLt hLen
    simp [insertSorted]
    exact hLen
  | succ i' ih =>
    intro arr isLt hLen
    simp [insertSorted]
    split <;> try apply hLen
    simp [*]

def insertionSortLoop [Ord α] (arr : Array α) (i : Nat) : Array α :=
  if h : i < arr.size then
    have : (insertSorted arr ⟨i, h⟩).size - (i + 1) < arr.size - i := by
      let fact := insert_sorted_size_eq arr.size i arr h rfl
      rw [fact]
      omega
    insertionSortLoop (insertSorted arr ⟨i, h⟩) (i + 1)
  else
    arr
termination_by arr.size - i


def insertionSort [Ord α] (arr : Array α) : Array α :=
   insertionSortLoop arr 0