-- Sorted lists of natural numbers: the head bounds the tail, sortedness
-- survives prefixes, suffixes, filters, and ordered concatenation,
-- insertion sort produces a sorted permutation, and a sorted arrangement
-- of given entries is unique
-- Part of Cumulus, the Grace17 mathematics library.
-- Copyright (c) 2026 Grant Passmore.
-- Released under the Apache License 2.0. See LICENSE in the repository root.
-- Authors: Grant Passmore

module Grace.Cumulus.List.Sorted
import Prefix

-- Entries are nondecreasing along positions.
predicate ListSorted (xs : Set) :=
  ∀ i, ∀ j, Nat(i) → Nat(j) → i < j → j < length{ℕ}(xs) → listCoefficient{ℕ}(0, xs, i) ≤ listCoefficient{ℕ}(0, xs, j)

theorem ListSorted.nil
  shows ListSorted([])
proof
  unfold ListSorted
  intros i j iNat jNat below bound
  rw length.nil at bound
  exfalso
  arith

theorem ListSorted.singleton (x : Nat)
  shows ListSorted(x :: [])
proof
  unfold ListSorted
  intros i j iNat jNat below bound
  have xIn : x ∈ ℕ := nat.omega.forward with *
  rw length.singleton at bound
  exfalso
  arith

theorem ListSorted.head_le (a : Nat) (xs : List(ℕ))
  requires sorted : ListSorted(a :: xs)
  shows ∀ y ∈ listSupport{ℕ}(xs), a ≤ y
proof
  intros y hy
  have aIn : a ∈ ℕ := nat.omega.forward with *
  obtain ⟨k, kNat, bound, value⟩ := listSupport.mem_iff_index.forward at ℕ, 0, y, xs with hy
  have kSucc : Nat(succ(k)) by typecheck
  have positive : 0 < succ(k) by arith
  have inside : succ(k) < length{ℕ}(a :: xs)
    rw length.cons
    arith
  unfold ListSorted at sorted
  have step := sorted with 0, succ(k), Naturals.constructor.zero, kSucc, positive, inside
  rw [listCoefficient.zero, listCoefficient.succ, value] at step
  exact step

theorem ListSorted.tail (a : Nat) (xs : List(ℕ))
  requires sorted : ListSorted(a :: xs)
  shows ListSorted(xs)
proof
  unfold ListSorted
  intros i j iNat jNat below bound
  have aIn : a ∈ ℕ := nat.omega.forward with *
  have iSucc : Nat(succ(i)) by typecheck
  have jSucc : Nat(succ(j)) by typecheck
  have belowSucc : succ(i) < succ(j) by arith
  have inside : succ(j) < length{ℕ}(a :: xs)
    rw length.cons
    arith
  unfold ListSorted at sorted
  have step := sorted with succ(i), succ(j), iSucc, jSucc, belowSucc, inside
  rw [listCoefficient.succ, listCoefficient.succ] at step
  exact step

theorem ListSorted.cons (a : Nat) (xs : List(ℕ))
  requires head : ∀ y ∈ listSupport{ℕ}(xs), a ≤ y
  requires sorted : ListSorted(xs)
  shows ListSorted(a :: xs)
proof
  unfold ListSorted
  intros i j iNat jNat below bound
  have aIn : a ∈ ℕ := nat.omega.forward with *
  rw length.cons at bound
  have shapeJ := natural.zero_or_successor at j with jNat
  cases shapeJ with zeroJ | successorJ
  . rw zeroJ at below
    exfalso
    arith
  . obtain ⟨q, qNat, equalJ⟩ := successorJ
    rw equalJ at bound
    have insideQ : q < length{ℕ}(xs) by arith
    have memberQ := listCoefficient.mem_support at ℕ, 0, xs, q with insideQ
    have shapeI := natural.zero_or_successor at i with iNat
    cases shapeI with zeroI | successorI
    . rw [zeroI, equalJ, listCoefficient.zero, listCoefficient.succ]
      exact head with listCoefficient{ℕ}(0, xs, q), memberQ
    . obtain ⟨p, pNat, equalI⟩ := successorI
      rw [equalI, equalJ] at below
      have belowPQ : p < q by arith
      unfold ListSorted at sorted
      have step := sorted with p, q, pNat, qNat, belowPQ, insideQ
      rw [equalI, equalJ, listCoefficient.succ, listCoefficient.succ]
      exact step

