// Copyright 2026 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///|
priv struct TimSortRun {
  len : Int
  start : Int
}

///|
/// The algorithm identifies strictly descending and non-descending subsequences, which are called
/// natural runs. There is a stack of pending runs yet to be merged. Each newly found run is pushed
/// onto the stack, and then some pairs of adjacent runs are merged until these two invariants are
/// satisfied:
///
/// 1. for every `i` in `1..runs.len()`: `runs[i - 1].len > runs[i].len`
/// 2. for every `i` in `2..runs.len()`: `runs[i - 2].len > runs[i - 1].len + runs[i].len`
///
/// The invariants ensure that the total running time is *O*(*n* \* log(*n*)) worst-case.
fn[T : Compare] timsort(arr : MutArrayView[T]) -> Unit {
  // Slices of up to this length get sorted using insertion sort.
  let max_insertion = 20

  // Short arrays get sorted in-place via insertion sort to avoid allocations.
  let len = arr.length()
  if len <= max_insertion {
    MutArrayView::insertion_sort(arr)
    return
  }
  let runs : Array[TimSortRun] = []
  for start = 0, end = 0; end < len; {
    let (streak_end, was_reversed) = find_streak(arr.mut_view(start~))
    let end = end + streak_end
    if was_reversed {
      arr.mut_view(start~, end~).rev_in_place()
    }
    // Insert some more elements into the run if it's too short. Insertion sort is faster than
    // merge sort on short sequences, so this significantly improves performance.
    let end = provide_sorted_batch(arr, start, end)
    runs.push({ start, len: end - start })
    while true {
      guard collapse(runs, len) is Some(r) else { break }
      let left = runs[r]
      let right = runs[r + 1]
      merge(arr.slice(left.start, right.start + right.len), left.len)
      runs[r + 1] = { start: left.start, len: left.len + right.len }
      runs.remove(r) |> ignore
    }
    continue end, end
  }
}

///|
fn[T : Compare] MutArrayView::insertion_sort(arr : MutArrayView[T]) -> Unit {
  for i in 1.. 0 && arr.unsafe_get(j) < arr.unsafe_get(j - 1); j = j - 1 {
      arr.swap(j, j - 1)
    }
  }
}

///|
/// Merges non-decreasing runs `arr[:mid]` and `arr[mid:]`. Copy `arr[mid:]` to buf and merge 
/// `buf` and `arr[:mid]` to `arr[:]`
fn[T : Compare] merge(arr : MutArrayView[T], mid : Int) -> Unit {
  let buf_len = arr.length() - mid
  let buf = UninitializedArray::make_and_blit(
    arr.buf(),
    allocate_len=buf_len,
    src_offset=arr.start() + mid,
    len=buf_len,
  )
  let buf_remaining = for p1 = mid - 1, p2 = buf_len - 1, p = mid + buf_len - 1; p1 >=
                         0 &&
                         p2 >= 0; {
    if arr.unsafe_get(p1) > buf[p2] {
      arr.unsafe_set(p, arr.unsafe_get(p1))
      continue p1 - 1, p2, p - 1
    } else {
      arr.unsafe_set(p, buf[p2])
      continue p1, p2 - 1, p - 1
    }
  } nobreak {
    p2
  }
  if buf_remaining >= 0 {
    UninitializedArray::unsafe_blit(
      arr.buf(),
      arr.start(),
      buf,
      0,
      buf_remaining + 1,
    )
  }
}

///|
/// Finds a streak of presorted elements starting at the beginning of the slice. Returns the first
/// value that is not part of said streak, and a bool denoting whether the streak was reversed.
/// Streaks can be increasing or decreasing.
fn[T : Compare] find_streak(arr : MutArrayView[T]) -> (Int, Bool) {
  let len = arr.length()
  if len < 2 {
    return (len, false)
  }
  let assume_reverse = arr.unsafe_get(1) < arr.unsafe_get(0)
  if assume_reverse {
    let end = for idx in 2..= arr.unsafe_get(idx - 1) {
        continue
      }
      break idx
    } nobreak {
      len
    }
    (end, false)
  }
}

