#![feature(allocator_api, extend_one, btreemap_alloc)]
//! [`BTreeSet`](std::collections::BTreeSet) but with duplicates.
use std::alloc::{Allocator, Global};
use std::collections::BTreeMap;
use std::collections::btree_map::{self, Entry};
use std::fmt::Debug;
use std::iter::FlatMap;

use smallvec::SmallVec;

#[derive(Clone)]
/// Like a btree set, but with duplicates.
/// Currently immutable cause i havent thought about it yet.
pub struct BTreeMultiset<
    T: ToK,
    const N: usize = 4,
    A: Allocator + Clone = Global,
> {
    inner: BTreeMap<T::K, SmallVec<[T; N]>, A>,
}
pub trait ToK {
    type K;
    fn to_key(&self) -> Self::K;
}
macro_rules! primitive {
    ($($t:ty)+) => { $(
        impl ToK for $t {
            type K = Self;

            fn to_key(&self) -> Self::K {
                *self
            }
        }
    )+ };
}
primitive!(i8 i16 i32 i64 i128 u8 u16 u32 u64 u128);
impl<T: Clone> ToK for Vec<T> {
    type K = Self;
    fn to_key(&self) -> Self::K {
        self.clone()
    }
}
impl<const N: usize, T: Debug + ToK<K: Ord>, A: Allocator + Clone> Debug
    for BTreeMultiset<T, N, A>
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_set().entries(self.iter()).finish()
    }
}
impl<const N: usize, T: ToK<K: Ord>, A: Allocator + Clone> IntoIterator
    for BTreeMultiset<T, N, A>
{
    type Item = T;

    type IntoIter = FlatMap<
        btree_map::IntoIter<<T as ToK>::K, SmallVec<[T; N]>, A>,
        SmallVec<[T; N]>,
        fn((<T as ToK>::K, SmallVec<[T; N]>)) -> SmallVec<[T; N]>,
    >;

    fn into_iter(self) -> Self::IntoIter {
        self.inner.into_iter().flat_map(|x| x.1)
    }
}
impl<const N: usize, T: ToK<K: Ord>> FromIterator<T>
    for BTreeMultiset<T, N>
{
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        let mut x = Self::default();
        x.extend(iter);
        x
    }
}
impl<const N: usize, T: ToK<K: Ord>, A: Allocator + Clone> Extend<T>
    for BTreeMultiset<T, N, A>
{
    #[inline]
    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
        iter.into_iter().for_each(move |x| {
            self.insert(x);
        });
    }

    #[inline]
    fn extend_one(&mut self, x: T) {
        self.insert(x);
    }
}
impl<const N: usize, T: ToK<K: Ord>> BTreeMultiset<T, N> {
    pub fn new() -> Self {
        Self { inner: BTreeMap::new() }
    }
}
impl<const N: usize, T: ToK<K: Ord>> Default for BTreeMultiset<T, N> {
    fn default() -> Self {
        Self::new()
    }
}
type Iter<'a, const N: usize, T> = FlatMap<
    btree_map::Iter<'a, <T as ToK>::K, SmallVec<[T; N]>>,
    &'a SmallVec<[T; N]>,
    fn((&'a <T as ToK>::K, &'a SmallVec<[T; N]>)) -> &'a SmallVec<[T; N]>,
>;
type Range<'a, const N: usize, T> = FlatMap<
    btree_map::Range<'a, <T as ToK>::K, SmallVec<[T; N]>>,
    &'a SmallVec<[T; N]>,
    fn((&'a <T as ToK>::K, &'a SmallVec<[T; N]>)) -> &'a SmallVec<[T; N]>,
>;
impl<const N: usize, T: ToK, A: Allocator + Clone> BTreeMultiset<T, N, A> {
    pub fn into_inner(self) -> BTreeMap<T::K, SmallVec<[T; N]>, A> {
        self.inner
    }
}
impl<const N: usize, T: ToK<K: Ord>, A: Allocator + Clone>
    BTreeMultiset<T, N, A>
{
    pub fn new_in(a: A) -> Self {
        Self { inner: BTreeMap::new_in(a) }
    }
    /// Gets an iterator that visits the elements in the [`BTreeMultiset`] in ascending order.
    /// Similar to `range(..)`.
    pub fn iter(&'_ self) -> Iter<'_, N, T> {
        self.inner.iter().flat_map(|x| x.1)
    }
    /// Adds a value to the multi set.
    pub fn insert(&mut self, t: T) -> bool {
        match self.inner.entry(t.to_key()) {
            Entry::Vacant(vacant_entry) => {
                vacant_entry.insert(SmallVec::from_iter([t]));
                false
            }
            Entry::Occupied(occupied_entry) => {
                occupied_entry.into_mut().push(t);
                true
            }
        }
    }

    /// Constructs a double-ended iterator over a sub-range of elements in the multiset.
    /// The simplest way is to use the range syntax `min..max`, thus `range(min..max)` will
    /// yield elements from min (inclusive) to max (exclusive).
    /// The range may also be entered as `(Bound<T>, Bound<T>)`, so for example
    /// `range((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
    /// range from 4 to 10.
    ///
    /// # Panics
    ///
    /// Panics if range `start > end`.
    /// Panics if range `start == end` and both bounds are `Excluded`.
    pub fn range(
        &'_ self,
        r: impl std::ops::RangeBounds<T::K>,
    ) -> Range<'_, N, T> {
        self.inner.range(r).flat_map(|x| x.1)
    }

    // pub fn range_mut<'a, R>(
    //     &'a mut self,
    //     r: impl std::ops::RangeBounds<K>,
    // ) -> FlatMap<
    //     RangeMut<'a, K, SmallVec<T>>,
    //     &'a mut SmallVec<T>,
    //     fn((&'a K, &'a mut SmallVec<T>)) -> &'a mut SmallVec<T>,
    // > {
    //     self.inner.range_mut(r).flat_map(|x| x.1)
    // }
}
#[test]
fn works() {
    let mut x: BTreeMultiset<i32> = BTreeMultiset::new();

    x.insert(4);
    x.insert(3);
    x.insert(7);
    x.insert(3);
    assert_eq!(x.iter().copied().collect::<Vec<_>>(), vec![3, 3, 4, 7]);
    assert_eq!(x.range(0..=4).copied().collect::<Vec<_>>(), vec![3, 3, 4]);
}