use criterion::{black_box, criterion_group, criterion_main, Criterion};

use receme::{
    catana::{Ana, Cata},
    functor::Functor,
    hylo::Hylo,
    tree::{DFTree, TEStrategy, Tree, TreeIndex},
};

#[derive(Debug, Clone)]
enum Expr<T> {
    Add(T, T),
    Lit(isize),
}

impl<T> Functor<T> for Expr<T> {
    type Target<S> = Expr<S>;

    fn fmap<S>(self, mut f: impl FnMut(T) -> S) -> Self::Target<S> {
        match self {
            Expr::Add(a, b) => Expr::Add(f(a), f(b)),
            Expr::Lit(value) => Expr::Lit(value),
        }
    }
}

#[derive(Debug, Clone)]
enum ExBox {
    Add(Box<ExBox>, Box<ExBox>),
    Lit(isize),
}

impl ExBox {
    fn lit(value: isize) -> Self {
        Self::Lit(value)
    }

    fn add(a: ExBox, b: ExBox) -> Self {
        Self::Add(Box::new(a), Box::new(b))
    }
}

pub fn bench(c: &mut Criterion) {
    fn construct_tree() -> Tree<Expr<TreeIndex>> {
        let strategy: TEStrategy = TEStrategy::UnsafeArena;

        let elements: Vec<Expr<TreeIndex>> = vec![
            Expr::Add(1, 2).fmap(TreeIndex::new),
            Expr::Lit(1),
            Expr::Add(3, 4).fmap(TreeIndex::new),
            Expr::Lit(3),
            Expr::Add(5, 6).fmap(TreeIndex::new),
            Expr::Lit(10),
            Expr::Lit(-14),
        ];

        Tree::new(elements, strategy)
    }

    fn construct_box_tree() -> Box<ExBox> {
        Box::new(ExBox::add(
            ExBox::lit(1),
            ExBox::add(ExBox::lit(3), ExBox::add(ExBox::lit(10), ExBox::lit(-14))),
        ))
    }

    let tree = construct_tree();

    let safe_tree = {
        let mut tree = tree.clone();
        tree.set_strategy(Default::default());
        tree
    };

    let depth_first_tree = {
        let mut tree = tree.clone();
        tree.set_strategy(TEStrategy::DepthFirst);
        tree
    };

    let boxtree = construct_box_tree();

    c.bench_function("bench cata", |b| {
        b.iter(|| {
            let result = tree.clone().cata(|expr: Expr<isize>| match expr {
                Expr::Add(a, b) => a + b,
                Expr::Lit(v) => v,
            });

            black_box(result);
        })
    });

    c.bench_function("bench cata safe", |b| {
        b.iter(|| {
            let result = safe_tree.clone().cata(|expr: Expr<isize>| match expr {
                Expr::Add(a, b) => a + b,
                Expr::Lit(v) => v,
            });

            black_box(result);
        })
    });

    c.bench_function("bench cata depth first", |b| {
        b.iter(|| {
            let result = depth_first_tree
                .clone()
                .cata(|expr: Expr<isize>| match expr {
                    Expr::Add(a, b) => a + b,
                    Expr::Lit(v) => v,
                });

            black_box(result);
        })
    });

    c.bench_function("bench cata loop", |b| {
        b.iter(|| {
            let tree = tree.clone();

            let mut stack: Vec<Result<usize, usize>> = vec![Ok(0)];
            let mut result_stack: Vec<isize> = vec![];

            while let Some(top) = stack.pop() {
                let top_is_ok_p = top.is_ok();

                let top = match top {
                    Ok(top) => top,
                    Err(top) => top,
                };

                let top_node = tree.nth(top).unwrap();

                match top_node {
                    Expr::Add(a, b) => {
                        if top_is_ok_p {
                            stack.push(Err(top));
                            stack.push(Ok(**b));
                            stack.push(Ok(**a));
                        } else {
                            let a = result_stack.pop().unwrap();
                            let b = result_stack.pop().unwrap();

                            result_stack.push(a + b);
                        }
                    }
                    Expr::Lit(v) => result_stack.push(*v),
                }
            }

            let result = result_stack.pop().unwrap();

            black_box(result);
        })
    });

    c.bench_function("bench ana box loop", |b| {
        b.iter(|| {
            let boxtree = boxtree.clone();

            let mut elements = vec![];

            let mut stack = vec![boxtree];

            while let Some(bt) = stack.pop() {
                match *bt {
                    ExBox::Add(a, b) => {
                        let len = elements.len();

                        elements.push(Expr::Add(TreeIndex::new(len + 1), TreeIndex::new(len + 2)));

                        stack.push(b);
                        stack.push(a);
                    }
                    ExBox::Lit(v) => elements.push(Expr::Lit(v)),
                }
            }

            let tree: Tree<Expr<TreeIndex>> = Tree::new(elements, Default::default());

            black_box(tree);
        })
    });

    c.bench_function("bench ana boxed", |b| {
        b.iter(|| {
            let boxtree = boxtree.clone();

            let tree = Tree::ana(boxtree, |bt: Box<ExBox>| match *bt {
                ExBox::Add(a, b) => Expr::Add(a, b),
                ExBox::Lit(v) => Expr::Lit(v),
            });

            black_box(tree);
        })
    });

    c.bench_function("bench ana depth first boxed", |b| {
        b.iter(|| {
            let boxtree = boxtree.clone();

            let tree = DFTree::ana(boxtree, |bt: Box<ExBox>| match *bt {
                ExBox::Add(a, b) => Expr::Add(a, b),
                ExBox::Lit(v) => Expr::Lit(v),
            });

            black_box(tree);
        })
    });

    c.bench_function("bench hylo", |b| {
        b.iter(|| {
            let boxtree = boxtree.clone();

            black_box(Tree::hylo(
                boxtree,
                |expr| match expr {
                    Expr::Add(a, b) => a + b,
                    Expr::Lit(v) => v,
                },
                |bt: Box<ExBox>| match *bt {
                    ExBox::Add(a, b) => Expr::Add(a, b),
                    ExBox::Lit(v) => Expr::Lit(v),
                },
            ))
        })
    });

    c.bench_function("bench hylo loop", |b| {
        b.iter(|| {
            let boxtree = boxtree.clone();

            let mut elements = vec![];

            let mut stack = vec![boxtree];

            while let Some(bt) = stack.pop() {
                match *bt {
                    ExBox::Add(a, b) => {
                        let len = elements.len();

                        elements.push(Expr::Add(TreeIndex::new(len + 1), TreeIndex::new(len + 2)));

                        stack.push(b);
                        stack.push(a);
                    }
                    ExBox::Lit(v) => elements.push(Expr::Lit(v)),
                }
            }

            let tree: Tree<Expr<TreeIndex>> = Tree::new(elements, Default::default());

            let mut stack: Vec<Result<usize, usize>> = vec![Ok(0)];
            let mut result_stack: Vec<isize> = vec![];

            while let Some(top) = stack.pop() {
                let top_is_ok_p = top.is_ok();
                let top = match top {
                    Ok(top) => top,
                    Err(top) => top,
                };

                let top_node = tree.nth(top).unwrap();

                match top_node {
                    Expr::Add(a, b) => {
                        if top_is_ok_p {
                            stack.push(Err(top));
                            stack.push(Ok(**b));
                            stack.push(Ok(**a));
                        } else {
                            let a = result_stack.pop().unwrap();
                            let b = result_stack.pop().unwrap();

                            result_stack.push(a + b);
                        }
                    }
                    Expr::Lit(v) => result_stack.push(*v),
                }
            }

            let result = result_stack.pop().unwrap();

            assert_eq!(result, 0);

            black_box(result);
        })
    });
}

criterion_group!(benches, bench);
criterion_main!(benches);