///|
fn[T : Compare] provide_sorted_batch(
  arr : MutArrayView[T],
  start : Int,
  end : Int,
) -> Int {
  let len = arr.length()
  // This value is a balance between least comparisons and best performance, as
  // influenced by for example cache locality.
  let min_insertion_run = 10
  // Insert some more elements into the run if it's too short. Insertion sort is faster than
  // merge sort on short sequences, so this significantly improves performance.
  let start_end_diff = end - start
  if start_end_diff < min_insertion_run && end < len {
    // v[start_found:end] are elements that are already sorted in the input. We want to extend
    // the sorted region to the left, so we push up MIN_INSERTION_RUN - 1 to the right. Which is
    // more efficient that trying to push those already sorted elements to the left.
    let sort_end = minimum(len, start + min_insertion_run)
    MutArrayView::insertion_sort(arr.slice(start, sort_end))
    sort_end
  } else {
    end
  }
}

// TimSort is infamous for its buggy implementations, as described here:
// http://envisage-project.eu/timsort-specification-and-verification/
//
// This function correctly checks invariants for the top four runs. Additionally, if the top
// run starts at index 0, it will always demand a merge operation until the stack is fully
// collapsed, in order to complete the sort.

///|
fn collapse(runs : Array[TimSortRun], stop : Int) -> Int? {
  let n : Int = runs.length()
  if n >= 2 &&
    (
      runs[n - 1].start + runs[n - 1].len == stop ||
      runs[n - 2].len <= runs[n - 1].len ||
      (n >= 3 && runs[n - 3].len <= runs[n - 2].len + runs[n - 1].len) ||
      (n >= 4 && runs[n - 4].len <= runs[n - 3].len + runs[n - 2].len)
    ) {
    if n >= 3 && runs[n - 3].len < runs[n - 1].len {
      Some(n - 3)
    } else {
      Some(n - 2)
    }
  } else {
    None
  }
}

///|
/// Sorts the array
/// 
/// It's an in-place, unstable sort(it will reorder equal elements). The time complexity is O(n log n) in the worst case.
/// 
/// # Example
/// 
/// ```mbt check
/// test {
///   let arr = [5, 4, 3, 2, 1]
///   arr.sort()
///   @test.assert_eq(arr, [1, 2, 3, 4, 5])
/// }
/// ```
pub fn[T : Compare] FixedArray::sort(self : FixedArray[T]) -> Unit {
  self.mut_view().sort()
}

///|
fn[T] MutArrayView::slice(
  arr : MutArrayView[T],
  start : Int,
  end : Int,
) -> MutArrayView[T] {
  arr.mut_view(start~, end~)
}

///|
fn[T] MutArrayView::swap(arr : MutArrayView[T], i : Int, j : Int) -> Unit {
  // Callers are internal sort helpers that only pass in-bounds indices.
  let temp = arr.unsafe_get(i)
  arr.unsafe_set(i, arr.unsafe_get(j))
  arr.unsafe_set(j, temp)
}