theorem ListSorted.cons_iff (a : Nat) (xs : List(ℕ))
  shows ListSorted(a :: xs) ↔ ((∀ y ∈ listSupport{ℕ}(xs), a ≤ y) ∧ ListSorted(xs))
proof
  constructor
  . intro sorted
    have head := ListSorted.head_le with sorted
    have tail := ListSorted.tail with sorted
    exact ⟨head, tail⟩
  . intro both
    obtain ⟨head, tail⟩ := both
    exact ListSorted.cons with head, tail

theorem ListSorted.cons_cons_le (a b : Nat) (xs : List(ℕ))
  requires sorted : ListSorted(a :: b :: xs)
  shows a ≤ b
proof
  have bIn : b ∈ ℕ := nat.omega.forward with *
  have head := ListSorted.head_le with sorted
  have member := listSupport.mem_cons_self at ℕ, b, xs with bIn
  exact head with b, member

theorem ListSorted.min_head (a : Nat) (xs : List(ℕ))
  requires sorted : ListSorted(a :: xs)
  shows ∀ y ∈ listSupport{ℕ}(a :: xs), a ≤ y
proof
  intros y hy
  have aIn : a ∈ ℕ := nat.omega.forward with *
  have head := ListSorted.head_le with sorted
  have split := listSupport.mem_cons_iff.forward with aIn, hy
  cases split with same | rest
  . rw same
    exact LessEqual.reflexive at a
  . exact head with y, rest

theorem ListSorted.replicate (x : Nat) (n : Nat)
  shows ListSorted(listReplicate(n, x))
proof
  induction n
  | zero =>
    rw listReplicate.zero
    exact ListSorted.nil
  | succ(k) using ih =>
    have xIn : x ∈ ℕ := nat.omega.forward with *
    have list := listReplicate.closed at ℕ, x, k with xIn
    have head : ∀ y ∈ listSupport{ℕ}(listReplicate(k, x)), x ≤ y
      intros y hy
      have equal := listReplicate.mem_support with xIn, hy
      rw equal
      exact LessEqual.reflexive at x
    rw listReplicate.succ
    exact ListSorted.cons with head, ih

theorem ListSorted.take (xs : List(ℕ)) (n : Nat)
  requires sorted : ListSorted(xs)
  shows ListSorted(listTake{ℕ}(n, xs))
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat sorted
    rw listTake.nil
    exact ListSorted.nil
  | List.cons(a, tail) using ih =>
    intros n nNat sorted
    have aIn : a ∈ ℕ by assumption
    have aNat := nat.omega.backward with aIn
    have shape := natural.zero_or_successor at n with nNat
    cases shape with zero | successor
    . have whole : List(ℕ, a :: tail) by typecheck
      rw [zero, listTake.zero]
      exact ListSorted.nil
    . obtain ⟨k, kNat, equal⟩ := successor
      have headAll := ListSorted.head_le with sorted
      have tailSorted := ListSorted.tail with sorted
      have inner := ih with k, kNat, tailSorted
      have taken := listTake.closed at ℕ, tail, k with *
      have subset := listTake.support_subset at ℕ, tail, k with *
      have head : ∀ y ∈ listSupport{ℕ}(listTake{ℕ}(k, tail)), a ≤ y
        intros y hy
        have member := subset with y, hy
        exact headAll with y, member
      rw [equal, listTake.succ_cons]
      exact ListSorted.cons with head, inner

theorem ListSorted.drop (xs : List(ℕ)) (n : Nat)
  requires sorted : ListSorted(xs)
  shows ListSorted(listDrop{ℕ}(n, xs))
proof
  induction xs generalizing n
  | List.nil =>
    intros n nNat sorted
    rw listDrop.nil
    exact ListSorted.nil
  | List.cons(a, tail) using ih =>
    intros n nNat sorted
    have aIn : a ∈ ℕ by assumption
    have aNat := nat.omega.backward with aIn
    have shape := natural.zero_or_successor at n with nNat
    cases shape with zero | successor
    . have whole : List(ℕ, a :: tail) by typecheck
      rw [zero, listDrop.zero]
      exact sorted
    . obtain ⟨k, kNat, equal⟩ := successor
      have tailSorted := ListSorted.tail with sorted
      rw [equal, listDrop.succ_cons]
      exact ih with k, kNat, tailSorted

