use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use monavec_core::{Config, Metric, MonaVec};
use monavec_core::quantize::{scalar, QuantPipeline};

fn make_unit_vec(dim: usize, seed: u64) -> Vec<f32> {
    use rand::SeedableRng;
    use rand::RngCore;
    use rand_chacha::ChaCha20Rng;
    let mut rng = ChaCha20Rng::seed_from_u64(seed);
    let mut v: Vec<f32> = (0..dim).map(|_| {
        let u = rng.next_u32() as f32 / u32::MAX as f32;
        u * 2.0 - 1.0
    }).collect();
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    v.iter_mut().for_each(|x| *x /= norm);
    v
}

fn bench_score_scalar_vs_simd(c: &mut Criterion) {
    let mut group = c.benchmark_group("score_4bit");

    for &dim in &[384usize, 768, 1536, 3072] {
        group.throughput(Throughput::Elements(dim as u64));

        let query   = make_unit_vec(dim, 0);
        let db_vec  = make_unit_vec(dim, 1);
        let packed  = scalar::encode_4bit(&db_vec);

        group.bench_with_input(BenchmarkId::new("scalar", dim), &dim, |b, _| {
            b.iter(|| scalar::score_4bit(black_box(&query), black_box(&packed)))
        });

        group.bench_with_input(BenchmarkId::new("dispatch", dim), &dim, |b, _| {
            b.iter(|| monavec_core::simd::score_4bit(black_box(&query), black_box(&packed)))
        });
    }

    group.finish();
}

fn bench_index_search(c: &mut Criterion) {
    let mut group = c.benchmark_group("index_search");

    for &(n, dim) in &[(1_000usize, 1536usize), (10_000, 768), (100_000, 384)] {
        let label = format!("{n}x{dim}");
        group.throughput(Throughput::Elements(n as u64));

        let config = Config::new(dim, Metric::Cosine, 4);
        let mut index = MonaVec::new(config).unwrap();

        let ids: Vec<u64>      = (0..n as u64).collect();
        let vecs: Vec<Vec<f32>> = (0..n).map(|i| make_unit_vec(dim, i as u64)).collect();
        index.add(&ids, &vecs).unwrap();

        let query = make_unit_vec(dim, 9999);

        group.bench_with_input(BenchmarkId::new("top10", &label), &label, |b, _| {
            b.iter(|| index.search(black_box(&query), 10, None).unwrap())
        });
    }

    group.finish();
}

criterion_group!(benches, bench_score_scalar_vs_simd, bench_index_search);
criterion_main!(benches);