///|
fn[T] MutArrayView::rev_in_place(arr : MutArrayView[T]) -> Unit {
  let len = arr.length()
  let mid_len = len / 2
  for i in 0.. Unit {
  let bubble_sort_len = 16
  for limit = limit, arr = arr, pred = pred, was_partitioned = true, balanced = true {
    let len = arr.length()
    if len <= bubble_sort_len {
      if len >= 2 {
        fixed_bubble_sort(arr)
      }
      return
    }
    // Too many imbalanced partitions may lead to O(n^2) performance in quick sort.
    // If the limit is reached, use heap sort to ensure O(n log n) performance.
    if limit == 0 {
      fixed_heap_sort(arr)
      return
    }
    let (pivot_index, likely_sorted) = fixed_choose_pivot(arr)
    // Try bubble sort if the array is likely already sorted.
    if was_partitioned && balanced && likely_sorted {
      if fixed_try_bubble_sort(arr) {
        return
      }
    }
    let (pivot, partitioned) = fixed_partition(arr, pivot_index)
    let was_partitioned = partitioned
    let balanced = minimum(pivot, len - pivot) >= len / 8
    let limit = if !balanced { limit - 1 } else { limit }
    if pred is Some(p) {
      // pred is less than all elements in arr
      // If pivot equals to pred, then we can skip all elements that are equal to pred.
      if p == arr.unsafe_get(pivot) {
        let i = for i = pivot; i < len && p == arr.unsafe_get(i); {
          continue i + 1
        } nobreak {
          i
        }
        continue limit, arr.slice(i, len), pred, was_partitioned, balanced
      }
    }
    let left = arr.slice(0, pivot)
    let right = arr.slice(pivot + 1, len)
    // Reduce the stack depth by only call fixed_quick_sort on the smaller fixed_partition.
    if left.length() < right.length() {
      fixed_quick_sort(left, pred, limit)
      continue limit,
        right,
        Some(arr.unsafe_get(pivot)),
        was_partitioned,
        balanced
    } else {
      fixed_quick_sort(right, Some(arr.unsafe_get(pivot)), limit)
      continue limit, left, pred, was_partitioned, balanced
    }
  }
}

///|
fn fixed_get_limit(len : Int) -> Int {
  for len = len, limit = 0; len > 0; {
    continue len / 2, limit + 1
  } nobreak {
    limit
  }
}

///|
/// Try to sort the array with bubble sort.
/// 
/// It will only tolerate at most 8 unsorted elements. The time complexity is O(n).
/// 
/// Returns whether the array is sorted.
fn[T : Compare] fixed_try_bubble_sort(arr : MutArrayView[T]) -> Bool {
  let max_tries = 8
  for i in 1.. 0 &&
                    arr.unsafe_get(j - 1) > arr.unsafe_get(j); {
      arr.swap(j, j - 1)
      continue j - 1, false
    } nobreak {
      sorted
    }
    if !sorted {
      let tries = tries + 1
      if tries > max_tries {
        break false
      }
      continue tries
    } else {
      continue tries
    }
  } nobreak {
    true
  }
}

///|
/// Try to sort the array with bubble sort.
/// 
/// It will only tolerate at most 8 unsorted elements. The time complexity is O(n).
/// 
/// Returns whether the array is sorted.
fn[T : Compare] fixed_bubble_sort(arr : MutArrayView[T]) -> Unit {
  for i in 1.. 0 && arr.unsafe_get(j - 1) > arr.unsafe_get(j); j = j - 1 {
      arr.swap(j, j - 1)
    }
  }
}

///|
test "fixed_try_bubble_sort" {
  let arr : FixedArray[_] = [8, 7, 6, 5, 4, 3, 2, 1]
  let sorted = fixed_try_bubble_sort(arr.mut_view())
  inspect(sorted, content="true")
  assert_true(arr == [1, 2, 3, 4, 5, 6, 7, 8])
}

///|
fn[T : Compare] fixed_partition(
  arr : MutArrayView[T],
  pivot_index : Int,
) -> (Int, Bool) {
  arr.swap(pivot_index, arr.length() - 1)
  let pivot = arr.unsafe_get(arr.length() - 1)
  let (i, partitioned) = for
    j in 0..<(arr.length() - 1)
    i = 0, partitioned = true {
    if arr.unsafe_get(j) < pivot {
      if i != j {
        arr.swap(i, j)
        continue i + 1, false
      } else {
        continue i + 1, partitioned
      }
    } else {
      continue i, partitioned
    }
  } nobreak {
    (i, partitioned)
  }
  arr.swap(i, arr.length() - 1)
  (i, partitioned)
}