theorem ListSorted.filter (S : Set) (xs : List(ℕ))
  requires sorted : ListSorted(xs)
  shows ListSorted(listFilter{ℕ}(S, xs))
proof
  induction xs
  | List.nil =>
    intro sorted
    rw listFilter.nil
    exact ListSorted.nil
  | List.cons(a, tail) using ih =>
    intro sorted
    have aIn : a ∈ ℕ by assumption
    have aNat := nat.omega.backward with aIn
    have headAll := ListSorted.head_le with sorted
    have tailSorted := ListSorted.tail with sorted
    have inner := ih with tailSorted
    have filtered := listFilter.closed at ℕ, S, tail with *
    have alternatives : a ∈ S ∨ ¬(a ∈ S) by tauto
    cases alternatives with member | absent
    . have subset := listFilter.support_subset at ℕ, S, tail with *
      have head : ∀ y ∈ listSupport{ℕ}(listFilter{ℕ}(S, tail)), a ≤ y
        intros y hy
        have inTail := subset with y, hy
        exact headAll with y, inTail
      rw listFilter.cons_member
      exact ListSorted.cons with head, inner
    . rw listFilter.cons_other
      exact inner

theorem ListSorted.append (xs ys : List(ℕ))
  requires first : ListSorted(xs)
  requires second : ListSorted(ys)
  requires between : ∀ x ∈ listSupport{ℕ}(xs), ∀ y ∈ listSupport{ℕ}(ys), x ≤ y
  shows ListSorted(append(ys, xs))
proof
  induction xs
  | List.nil =>
    intros first between
    have ysList : List(ℕ, ys) by typecheck
    rw append.nil
    exact second
  | List.cons(a, tail) using ih =>
    intros first between
    have aIn : a ∈ ℕ by assumption
    have aNat := nat.omega.backward with aIn
    have headAll := ListSorted.head_le with first
    have tailSorted := ListSorted.tail with first
    have headIn := listSupport.mem_cons_self at ℕ, a, tail with aIn
    have restBetween : ∀ x ∈ listSupport{ℕ}(tail), ∀ y ∈ listSupport{ℕ}(ys), x ≤ y
      intros x hx y hy
      have outer := listSupport.mem_cons_of_mem with aIn, hx
      exact between with x, outer, y, hy
    have inner := ih with tailSorted, restBetween
    have appended := append.closed at ℕ, tail, ys with *
    have head : ∀ y ∈ listSupport{ℕ}(append(ys, tail)), a ≤ y
      intros y hy
      have split := listSupport.mem_append_iff.forward at ℕ, y, tail, ys with hy
      cases split with inTail | inYs
      . exact headAll with y, inTail
      . exact between with a, headIn, y, inYs
    rw append.cons
    exact ListSorted.cons with head, inner

-- A sorted arrangement of given entries is unique.
theorem ListSorted.permutation_unique (xs ys : List(ℕ))
  requires first : ListSorted(xs)
  requires second : ListSorted(ys)
  requires permutation : ListPermutation(ℕ, xs, ys)
  shows xs = ys
proof
  induction xs generalizing ys
  | List.nil =>
    intros ys ysList first second permutation
    have back := ListPermutation.symmetric with permutation
    have empty := ListPermutation.nil with back
    rw empty
    rfl
  | List.cons(a, tail) using ih =>
    intros ys ysList first second permutation
    have aIn : a ∈ ℕ by assumption
    have aNat := nat.omega.backward with aIn
    have shape := List.cases with ysList
    cases shape with nil | cons
    . rw nil at permutation
      have empty := ListPermutation.nil with permutation
      have distinct := List.cons_ne_nil at ℕ, a, tail with aIn
      exfalso
      exact distinct with empty
    . obtain ⟨b, bIn, rest, restList, equal⟩ := cons
      have bNat := nat.omega.backward with bIn
      rw equal at permutation
      rw equal at second
      have headIn := listSupport.mem_cons_self at ℕ, a, tail with aIn
      have aInYs := ListPermutation.mem_support with permutation, headIn
      have back := ListPermutation.symmetric with permutation
      have headInYs := listSupport.mem_cons_self at ℕ, b, rest with bIn
      have bInXs := ListPermutation.mem_support with back, headInYs
      have minA := ListSorted.min_head with first
      have minB := ListSorted.min_head with second
      have aLeB := minA with b, bInXs
      have bLeA := minB with a, aInYs
      have same := LessEqual.antisymmetric with aLeB, bLeA
      rw ← same at permutation
      have tails := ListPermutation.cons_cancel with aIn, permutation
      have tailSorted := ListSorted.tail with first
      rw ← same at second
      have restSorted := ListSorted.tail with second
      have inner := ih with rest, restList, tailSorted, restSorted, tails
      rw [equal, ← same, inner]
      rfl

