use std::cmp::{Ordering, Reverse};
use std::collections::BinaryHeap;

#[derive(Clone, Debug)]
#[repr(transparent)]
pub struct Arbitrary<T>(pub T);

impl<T> PartialEq for Arbitrary<T> {
    #[inline(always)]
    fn eq(&self, _other: &Self) -> bool {
        true
    }
}

impl<T> Eq for Arbitrary<T> {}

impl<T> PartialOrd for Arbitrary<T> {
    #[inline(always)]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<T> Ord for Arbitrary<T> {
    #[inline(always)]
    fn cmp(&self, _other: &Self) -> Ordering {
        Ordering::Equal
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DynamicOrd<T: Ord> {
    inner: T,
    reverse: bool,
}

impl<T: Ord> DynamicOrd<T> {
    pub fn new(inner: T, reverse: bool) -> Self {
        Self { inner, reverse }
    }
}

impl<T: Ord> PartialOrd for DynamicOrd<T> {
    #[inline(always)]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<T: Ord> Ord for DynamicOrd<T> {
    #[inline(always)]
    fn cmp(&self, other: &Self) -> Ordering {
        debug_assert!(self.reverse == other.reverse);

        let ordering = self.inner.cmp(&other.inner);

        if self.reverse {
            ordering.reverse()
        } else {
            ordering
        }
    }
}

// A specialized heap handy to compute top-k in O(log n) time
// but only O(k) memory.
// It is a max-heap by default to fit rust's standard library
// choices.
#[derive(Debug, Clone)]
pub struct TopKHeap<T> {
    capacity: usize,
    heap: BinaryHeap<Reverse<T>>,
}

impl<T: Ord> TopKHeap<T> {
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            capacity,
            heap: BinaryHeap::with_capacity(capacity),
        }
    }

    pub fn push(&mut self, item: T) -> bool {
        let heap = &mut self.heap;

        if heap.len() < self.capacity {
            heap.push(Reverse(item));

            return true;
        } else {
            let worst_item = heap.peek().unwrap();

            if item > worst_item.0 {
                heap.pop();
                heap.push(Reverse(item));
                return true;
            }
        }

        false
    }

    pub fn into_sorted_vec(mut self) -> Vec<T> {
        let l = self.heap.len();

        let mut items = Vec::with_capacity(l);
        let uninit = items.spare_capacity_mut();

        let mut i: usize = l;

        while let Some(Reverse(item)) = self.heap.pop() {
            i -= 1;
            uninit[i].write(item);
        }

        unsafe {
            items.set_len(l);
        }

        items
    }
}

impl<T: Ord> Extend<T> for TopKHeap<T> {
    // Required method
    fn extend<I>(&mut self, iter: I)
    where
        I: IntoIterator<Item = T>,
    {
        for item in iter {
            self.push(item);
        }
    }
}

#[derive(Clone, Debug)]
pub struct TopKHeapMap<T, V> {
    capacity: usize,
    heap: BinaryHeap<(Reverse<T>, Arbitrary<V>)>,
}

impl<T: Ord, V> TopKHeapMap<T, V> {
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            capacity,
            heap: BinaryHeap::with_capacity(capacity),
        }
    }

    pub fn capacity(&self) -> usize {
        self.capacity
    }

    pub fn clear(&mut self) {
        self.heap.clear();
    }

    pub fn into_unordered_iter(self) -> impl Iterator<Item = (T, V)> {
        self.heap
            .into_iter()
            .map(|(Reverse(k), Arbitrary(v))| (k, v))
    }

    pub fn push_with<F>(&mut self, item: T, callback: F) -> bool
    where
        F: FnOnce() -> V,
    {
        let heap = &mut self.heap;

        if heap.len() < self.capacity {
            heap.push((Reverse(item), Arbitrary(callback())));

            return true;
        } else {
            let worst_item = heap.peek().unwrap();

            if item > worst_item.0.0 {
                heap.pop();
                heap.push((Reverse(item), Arbitrary(callback())));
                return true;
            }
        }

        false
    }

    pub fn into_sorted_vec(mut self) -> Vec<(T, V)> {
        let l = self.heap.len();

        let mut items = Vec::with_capacity(l);
        let uninit = items.spare_capacity_mut();

        let mut i: usize = l;

        while let Some((Reverse(item), Arbitrary(value))) = self.heap.pop() {
            i -= 1;
            uninit[i].write((item, value));
        }

        unsafe {
            items.set_len(l);
        }

        items
    }
}

impl<T: Ord + Clone, V: Clone> TopKHeapMap<T, V> {
    pub fn to_sorted_vec(&self) -> Vec<(T, V)> {
        self.clone().into_sorted_vec()
    }
}

#[derive(Clone, Debug)]
pub struct TopKHeapMapWithTies<T, V> {
    capacity: usize,
    heap: BinaryHeap<(Reverse<T>, Arbitrary<V>)>,
    ties: Option<Vec<(T, V)>>,
}