///|
/// Choose a pivot index for quick sort.
/// 
/// It avoids worst case performance by choosing a pivot that is likely to be close to the median.
/// 
/// Returns the pivot index and whether the array is likely sorted.
fn[T : Compare] fixed_choose_pivot(arr : MutArrayView[T]) -> (Int, Bool) {
  let len = arr.length()
  let use_median_of_medians = 50
  let max_swaps = 4 * 3
  let mut swaps = 0
  let b = len / 4 * 2
  if len >= 8 {
    let a = len / 4 * 1
    let c = len / 4 * 3
    let sort_2 = (a : Int, b : Int) => {
      if arr.unsafe_get(a) > arr.unsafe_get(b) {
        arr.swap(a, b)
        swaps += 1
      }
    }
    let sort_3 = (a : Int, b : Int, c : Int) => {
      sort_2(a, b)
      sort_2(b, c)
      sort_2(a, b)
    }
    if len > use_median_of_medians {
      sort_3(a - 1, a, a + 1)
      sort_3(b - 1, b, b + 1)
      sort_3(c - 1, c, c + 1)
    }
    sort_3(a, b, c)
  }
  if swaps == max_swaps {
    arr.rev_in_place()
    (len - b - 1, true)
  } else {
    (b, swaps == 0)
  }
}

///|
fn[T : Compare] fixed_heap_sort(arr : MutArrayView[T]) -> Unit {
  let len = arr.length()
  for i in (len / 2)>..0 {
    fixed_sift_down(arr, i)
  }
  for i in len>..1 {
    arr.swap(0, i)
    fixed_sift_down(arr.slice(0, i), 0)
  }
}

///|
fn[T : Compare] fixed_sift_down(arr : MutArrayView[T], index : Int) -> Unit {
  let len = arr.length()
  for index = index, child = index * 2 + 1; child < len; {
    let child = if child + 1 < len &&
      arr.unsafe_get(child) < arr.unsafe_get(child + 1) {
      child + 1
    } else {
      child
    }
    if arr.unsafe_get(index) >= arr.unsafe_get(child) {
      return
    }
    arr.swap(index, child)
    continue child, child * 2 + 1
  }
}

///|
fn fixed_test_sort(f : (MutArrayView[Int]) -> Unit) -> Unit raise {
  let arr : FixedArray[_] = [5, 4, 3, 2, 1]
  f(arr.mut_view())
  assert_true(arr == [1, 2, 3, 4, 5])
  let arr : FixedArray[_] = [5, 5, 5, 5, 1]
  f(arr.mut_view())
  assert_true(arr == [1, 5, 5, 5, 5])
  let arr : FixedArray[_] = [1, 2, 3, 4, 5]
  f(arr.mut_view())
  assert_true(arr == [1, 2, 3, 4, 5])
  let arr = FixedArray::make(1000, 0)
  for i in 0..<1000 {
    arr[i] = 1000 - i - 1
  }
  for step in 1..<100 {
    let i = step * 10
    arr.swap(i, i - 1)
  }
  f(arr.mut_view())
  let expected = FixedArray::make(1000, 0)
  for i in 0..<1000 {
    expected[i] = i
  }
  assert_true(arr == expected)
}

///|
test "fixed_heap_sort" {
  fixed_test_sort(arr => fixed_heap_sort(arr.mut_view()))
}

///|
test "fixed_bubble_sort" {
  fixed_test_sort(arr => fixed_bubble_sort(arr.mut_view()))
}

///|
test "sort" {
  fixed_test_sort(arr => arr.sort())
}

///|
test "stable_sort" {
  let arr : FixedArray[_] = [5, 1, 3, 4, 2]
  arr.mut_view().stable_sort()
  assert_true(arr == [1, 2, 3, 4, 5])
  let arr = FixedArray::make(1000, 0)
  for i in 0..<1000 {
    arr[i] = 1000 - i - 1
  }
  for step in 1..<100 {
    let i = step * 10
    arr.swap(i, i - 1)
  }
  arr.mut_view().stable_sort()
  let expected = FixedArray::make(1000, 0)
  for i in 0..<1000 {
    expected[i] = i
  }
  assert_true(arr == expected)
}

///|
test "stable_sort_complex" {
  let run_lens = [86, 64, 21, 20, 22]
  let total_len = run_lens.fold(init=0, (acc, x) => acc + x)
  let arr = FixedArray::make(total_len, 0)
  for i in 0.. Int {
  if x > y {
    y
  } else {
    x
  }
}