-- Insertion into a sorted list, and insertion sort.
theorem memberTest.successor_le (x y : Nat)
  requires le : x ≤ y
  shows memberTest(succ(y), x) = Bool.true
proof
  have below : x < succ(y) by arith
  have member := membership.of_Less at x, succ(y) with *, below
  exact memberTest.true with member

theorem memberTest.successor_gt (x y : Nat)
  requires gt : y < x
  shows memberTest(succ(y), x) = Bool.false
proof
  have absent : ¬(x ∈ succ(y))
    intro member
    have below := Less.of_membership at x, succ(y) with *, member
    have le := Less.successor_upper with below
    have strict : x < x by arith
    exact Less.irreflexive with *, strict
  exact memberTest.false with absent

def listInsert (x : Nat, xs : List(ℕ)) → Set
  | x, [] => x :: []
  | x, y :: rest =>
    boolChoice(memberTest(succ(y), x), x :: y :: rest, y :: listInsert(x, rest))

theorem listInsert.nil (x : Nat)
  shows listInsert(x, []) = x :: []
proof
  simp [listInsert]

theorem listInsert.cons_le (x y : Nat) (rest : List(ℕ))
  requires le : x ≤ y
  shows listInsert(x, y :: rest) = x :: y :: rest
proof
  have test := memberTest.successor_le with le
  simp [listInsert, test, boolChoice]

theorem listInsert.cons_gt (x y : Nat) (rest : List(ℕ))
  requires gt : y < x
  shows listInsert(x, y :: rest) = y :: listInsert(x, rest)
proof
  have test := memberTest.successor_gt with gt
  simp [listInsert, test, boolChoice]

theorem listInsert.closed (x : Nat) (xs : List(ℕ))
  shows List(ℕ, listInsert(x, xs))
proof
  induction xs
  | List.nil =>
    have xIn : x ∈ ℕ := nat.omega.forward with *
    rw listInsert.nil
    typecheck
  | List.cons(y, rest) using ih =>
    have xIn : x ∈ ℕ := nat.omega.forward with *
    have yIn : y ∈ ℕ by assumption
    have yNat := nat.omega.backward with yIn
    have alternatives : x ≤ y ∨ ¬(x ≤ y) by tauto
    cases alternatives with le | notLe
    . rw listInsert.cons_le
      typecheck
    . have gt := Less.of_not_LessEqual at x, y with notLe
      rw listInsert.cons_gt
      typecheck

theorem listInsert.permutation (x : Nat) (xs : List(ℕ))
  shows ListPermutation(ℕ, listInsert(x, xs), x :: xs)
proof
  induction xs
  | List.nil =>
    have xIn : x ∈ ℕ := nat.omega.forward with *
    rw listInsert.nil
    exact ListPermutation.reflexive at ℕ, x :: [] with *
  | List.cons(y, rest) using ih =>
    have xIn : x ∈ ℕ := nat.omega.forward with *
    have yIn : y ∈ ℕ by assumption
    have yNat := nat.omega.backward with yIn
    have inserted := listInsert.closed at x, rest with *
    have whole : List(ℕ, y :: rest) by typecheck
    have alternatives : x ≤ y ∨ ¬(x ≤ y) by tauto
    cases alternatives with le | notLe
    . rw listInsert.cons_le
      exact ListPermutation.reflexive at ℕ, x :: y :: rest with *
    . have gt := Less.of_not_LessEqual at x, y with notLe
      rw listInsert.cons_gt
      have step := ListPermutation.cons at ℕ, y, listInsert(x, rest), x :: rest with yIn, ih
      have swap := ListPermutation.swap at ℕ, y, x, rest with yIn, xIn
      exact ListPermutation.transitive with step, swap

