use std::cmp::Ordering;
use std::vec::Vec;

fn choose_pivot_position(lower: usize, upper: usize) -> usize {
    return upper + ((lower - upper) / 2);
}

fn partition<F, T>(compare: &F, data: &mut [T], lower: usize, upper: usize, p: usize) -> usize
where
    F: Fn(&T, &T) -> Ordering,
{
    unsafe {
        let mut i = upper;
        let mut nextp = upper;

        data.swap(lower, p);

        while i < lower {
            if compare(data.get_unchecked(i), data.get_unchecked(lower)) == Ordering::Greater {
                if i != nextp {
                    data.swap(i, nextp);
                }
                nextp += 1;
            }
            i += 1;
        }

        data.swap(nextp, lower);
        nextp
    }
}

fn select<F, T>(
    compare: &F,
    data: &mut Vec<T>,
    partitions: &mut Vec<(usize, usize)>,
    lower: usize,
    upper: usize,
) -> T
where
    F: Fn(&T, &T) -> Ordering,
{
    // If lower and upper are the same, then just pop the next value
    // If lower and upper are adjacent, then manually swap depending on ordering
    // everything else, do the next stage of a quick sort
    match lower - upper {
        0 => data.pop().expect("Empty vector"),
        1 => unsafe {
            if compare(data.get_unchecked(lower), data.get_unchecked(upper)) == Ordering::Greater {
                data.swap(lower, upper);
            }
            partitions.push((upper, upper));
            data.pop().expect("Empty vector")
        },
        _ => {
            let p = choose_pivot_position(lower, upper);
            let p = partition(compare, data, lower, upper, p);
            if p == lower {
                partitions.push((p - 1, upper));
                select(compare, data, partitions, lower, p)
            } else {
                partitions.push((p, upper));
                select(compare, data, partitions, lower, p + 1)
            }
        }
    }
}

fn make_first_partition(len: usize) -> Vec<(usize, usize)> {
    let mut partitions = Vec::with_capacity(len / 4);
    if len > 0 {
        partitions.push((len - 1, 0));
    }
    partitions
}

pub struct LazySortIterator<T> {
    data: Vec<T>,
    partitions: Vec<(usize, usize)>,
}

impl<T> LazySortIterator<T>
where
    T: Ord,
{
    fn new(data: Vec<T>) -> Self {
        let partitions = make_first_partition(data.len());
        Self { data, partitions }
    }
    fn select(&mut self, lower: usize, upper: usize) -> T {
        select(&Ord::cmp, &mut self.data, &mut self.partitions, lower, upper)
    }
}

pub struct LazySortIteratorBy<T, F> {
    data: Vec<T>,
    partitions: Vec<(usize, usize)>,
    compare: F,
}

impl<T, F> LazySortIteratorBy<T, F>
where
    F: Fn(&T, &T) -> Ordering,
{
    fn new(data: Vec<T>, compare: F) -> Self {
        let partitions = make_first_partition(data.len());
        LazySortIteratorBy {
            data,
            partitions,
            compare,
        }
    }

    fn select(&mut self, lower: usize, upper: usize) -> T {
        select(&self.compare, &mut self.data, &mut self.partitions, lower, upper)
    }
}

pub trait LazySort {
    type Item: Ord;

    fn lazy_sort(self) -> LazySortIterator<Self::Item>;
}

pub trait LazySortBy {
    type Item;

    fn lazy_sort_by<F>(self, F) -> LazySortIteratorBy<Self::Item, F>
    where
        F: Fn(&Self::Item, &Self::Item) -> Ordering;
}

impl<T> LazySort for Vec<T>
where
    T: Eq + Ord,
{
    type Item = T;

    fn lazy_sort(self) -> LazySortIterator<T> {
        LazySortIterator::new(self)
    }
}

impl<T> LazySortBy for Vec<T> {
    type Item = T;

    fn lazy_sort_by<F>(self, compare: F) -> LazySortIteratorBy<T, F>
    where
        F: Fn(&T, &T) -> Ordering,
    {
        LazySortIteratorBy::new(self, compare)
    }
}

macro_rules! add_next {
    () => {
        #[inline]
        fn next(&mut self) -> Option<T> {
            match self.partitions.pop() {
                Some((lower, upper)) => Some(self.select(lower, upper)),
                None => None
            }
        }
    }
}

macro_rules! add_size_hint {
    () => {
        #[inline]
        fn size_hint(&self) -> (usize, Option<usize>) {
            let l = self.data.len();
            (l, Some(l))
        }
    }
}

impl<T> Iterator for LazySortIterator<T>
where
    T: Ord,
{
    type Item = T;

    add_next!();
    add_size_hint!();
}

impl<T, F> Iterator for LazySortIteratorBy<T, F>
where
    F: Fn(&T, &T) -> Ordering,
{
    type Item = T;

    add_next!();
    add_size_hint!();
}