impl<T: Ord, V> TopKHeapMapWithTies<T, V> {
    pub fn with_capacity(capacity: usize, ties: bool) -> Self {
        Self {
            capacity,
            heap: BinaryHeap::with_capacity(capacity),
            ties: ties.then(Vec::new),
        }
    }

    pub fn len(&self) -> usize {
        self.heap.len() + self.ties.as_ref().map(|ties| ties.len()).unwrap_or(0)
    }

    pub fn push_with<F>(&mut self, item: T, callback: F) -> bool
    where
        F: FnOnce() -> V,
    {
        let heap = &mut self.heap;

        if heap.len() < self.capacity {
            heap.push((Reverse(item), Arbitrary(callback())));

            return true;
        } else {
            let worst_item = heap.peek().unwrap();

            match item.cmp(&worst_item.0.0) {
                Ordering::Greater => {
                    let mut worst_item_popped = heap.pop().unwrap();

                    heap.push((Reverse(item), Arbitrary(callback())));

                    if let Some(ties) = &mut self.ties {
                        ties.clear();
                    }

                    if worst_item_popped.0.0 == heap.peek().unwrap().0.0 {
                        std::mem::swap(&mut heap.peek_mut().unwrap().1, &mut worst_item_popped.1);

                        if let Some(ties) = &mut self.ties {
                            ties.push((worst_item_popped.0.0, worst_item_popped.1.0));
                        }
                    }

                    return true;
                }
                Ordering::Equal => {
                    if let Some(ties) = &mut self.ties {
                        ties.push((item, callback()));
                        return true;
                    }

                    return false;
                }
                _ => (),
            };
        }

        false
    }

    #[inline(always)]
    pub fn pop(&mut self) -> Option<(T, V)> {
        self.heap
            .pop()
            .map(|(Reverse(item), Arbitrary(value))| (item, value))
    }

    pub fn into_sorted_vec(mut self) -> Vec<(T, V)> {
        let l = self.len();
        let hl = self.heap.len();

        let mut items = Vec::with_capacity(l);
        let uninit = items.spare_capacity_mut();

        let mut i: usize = hl;

        while let Some((Reverse(item), Arbitrary(value))) = self.heap.pop() {
            i -= 1;
            uninit[i].write((item, value));
        }

        i = hl;

        if let Some(ties) = self.ties {
            for pair in ties {
                uninit[i].write(pair);
                i += 1;
            }
        }

        unsafe {
            items.set_len(l);
        }

        items
    }

    pub fn merge(&mut self, other: Self) {
        debug_assert!(
            self.capacity == other.capacity && self.ties.is_some() == other.ties.is_some()
        );

        for (Reverse(k), Arbitrary(v)) in other.heap.into_iter() {
            self.push_with(k, || v);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_numbers() {
        let mut heap = TopKHeap::with_capacity(3);
        heap.extend([1, 2, 3, 4, 5, 6]);

        assert_eq!(heap.into_sorted_vec(), vec![6, 5, 4]);
    }

    #[test]
    fn test_reverse_numbers() {
        let mut heap = TopKHeap::with_capacity(3);
        heap.extend([1, 2, 3, 4, 5, 6].into_iter().map(Reverse));

        assert_eq!(
            heap.into_sorted_vec()
                .iter()
                .map(|n| n.0)
                .collect::<Vec<_>>(),
            vec![1, 2, 3]
        );
    }

    #[test]
    fn test_map() {
        let mut heap = TopKHeapMap::with_capacity(2);
        heap.push_with(1, || "one");
        heap.push_with(2, || "two");
        heap.push_with(3, || "three");

        assert_eq!(
            heap.clone().into_sorted_vec(),
            vec![(3, "three"), (2, "two")]
        );
    }

    #[test]
    fn test_map_with_ties() {
        let mut heap = TopKHeapMapWithTies::with_capacity(2, true);
        heap.push_with(1, || "one");
        heap.push_with(2, || "two");
        heap.push_with(3, || "three");

        assert_eq!(
            heap.clone().into_sorted_vec(),
            vec![(3, "three"), (2, "two")]
        );

        // Final ties
        let mut heap = TopKHeapMapWithTies::with_capacity(2, true);
        heap.push_with(1, || "one");
        heap.push_with(2, || "two");
        heap.push_with(3, || "three");
        heap.push_with(2, || "four");
        heap.push_with(2, || "five");

        assert_eq!(
            heap.clone().into_sorted_vec(),
            vec![(3, "three"), (2, "two"), (2, "four"), (2, "five")]
        );
    }

    #[test]
    fn test_map_with_ties_special_order() {
        let mut heap = TopKHeapMapWithTies::with_capacity(2, true);
        heap.push_with(1, || "one");
        heap.push_with(2, || "two");
        heap.push_with(2, || "three");
        heap.push_with(3, || "four");

        assert_eq!(
            heap.clone().into_sorted_vec(),
            vec![(3, "four"), (2, "two"), (2, "three")]
        );
    }
}
