use ::core::num::Wrapping;
use ::mice::backend_support::TyTarget;
use ::rand::{Rng, RngCore};
use criterion::{black_box, criterion_group, criterion_main, BenchmarkGroup, Criterion};

/// Dice expressions we're using for benchmarks.
const DICE_EXPRESSIONS: &[(&str, &str, fn(&mut dyn RngCore) -> i64)] = &[
    ("Small Expression", "2d6", |rng| {
        (0..2).fold(Wrapping(0i64), |a, _| a + Wrapping(rng.gen_range(1..=6))).0
    }),
    ("Medium Expression", "2d6 + 4d4 + 9", |rng| {
        let die_2d6 = (0..2).fold(Wrapping(0i64), |a, _| a + Wrapping(rng.gen_range(1..=6)));
        let die_4d4 = (0..4).fold(Wrapping(0i64), |a, _| a + Wrapping(rng.gen_range(1..=4)));
        let nine = Wrapping(9);
        (die_2d6 + die_4d4 + nine).0
    }),
    ("Large Expression", "3d9 + 9d4 - 2d1 + 40d7", |rng| {
        let die_3d9 = (0..3).fold(Wrapping(0i64), |a, _| a + Wrapping(rng.gen_range(1..=9)));
        let die_9d4 = (0..9).fold(Wrapping(0i64), |a, _| a + Wrapping(rng.gen_range(1..=4)));
        let die_2d1 = Wrapping(2);
        let die_40d7 = (0..40).fold(Wrapping(0), |a, _| a + Wrapping(rng.gen_range(1..=7)));
        (die_3d9 + die_9d4 - die_2d1 + die_40d7).0
    }),
    ("Absurdly Large Expression", "3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7 + 3d9 + 9d4 - 2d1 + 40d7", |_| 0),
    ("Expensive But Short Expression", "10000d100", |_| 0)
];

fn roll_startup_benchmark(c: &mut Criterion) {
    let rollers = |mut group: BenchmarkGroup<_>, program_text: &str| {
        group.bench_function("New AST Parser", |b| {
            b.iter(|| {
                let _ = black_box(
                    ::mice::parse::parse_expression::<TyTarget::AstInterp>(program_text.as_bytes())
                        .unwrap(),
                );
            });
        });
        group.bench_function("New Stack Machine Compiler", |b| {
            b.iter(|| {
                let _ = black_box({
                    let (_, (_, ast)) =
                        ::mice::parse::parse_expression::<TyTarget::Stack>(program_text.as_bytes())
                            .unwrap();
                    ::mice::stack::compile(&ast)
                });
            });
        });
        group.bench_function("MIR Stack Compiler", |b| {
            b.iter(|| {
                let (_, (_, ast)) = ::mice::parse::parse_expression::<TyTarget::Mir::Stack>(
                    program_text.as_bytes(),
                )
                .unwrap();
                let mut mir = ::mice::mir::lower(&ast).unwrap();
                ::mice::mir::opt::destroy_metadata(&mut mir);
                ::mice::mir::stack::lower(&mir)
            });
        });
    };
    for (title, expression, _rust) in DICE_EXPRESSIONS {
        rollers(
            c.benchmark_group(format!("Startup - {}", title)),
            black_box(*expression),
        );
    }
}

fn rolling_benchmark(c: &mut Criterion) {
    use ::rand::SeedableRng;
    let mut rng = ::rand::rngs::SmallRng::from_entropy();

    let mut rollers = |mut group: BenchmarkGroup<_>,
                       program_text: &str,
                       rust: fn(&mut dyn RngCore) -> i64| {
        let (_, (_, program)) = black_box(
            mice::parse::parse_expression::<TyTarget::AstInterp>(program_text.as_bytes()).unwrap(),
        );
        let stack_program = black_box(mice::stack::compile(&program));
        let mut mir = black_box(::mice::mir::lower(&program).unwrap());
        ::mice::mir::opt::destroy_metadata(&mut mir);
        let mir_stack_program = black_box(::mice::mir::stack::lower(&mir));
        group.bench_function("New AST Roller", |b| {
            b.iter(|| black_box(mice::interp::interpret(&mut rng, &program)));
        });
        group.bench_function("New Stack Machine Roller", |b| {
            b.iter(|| {
                let mut machine = ::mice::stack::Machine::new();
                let _ = black_box(machine.eval_with(&mut rng, &stack_program));
            });
        });
        group.bench_function("MIR Stack VM", |b| {
            b.iter(|| {
                let mut machine = ::mice::mir::stack::Machine::new(&mir_stack_program, 0);
                let _ = black_box(machine.interpret(&mut rng));
                let _ = black_box(machine);
            });
        });
        group.bench_function("Rust", |b| {
            b.iter(|| {
                let _ = black_box(rust(&mut rng));
            });
        });
    };

    let mut wg = |gn, t, rust| rollers(c.benchmark_group(gn), t, rust);
    for (title, expression, rust) in DICE_EXPRESSIONS {
        wg(format!("Execution - {}", title), expression, *rust);
    }
}

fn end_to_end_rolling_benchmark(c: &mut Criterion) {
    use ::rand::SeedableRng;
    let mut rng = ::rand::rngs::SmallRng::from_entropy();
    let mut rollers = |mut group: BenchmarkGroup<_>, program_text: &str| {
        group.bench_function("New AST Roller", |b| {
            b.iter(|| {
                let (_, (_, program)) =
                    mice::parse::parse_expression::<TyTarget::AstInterp>(program_text.as_bytes())
                        .unwrap();
                black_box(mice::interp::interpret(&mut rng, &program))
            });
        });
        group.bench_function("New Stack Machine Roller", |b| {
            b.iter(|| {
                let (_, (_, program)) =
                    mice::parse::parse_expression::<TyTarget::Stack>(program_text.as_bytes())
                        .unwrap();
                let stack_program = mice::stack::compile(&program);
                let mut machine = ::mice::stack::Machine::new();
                let _ = black_box(machine.eval_with(&mut rng, &stack_program));
            });
        });
        group.bench_function("MIR Stack VM", |b| {
            b.iter(|| {
                let (_, (_, program)) =
                    mice::parse::parse_expression::<TyTarget::Mir::Stack>(program_text.as_bytes())
                        .unwrap();
                let mut mir = ::mice::mir::lower(&program).unwrap();
                ::mice::mir::opt::destroy_metadata(&mut mir);
                let stack_program = ::mice::mir::stack::lower(&mir);
                let mut machine = ::mice::mir::stack::Machine::new(&stack_program, 0);
                let _ = black_box(machine.interpret(&mut rng));
                let _ = black_box(machine);
            });
        });
    };

    for (title, expression, _rust) in DICE_EXPRESSIONS {
        rollers(
            c.benchmark_group(format!("End to End - {}", title)),
            black_box(expression),
        );
    }
}

criterion_group!(startup_benches, roll_startup_benchmark);
criterion_group!(benches, rolling_benchmark);
criterion_group!(end_to_end_benches, end_to_end_rolling_benchmark);
criterion_main!(startup_benches, benches, end_to_end_benches);