theorem listInsert.length (x : Nat) (xs : List(ℕ))
  shows length{ℕ}(listInsert(x, xs)) = succ(length{ℕ}(xs))
proof
  have xIn : x ∈ ℕ := nat.omega.forward with *
  have inserted := listInsert.closed at x, xs with *
  have whole : List(ℕ, x :: xs) by typecheck
  have permutation := listInsert.permutation at x, xs with *
  have lengths := ListPermutation.length with permutation
  rw [lengths, length.cons]
  rfl

theorem listInsert.count (x z : Nat) (xs : List(ℕ))
  shows listCount{ℕ}(z, listInsert(x, xs)) = listCount{ℕ}(z, x :: xs)
proof
  have xIn : x ∈ ℕ := nat.omega.forward with *
  have inserted := listInsert.closed at x, xs with *
  have whole : List(ℕ, x :: xs) by typecheck
  have permutation := listInsert.permutation at x, xs with *
  exact ListPermutation.count with permutation

theorem listInsert.support (x : Nat) (xs : List(ℕ))
  shows listSupport{ℕ}(listInsert(x, xs)) = insert(x, listSupport{ℕ}(xs))
proof
  have xIn : x ∈ ℕ := nat.omega.forward with *
  have inserted := listInsert.closed at x, xs with *
  have whole : List(ℕ, x :: xs) by typecheck
  have permutation := listInsert.permutation at x, xs with *
  have support := ListPermutation.support with permutation
  rw [support, listSupport.cons]
  rfl

theorem listInsert.sorted (x : Nat) (xs : List(ℕ))
  requires sorted : ListSorted(xs)
  shows ListSorted(listInsert(x, xs))
proof
  induction xs
  | List.nil =>
    intro sorted
    rw listInsert.nil
    exact ListSorted.singleton at x with *
  | List.cons(y, rest) using ih =>
    intro sorted
    have xIn : x ∈ ℕ := nat.omega.forward with *
    have yIn : y ∈ ℕ by assumption
    have yNat := nat.omega.backward with yIn
    have inserted := listInsert.closed at x, rest with *
    have whole : List(ℕ, y :: rest) by typecheck
    have headAll := ListSorted.head_le with sorted
    have restSorted := ListSorted.tail with sorted
    have alternatives : x ≤ y ∨ ¬(x ≤ y) by tauto
    cases alternatives with le | notLe
    . rw listInsert.cons_le
      have head : ∀ z ∈ listSupport{ℕ}(y :: rest), x ≤ z
        intros z hz
        have split := listSupport.mem_cons_iff.forward with yIn, hz
        cases split with same | inRest
        . rw same
          exact le
        . have subset := listSupport.subset at ℕ, rest with *
          have zIn := subset with z, inRest
          have zNat := nat.omega.backward with zIn
          have yz := headAll with z, inRest
          exact LessEqual.transitive with le, yz
      exact ListSorted.cons with head, sorted
    . have gt := Less.of_not_LessEqual at x, y with notLe
      have inner := ih with restSorted
      rw listInsert.cons_gt
      have head : ∀ z ∈ listSupport{ℕ}(listInsert(x, rest)), y ≤ z
        intros z hz
        rw listInsert.support at hz
        have split := insert.membership.forward with hz
        cases split with same | inRest
        . rw same
          exact Less.implies_LessEqual with gt
        . exact headAll with z, inRest
      exact ListSorted.cons with head, inner

def listInsertionSort (xs : List(ℕ)) → Set
  | [] => []
  | x :: rest => listInsert(x, listInsertionSort(rest))

theorem listInsertionSort.nil
  shows listInsertionSort([]) = []
proof
  simp [listInsertionSort]

theorem listInsertionSort.cons (x : Nat) (rest : List(ℕ))
  shows listInsertionSort(x :: rest) = listInsert(x, listInsertionSort(rest))
proof
  simp [listInsertionSort]

theorem listInsertionSort.closed (xs : List(ℕ))
  shows List(ℕ, listInsertionSort(xs))
