use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
extern crate iai;
use hicollections::BTreeSet as MySet;
use hicollections::{AvlTree, AvlTreeNode, RbTree, RbTreeNode};
use hipool::PoolAlloc;
use rand;
use std::cmp::{Ordering, PartialEq, PartialOrd};
use std::collections::BTreeSet;
use std::mem::{MaybeUninit, size_of};
use std::ptr::{self, NonNull};
use std::slice;
use std::time::Instant;

#[global_allocator]
static ALLOCATOR: PoolAlloc = PoolAlloc;

#[repr(C)]
struct AvlFoo {
    val: usize,
    node: AvlTreeNode,
}

impl AvlFoo {
    const fn new() -> Self {
        Self {
            val: 0,
            node: AvlTreeNode::new(),
        }
    }
}

impl Clone for AvlFoo {
    fn clone(&self) -> Self {
        Self {
            val: self.val,
            node: AvlTreeNode::new(),
        }
    }
}

impl Clone for RbFoo {
    fn clone(&self) -> Self {
        Self {
            val: self.val,
            node: RbTreeNode::new(),
        }
    }
}

#[repr(C)]
struct RbFoo {
    val: usize,
    node: RbTreeNode,
}

impl RbFoo {
    const fn new() -> Self {
        Self {
            val: 0,
            node: RbTreeNode::new(),
        }
    }
}

impl PartialEq for RbFoo {
    fn eq(&self, other: &Self) -> bool {
        self.val == other.val
    }
}

impl PartialOrd for RbFoo {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.val.partial_cmp(&other.val)
    }
}

impl PartialEq for AvlFoo {
    fn eq(&self, other: &Self) -> bool {
        self.val == other.val
    }
}

impl PartialOrd for AvlFoo {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.val.partial_cmp(&other.val)
    }
}

impl Eq for AvlFoo {}
impl Eq for RbFoo {}
impl Ord for AvlFoo {
    fn cmp(&self, other: &Self) -> Ordering {
        self.val.cmp(&other.val)
    }
}
impl Ord for RbFoo {
    fn cmp(&self, other: &Self) -> Ordering {
        self.val.cmp(&other.val)
    }
}

#[repr(C)]
struct BTreeFoo {
    val: usize,
    _1: MaybeUninit<AvlTreeNode>,
}

impl BTreeFoo {
    const fn new() -> Self {
        Self {
            val: 0,
            _1: MaybeUninit::uninit(),
        }
    }
}

impl Clone for BTreeFoo {
    fn clone(&self) -> Self {
        Self {
            val: self.val,
            _1: MaybeUninit::uninit(),
        }
    }
}

impl PartialOrd for BTreeFoo {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.val.partial_cmp(&other.val)
    }
}

impl PartialEq for BTreeFoo {
    fn eq(&self, other: &Self) -> bool {
        self.val == other.val
    }
}

impl Ord for BTreeFoo {
    fn cmp(&self, other: &Self) -> Ordering {
        self.val.cmp(&other.val)
    }
}

impl Eq for BTreeFoo {}

static mut DATA: [usize; 1000000] = [0; 1000000];
static mut CNT: usize = 0;

fn get_slice<'a>(len: usize) -> &'a [usize]
{
    let data = unsafe { &mut DATA };
    let len = if len > data.len() { data.len() } else { len };
    unsafe {
        if CNT == 0 {
            for item in data.iter_mut() {
                    *item = rand::random();
            }
            CNT = len;
        }
        &data[..len]
    }
}

struct AvlCtx {
    foos: Vec<AvlFoo>,
    tree: Box<AvlTree<AvlFoo>>,
}

impl AvlCtx {
    fn new<F>(cnt: usize, f: F) -> Self
    where
        F: FnOnce(*const AvlFoo) -> *const AvlTreeNode,
    {
        let mut foos = vec![AvlFoo::new(); cnt];
        for (foo, n) in foos.iter_mut().zip(get_slice(cnt)) {
            foo.val = *n;
        }

        Self {
            foos,
            tree: Box::new(AvlTree::new(f)),
        }
    }
    fn new_ord<F>(cnt: usize, f: F) -> Self
    where
        F: FnOnce(*const AvlFoo) -> *const AvlTreeNode,
    {
        let mut foos = vec![AvlFoo::new(); cnt];
        for (foo, n) in foos.iter_mut().zip(get_slice(cnt)) {
            foo.val = *n;
        }
        foos.sort();

        Self {
            foos,
            tree: Box::new(AvlTree::new(f)),
        }
    }
}

