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

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

impl Default for AvlFoo {
    fn default() -> Self {
        Self {
            val: rand::random(),
            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 Default for RbFoo {
    fn default() -> Self {
        Self {
            val: rand::random(),
            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)
    }
}

#[derive(Clone, Copy)]
struct BTreeFoo {
    val: usize,
    parent: Option<NonNull<BTreeFoo>>,
    left: Option<NonNull<BTreeFoo>>,
    right: Option<NonNull<BTreeFoo>>,
}

impl Default for BTreeFoo {
    fn default() -> Self {
        Self {
            val: 0,
            parent: None,
            left: None,
            right: None,
        }
    }
}

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 {}

fn static_foo(foo: &BTreeFoo) -> &'static BTreeFoo {
    unsafe { &* NonNull::from(foo).as_ptr() }
}

static mut DATA: Option<Vec<usize>> = None;

fn get_slice<'a>(len: usize) -> &'a mut [usize] {
    unsafe {
        if let Some(ref mut datas) = DATA {
            if datas.len() == len {
                return datas.as_mut_slice();
            }
        }
        DATA = Some(vec![0; len]);
        for item in DATA.as_mut().unwrap().as_mut_slice() {
            *item = rand::random();
        }
        return DATA.as_mut().unwrap().as_mut_slice();
    }
}

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

impl AvlCtx {
    fn new<F>(cnt: usize, f: F) -> Self
    where
        F: FnOnce(MaybeUninit<&AvlFoo>) -> MaybeUninit<&AvlTreeNode>,
    {
        let mut ret = Self {
            foos: vec![AvlFoo::default(); cnt],
            tree: Box::new(AvlTree::new(f)),
        };
        ret.foos.iter_mut().zip(get_slice(cnt)).for_each(|(foo, n)| foo.val = *n);
        ret
    }
    fn new_ord<F>(cnt: usize, f: F) -> Self
    where
        F: FnOnce(MaybeUninit<&AvlFoo>) -> MaybeUninit<&AvlTreeNode>,
    {
        let mut ret = Self {
            foos: vec![AvlFoo::default(); cnt],
            tree: Box::new(AvlTree::new(f)),
        };
        let slice = get_slice(cnt);
        slice.sort();
        ret.foos.iter_mut().zip(slice).for_each(|(foo, n)| foo.val = *n);
        ret
    }
}

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

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

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

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

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

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


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("avl_insert", i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().node)
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new("rb_insert", i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().node)
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        */
        group.bench_function(BenchmarkId::new("btree_insert", 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("my_btree_insert", 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("avl_find", i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().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("rb_find", i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().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("btree_find", 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.get(&static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new("my_btree_find", 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.get(&static_foo(foo));
                    });
                },
                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("avl_remove", i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().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("rb_remove", i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().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("btree_remove", 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("my_btree_remove", 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("avl_insert", i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new_ord(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().node)
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new("rb_insert", i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new_ord(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().node)
                    });
                    ctx
                },
                |ctx| {
                    ctx.foos.iter().for_each(|foo| {
                        ctx.tree.insert(foo, false);
                    });
                },
                BatchSize::SmallInput,
            );
        });
        */
        group.bench_function(BenchmarkId::new("btree_insert", 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("my_btree_insert", 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("avl_find", i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new_ord(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().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("rb_find", i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new_ord(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().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("btree_find", 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(&static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new("my_btree_find", 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(&static_foo(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("avl_remove", i), |b| {
            b.iter_batched_ref(
                || -> AvlCtx {
                    let mut ctx = AvlCtx::new_ord(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().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("rb_remove", i), |b| {
            b.iter_batched_ref(
                || -> RbCtx {
                    let mut ctx = RbCtx::new_ord(cnts[i], |foo| unsafe {
                        MaybeUninit::new(&foo.assume_init().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("btree_remove", 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(&static_foo(foo));
                    });
                },
                BatchSize::SmallInput,
            );
        });
        group.bench_function(BenchmarkId::new("my_btree_remove", 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.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);
