use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use physarum::{
    agent::Agent,
    grid::{combine, Grid},
    model,
};
use rand::{rngs::StdRng, SeedableRng};

// Benchmark agent movement and deposition
fn agent_benchmark(c: &mut Criterion) {
    let mut group = c.benchmark_group("Agent Tick");
    let n_agents = [1_000, 10_000, 100_000];

    for &n in &n_agents {
        group.bench_with_input(BenchmarkId::from_parameter(n), &n, |b, &n| {
            let mut rng = StdRng::seed_from_u64(42);
            let agents = (0..n).map(|_| Agent::new(256, 256, &mut rng)).collect();
            let mut grid = Grid::new(256, 256, &mut rng, agents);

            b.iter(|| {
                grid.tick();
            });
        });
    }
    group.finish();
}

// Benchmark grid diffusion (blur)
fn diffusion_benchmark(c: &mut Criterion) {
    let mut group = c.benchmark_group("Grid Diffusion");
    let sizes = [(256, 256), (512, 512)];
    let radii = [1, 3];

    for &(w, h) in &sizes {
        for &r in &radii {
            group.bench_with_input(
                BenchmarkId::new("diffuse", format!("{}x{}_r{}", w, h, r)),
                &(w, h, r),
                |b, &(w, h, r)| {
                    b.iter_batched(
                        || {
                            let mut rng = StdRng::seed_from_u64(42);
                            Grid::new(w, h, &mut rng, vec![])
                        },
                        |mut grid| grid.diffuse(r),
                        BatchSize::SmallInput,
                    );
                },
            );
        }
    }
    group.finish();
}

// Benchmark grid combining
fn combine_benchmark(c: &mut Criterion) {
    let mut group = c.benchmark_group("Combine Grids");
    let populations = [2, 4];

    for &np in &populations {
        group.bench_with_input(BenchmarkId::from_parameter(np), &np, |b, &np| {
            b.iter_batched(
                || {
                    let mut rng = StdRng::seed_from_u64(42);
                    let grids = (0..np)
                        .map(|_| Grid::new(256, 256, &mut rng, vec![]))
                        .collect::<Vec<_>>();
                    let attraction_table = vec![vec![1.0; np]; np];
                    (grids, attraction_table)
                },
                |(mut grids, table)| combine(&mut grids, &table),
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

// Benchmark full model step
fn model_step_benchmark(c: &mut Criterion) {
    let mut group = c.benchmark_group("Model Step");
    let params = [(256, 256, 2), (512, 512, 4)];

    for &(w, h, np) in &params {
        group.bench_with_input(
            BenchmarkId::new("step", format!("{}x{}_p{}", w, h, np)),
            &(w, h, np),
            |b, &(w, h, np)| {
                b.iter_batched(
                    || model::Model::new(w, h, 1 << 16, np, 1),
                    |mut model| model.step(),
                    BatchSize::SmallInput,
                );
            },
        );
    }
    group.finish();
}

criterion_group!(
    benches,
    agent_benchmark,
    diffusion_benchmark,
    combine_benchmark,
    model_step_benchmark
);
criterion_main!(benches);