-- see also the following similar solution:
-- https://jamesoswald.dev/posts/lean4-insertion-sort/

-- TYPES --
-- we use a List on Ints for the sake of readability
-- we could abstract to other types α that have a complete linear order [Ord α]
def IntList := List Int

-- ALGORITHMS --

-- algorithm for inserting a new item into an already sorted list
def insertSorted (i : Int) : IntList → IntList
    | [] => [i]
    | y :: ys => if i < y then
        i :: y :: ys
      else
        y :: (insertSorted i ys)

-- algorithm insertion sort
def insertionSort : IntList → IntList
  | [] => []
  | x :: xs => insertSorted x (insertionSort xs)

-- PROPERTY --
-- "List is sorted"
def sortedList : IntList → Prop
  | [] => true
  | _ :: [] => true
  | a :: b :: cs => (a ≤ b) ∧ sortedList (b :: cs)

-- THEOREM THAT LIST IS SORTED --

-- correctness proof for insertSorted
theorem iSRetainsSorting (xs: IntList) (invariant : sortedList xs) : ∀ x : Int, sortedList (insertSorted x xs) := by
  intro x
  induction xs with
  | nil => rfl
  | cons a bs th =>
    by_cases l : x < a
    · grind [sortedList, insertSorted]
    · match bs with
        | [] =>
          grind [sortedList, insertSorted]
        | b :: t =>
          grind [sortedList, insertSorted]

-- inductively use iSRetainsSorting to show correctness of insertionSort itself
theorem iSIsSorted : ∀ xs : IntList, sortedList (insertionSort xs) := by
  intro xs
  induction xs with
  | nil => simp [insertionSort, sortedList]
  | cons h t ih =>
    unfold insertionSort
    let betterFact := iSRetainsSorting (insertionSort t) ih h
    exact betterFact


-- THEOREM THAT OUTPUT LIST IS PERMUTATION OF INPUT LIST --

-- to access List.Perm as `~`
open List

theorem insertSortedPerm (x : Int) (xs : IntList) : (x :: xs) ~ (insertSorted x xs) := by
  induction xs with
  | nil => simp [insertSorted]
  | cons a bs th =>
    unfold insertSorted
    by_cases l : x < a
    · simp [l]
    · simp [l]
      let fact := List.Perm.swap a x bs
      let elaboratedTh := List.Perm.cons a th
      exact List.Perm.trans fact elaboratedTh

theorem iSPerm : ∀ xs : IntList, xs ~ insertionSort xs := by
  intro ys
  induction ys with
  | nil => unfold insertionSort; simp
  | cons a bs ih =>
    unfold insertionSort
    let fact₁ := List.Perm.cons a ih
    let fact₂ := insertSortedPerm a (insertionSort bs)
    exact List.Perm.trans fact₁ fact₂

-- EXAMPLES --
#eval insertSorted 3 [1,2,3,4,5]

example :  sortedList []      := by rfl
example :  sortedList [1]     := by rfl
example :  sortedList [1,2]   := by simp [sortedList]
example :  sortedList [1,2,3] := by simp [sortedList]
example : ¬sortedList [3,2]   := by simp [sortedList]
example : ¬sortedList [3,2,3] := by simp [sortedList]

#eval insertionSort [3,7,4,9,5]
#eval insertionSort [2,2,3,5,6,1,2,10,4,3]