struct RbCtx {
    foos: Vec<RbFoo>,
    tree: Box<RbTree<RbFoo>>,
}

impl RbCtx {
    fn new_ord<F>(cnt: usize, f: F) -> Self
    where
        F: FnOnce(*const RbFoo) -> *const RbTreeNode,
    {
        let mut foos = vec![RbFoo::new(); cnt];
        for (foo, n) in foos.iter_mut().zip(get_slice(cnt)) {
            foo.val = *n;
        }
        foos.sort();
        Self {
            foos,
            tree: Box::new(RbTree::new(f)),
        }
    }
    fn new<F>(cnt: usize, f: F) -> Self
    where
        F: FnOnce(*const RbFoo) -> *const RbTreeNode,
    {
        let mut foos = vec![RbFoo::new(); cnt];
        for (foo, n) in foos.iter_mut().zip(get_slice(cnt)) {
            foo.val = *n;
        }
        Self {
            foos,
            tree: Box::new(RbTree::new(f)),
        }
    }
}

fn static_foo(foo: &BTreeFoo) -> &'static BTreeFoo {
    let ptr = foo as *const BTreeFoo;
    unsafe { &*ptr }
}

struct BTreeCtx {
    foos: Vec<BTreeFoo>,
    tree: Box<BTreeSet<&'static BTreeFoo>>,
}

impl BTreeCtx {
    fn new_ord(cnt: usize) -> Self {
        let mut foos = vec![BTreeFoo::new(); cnt];
        for (foo, n) in foos.iter_mut().zip(get_slice(cnt)) {
            foo.val = *n;
        }
        foos.sort();
        Self {
            foos,
            tree: Box::new(BTreeSet::new()),
        }
    }
    fn new(cnt: usize) -> Self {
        let mut foos = vec![BTreeFoo::new(); cnt];
        for (foo, n) in foos.iter_mut().zip(get_slice(cnt)) {
            foo.val = *n;
        }
        Self {
            foos,
            tree: Box::new(BTreeSet::new()),
        }
    }
}

struct SetCtx {
    foos: Vec<BTreeFoo>,
    tree: Box<MySet<&'static BTreeFoo>>,
}

impl SetCtx {
    fn new_ord(cnt: usize) -> Self {
        let mut foos = vec![BTreeFoo::new(); cnt];
        for (foo, n) in foos.iter_mut().zip(get_slice(cnt)) {
            foo.val = *n;
        }
        foos.sort();
        Self {
            foos,
            tree: Box::new(MySet::new(0)),
        }
    }
    fn new(cnt: usize) -> Self {
        let mut foos = vec![BTreeFoo::new(); cnt];
        for (foo, n) in foos.iter_mut().zip(get_slice(cnt)) {
            foo.val = *n;
        }
        Self {
            foos,
            tree: Box::new(MySet::new(0)),
        }
    }
}

