use criterion::{criterion_group, criterion_main, Criterion};
use std::collections::BTreeSet;

type Pos = (usize, usize);

#[derive(Clone, Hash, PartialEq, Eq)]
struct State {
    pos: Pos,
    visited: BTreeSet<Pos>,
}

fn successors(state: &State, x_max: usize, y_max: usize) -> Vec<State> {
    {
        let mut states = vec![];
        if state.pos.0 > 0 {
            states.push((state.pos.0 - 1, state.pos.1));
        }
        if state.pos.0 < x_max {
            states.push((state.pos.0 + 1, state.pos.1));
        }
        if state.pos.1 > 0 {
            states.push((state.pos.0, state.pos.1 - 1));
        }
        if state.pos.1 < y_max {
            states.push((state.pos.0, state.pos.1 + 1));
        }
        states
    }
    .into_iter()
    .map(|s| State {
        pos: s,
        visited: {
            let mut visited = state.visited.to_owned();
            visited.insert(s);
            visited
        },
    })
    .collect()
}

fn heuristic(state: &State, x_max: usize, y_max: usize) -> usize {
    x_max * y_max - state.visited.len()
}

fn bfs(init: &State, x_max: usize, y_max: usize) -> Option<Vec<State>> {
    searchlib::bfs::solve(
        init,
        |s| successors(s, x_max, y_max),
        |s| s.visited.len() == x_max * y_max,
    )
}

fn dfs(init: &State, x_max: usize, y_max: usize) -> Option<Vec<State>> {
    searchlib::dfs::solve(
        init,
        |s| successors(s, x_max, y_max),
        |s| s.visited.len() == x_max * y_max,
    )
}

fn gbfs(init: &State, x_max: usize, y_max: usize) -> Option<Vec<State>> {
    searchlib::gbfs::solve(
        init,
        |s| successors(s, x_max, y_max),
        |s| s.visited.len() == x_max * y_max,
        |s| heuristic(s, x_max, y_max),
    )
}

fn lgbfs(init: &State, x_max: usize, y_max: usize) -> Option<Vec<State>> {
    searchlib::gbfs::solve(
        init,
        |s| successors(s, x_max, y_max),
        |s| s.visited.len() == x_max * y_max,
        |s| heuristic(s, x_max, y_max),
    )
}

fn hcs(init: &State, x_max: usize, y_max: usize) -> Option<Vec<State>> {
    searchlib::hcs::solve(
        init,
        |s| successors(s, x_max, y_max),
        |s| s.visited.len() == x_max * y_max,
        |s| heuristic(s, x_max, y_max),
    )
}

fn ehcs(init: &State, x_max: usize, y_max: usize) -> Option<Vec<State>> {
    searchlib::ehcs::solve(
        init,
        |s| successors(s, x_max, y_max),
        |s| s.visited.len() == x_max * y_max,
        |s| heuristic(s, x_max, y_max),
    )
}

fn criterion_benchmark(c: &mut Criterion) {
    let init = State {
        pos: (0, 0),
        visited: BTreeSet::from([(0, 0)]),
    };
    let mut group = c.benchmark_group("visit_all");
    group.noise_threshold(0.05);
    group.bench_function("BFS", |b| b.iter(|| bfs(&init, 2, 3)));
    group.bench_function("DFS", |b| b.iter(|| dfs(&init, 3, 3)));
    group.bench_function("GBFS", |b| b.iter(|| gbfs(&init, 6, 6)));
    group.bench_function("LGBFS", |b| b.iter(|| lgbfs(&init, 6, 6)));
    group.bench_function("HCS", |b| b.iter(|| hcs(&init, 6, 6)));
    group.bench_function("EHCS", |b| b.iter(|| ehcs(&init, 6, 6)));
    group.finish();
}

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