proof
  induction xs
  | List.nil =>
    rw listInsertionSort.nil
    typecheck
  | List.cons(x, rest) using ih =>
    have xIn : x ∈ ℕ by assumption
    have xNat := nat.omega.backward with xIn
    rw listInsertionSort.cons
    exact listInsert.closed at x, listInsertionSort(rest) with *

theorem listInsertionSort.permutation (xs : List(ℕ))
  shows ListPermutation(ℕ, listInsertionSort(xs), xs)
proof
  induction xs
  | List.nil =>
    rw listInsertionSort.nil
    exact ListPermutation.reflexive at ℕ, [] with *
  | List.cons(x, rest) using ih =>
    have xIn : x ∈ ℕ by assumption
    have xNat := nat.omega.backward with xIn
    have sortedRest := listInsertionSort.closed at rest with *
    have inserted := listInsert.closed at x, listInsertionSort(rest) with *
    have consSorted : List(ℕ, x :: listInsertionSort(rest)) by typecheck
    rw listInsertionSort.cons
    have step := listInsert.permutation at x, listInsertionSort(rest) with *
    have consStep := ListPermutation.cons at ℕ, x, listInsertionSort(rest), rest with xIn, ih
    exact ListPermutation.transitive with step, consStep

theorem listInsertionSort.sorted (xs : List(ℕ))
  shows ListSorted(listInsertionSort(xs))
proof
  induction xs
  | List.nil =>
    rw listInsertionSort.nil
    exact ListSorted.nil
  | List.cons(x, rest) using ih =>
    have xIn : x ∈ ℕ by assumption
    have xNat := nat.omega.backward with xIn
    have sortedRest := listInsertionSort.closed at rest with *
    rw listInsertionSort.cons
    exact listInsert.sorted with ih

theorem listInsertionSort.length (xs : List(ℕ))
  shows length{ℕ}(listInsertionSort(xs)) = length{ℕ}(xs)
proof
  have sorted := listInsertionSort.closed at xs with *
  have permutation := listInsertionSort.permutation at xs with *
  exact ListPermutation.length with permutation

theorem listInsertionSort.support (xs : List(ℕ))
  shows listSupport{ℕ}(listInsertionSort(xs)) = listSupport{ℕ}(xs)
proof
  have sorted := listInsertionSort.closed at xs with *
  have permutation := listInsertionSort.permutation at xs with *
  exact ListPermutation.support with permutation

theorem listInsertionSort.count (x : Set) (xs : List(ℕ))
  shows listCount{ℕ}(x, listInsertionSort(xs)) = listCount{ℕ}(x, xs)
proof
  have sorted := listInsertionSort.closed at xs with *
  have permutation := listInsertionSort.permutation at xs with *
  exact ListPermutation.count with permutation

theorem listInsertionSort.of_sorted (xs : List(ℕ))
  requires sorted : ListSorted(xs)
  shows listInsertionSort(xs) = xs
proof
  have list := listInsertionSort.closed at xs with *
  have resultSorted := listInsertionSort.sorted at xs with *
  have permutation := listInsertionSort.permutation at xs with *
  exact ListSorted.permutation_unique with resultSorted, sorted, permutation

theorem listInsertionSort.idempotent (xs : List(ℕ))
  shows listInsertionSort(listInsertionSort(xs)) = listInsertionSort(xs)
proof
  have list := listInsertionSort.closed at xs with *
  have resultSorted := listInsertionSort.sorted at xs with *
  exact listInsertionSort.of_sorted with resultSorted

theorem listInsertionSort.of_permutation (xs ys : List(ℕ))
  requires permutation : ListPermutation(ℕ, xs, ys)
  shows listInsertionSort(xs) = listInsertionSort(ys)
proof
  have leftList := listInsertionSort.closed at xs with *
  have rightList := listInsertionSort.closed at ys with *
  have leftSorted := listInsertionSort.sorted at xs with *
  have rightSorted := listInsertionSort.sorted at ys with *
  have left := listInsertionSort.permutation at xs with *
  have right := listInsertionSort.permutation at ys with *
  have rightBack := ListPermutation.symmetric with right
  have middle := ListPermutation.transitive with left, permutation
  have whole := ListPermutation.transitive with middle, rightBack
  exact ListSorted.permutation_unique with leftSorted, rightSorted, whole