fn tree_rand_insert_benches(c: &mut Criterion) {
    let cnts = [1, 10, 100_usize, 1000, 10000, 100000];
    let mut group = c.benchmark_group("tree_insert_rand");
    for i in 0..cnts.len() {
        group.bench_function(BenchmarkId::new(format!("avl_insert_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("rb_insert_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("btree_insert_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> BTreeCtx {
                    let mut ctx = BTreeCtx::new(cnts[i]);
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });

        group.bench_function(BenchmarkId::new(format!("mytree_insert_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> SetCtx {
                    let mut ctx = SetCtx::new(cnts[i]);
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

fn tree_rand_find_benches(c: &mut Criterion) {
    let cnts = [1, 10, 50, 100_usize, 1000, 10000, 100000];
    let mut group = c.benchmark_group("tree_find_rand");
    for i in 0..cnts.len() {
        group.bench_function(BenchmarkId::new(format!("avl_find_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        assert!(ctx.tree.find(foo).is_some());
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("rb_find_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        assert!(ctx.tree.find(foo).is_some());
                        //ctx.tree.find(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("btree_find_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> BTreeCtx {
                    let mut ctx = BTreeCtx::new(cnts[i]);
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        assert!(ctx.tree.get(foo).is_some());
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("mytree_find_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> SetCtx {
                    let mut ctx = SetCtx::new(cnts[i]);
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        assert!(ctx.tree.get(foo).is_some());
                    });
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

fn tree_rand_remove_benches(c: &mut Criterion) {
    let cnts = [100_usize, 1000, 10000, 100000];
    let mut group = c.benchmark_group("tree_remove_rand");
    for i in 0..cnts.len() {
        group.bench_function(BenchmarkId::new(format!("avl_remove_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.remove(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("rb_remove_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.remove(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("btree_remove_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> BTreeCtx {
                    let mut ctx = BTreeCtx::new(cnts[i]);
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.remove(static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("mytree_remove_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> SetCtx {
                    let mut ctx = SetCtx::new(cnts[i]);
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.remove(static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

fn tree_ord_insert_benches(c: &mut Criterion) {
    let cnts = [100_usize, 1000, 10000, 100000];
    let mut group = c.benchmark_group("tree_insert_ord");
    for i in 0..cnts.len() {
        group.bench_function(BenchmarkId::new(format!("avl_insert_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx =
                        AvlCtx::new_ord(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("rb_insert_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx =
                        RbCtx::new_ord(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("btree_insert_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> BTreeCtx {
                    let mut ctx = BTreeCtx::new_ord(cnts[i]);
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("mytree_insert_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> SetCtx {
                    let mut ctx = SetCtx::new_ord(cnts[i]);
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

fn tree_ord_find_benches(c: &mut Criterion) {
    let cnts = [1, 10, 50, 100_usize, 1000, 10000, 100000];
    let mut group = c.benchmark_group("tree_find_ord");
    for i in 0..cnts.len() {
        group.bench_function(BenchmarkId::new(format!("avl_find_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx =
                        AvlCtx::new_ord(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.find(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("rb_find_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx =
                        RbCtx::new_ord(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.find(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("btree_find_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> BTreeCtx {
                    let mut ctx = BTreeCtx::new_ord(cnts[i]);
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.get(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("mytree_find_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> SetCtx {
                    let mut ctx = SetCtx::new_ord(cnts[i]);
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.get(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

fn tree_ord_remove_benches(c: &mut Criterion) {
    let cnts = [100_usize, 1000, 10000, 100000];
    let mut group = c.benchmark_group("tree_remove_ord");
    for i in 0..cnts.len() {
        group.bench_function(BenchmarkId::new(format!("avl_remove_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx =
                        AvlCtx::new_ord(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.remove(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("rb_remove_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx =
                        RbCtx::new_ord(cnts[i], |foo| unsafe { ptr::addr_of!((*foo).node) });
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.remove(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("btree_remove_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> BTreeCtx {
                    let mut ctx = BTreeCtx::new_ord(cnts[i]);
                    {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(static_foo(foo));
                    });
                    }
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.remove(foo);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new(format!("mytree_remove_{}", cnts[i]), i), |b| {
            b.iter_batched_ref(
                || -> SetCtx {
                    let mut ctx = SetCtx::new_ord(cnts[i]);
                    for foo in &ctx.foos {
                        ctx.tree.insert(static_foo(foo));
                    }
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.remove(static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

//iai::main!(test_init, test_avl_insert, test_avl_find, test_avl_remove);

criterion_group!(insert, tree_rand_insert_benches, tree_ord_insert_benches);
criterion_group!(find, tree_rand_find_benches, tree_ord_find_benches);
criterion_group!(remove, tree_rand_remove_benches, tree_ord_remove_benches);

criterion_main!(find, insert, remove);
//criterion_main!(find);
