use crate::utils::{gt_by, lt_by};
use std::cmp::Ordering;

/// Sorts a slice of data using the [quick sort](https://en.wikipedia.org/wiki/Quicksort) in `O(n×log(n))`.
///
/// # Examples
///
/// ```
/// # use sort_steps::{quick_sort};
/// let numbers = [5, 9, 3, 6, 8, 2, 1, 7, 4];
/// println!("Quick Sort Steps:");
/// for (i, v) in quick_sort(&numbers).enumerate() {
///     println!("#{:02}: {:?}", i, v);
/// }
/// ```
pub fn quick_sort<T>(v: &[T]) -> impl Iterator<Item = Vec<T>>
where
    T: PartialOrd + Clone,
{
    quick_sort_by(v, |a, b| a.partial_cmp(b))
}

/// Sorts a slice of data using the [quick sort](https://en.wikipedia.org/wiki/Quicksort) with custom comparator.
///
/// # Examples
///
/// ```
/// # use sort_steps::{quick_sort_by};
/// let numbers = [5, 9, 3, 6, 8, 2, 1, 7, 4];
/// println!("Quick Sort Steps:");
/// for (i, v) in quick_sort_by(&numbers, |a, b| a.partial_cmp(b)).enumerate() {
///     println!("#{:02}: {:?}", i, v);
/// }
/// ```
pub fn quick_sort_by<T, F>(v: &[T], compare: F) -> impl Iterator<Item = Vec<T>>
where
    F: Fn(&T, &T) -> Option<Ordering> + Copy,
    T: Clone,
{
    let mut v = v.to_vec();
    let mut result = vec![v.to_vec()];
    let high = v.len() - 1;
    quick_sort_helper(&mut v, 0, high, compare, &mut result);
    result.into_iter()
}

fn quick_sort_helper<T, F>(v: &mut [T], low: usize, high: usize, compare: F, steps: &mut Vec<Vec<T>>)
where
    F: Fn(&T, &T) -> Option<Ordering> + Copy,
    T: Clone,
{
    if low < high {
        let pivot_index = partition(v, low, high, compare);
        steps.push(v.to_vec());
        quick_sort_helper(v, low, pivot_index, compare, steps);
        quick_sort_helper(v, pivot_index + 1, high, compare, steps);
    }
}

fn partition<T, F>(v: &mut [T], low: usize, high: usize, compare: F) -> usize
where
    F: Fn(&T, &T) -> Option<Ordering> + Copy,
    T: Clone,
{
    let pivot = v[low].clone();
    let mut i = low;
    let mut j = high;

    loop {
        while lt_by(&v[i], &pivot, compare) {
            i += 1;
        }
        while gt_by(&v[j], &pivot, compare) {
            j -= 1;
        }
        if i >= j {
            return j;
        }
        v.swap(i, j);
    }
}