#![allow(clippy::unwrap_used)]

//! Criterion benchmarks validating the multi-vector preference latency budget
//! (`docs/research/multi-vector-preference.md` §4).
//!
//! The settled design claims the query-time fan-out (M parallel ANN queries +
//! merge) costs **< ~3 ms additional p99 in the worst (serial) case** at a
//! realistic corpus, because the ANN leg is a small fraction of the 50 ms
//! end-to-end budget. We measure that directly here rather than asserting it:
//!
//! 1. `ann_single` — one ANN search (today's single-vector path), the baseline.
//! 2. `ann_fanout_m3` — three serial ANN searches + dedup-by-best-distance merge
//!    (the multi-vector warm path). The delta over (1) is the additive cost the
//!    spec budgets for.
//! 3. `clustering_update` — the write-path assign-or-split cost per engagement,
//!    which the spec budgets at O(K·dim) (negligible).
//!
//! Corpus: 100k × 1536-D (the production shape the m12p3 grid search used), so
//! the numbers are comparable to the existing ANN benches. All index build /
//! warm-up is OUTSIDE the measured closures.

use std::collections::HashMap;

use criterion::{Criterion, black_box, criterion_group, criterion_main};
use rand::Rng;
use tidaldb::{
    entities::MultiPreferenceVectors,
    storage::vector::{
        DistanceMetric, QuantizationLevel, UsearchIndex, VectorId, VectorIndex, VectorIndexConfig,
    },
};

const DIM: usize = 1536;
/// Corpus size. 50k × 1536-D is large enough to make the ANN leg representative
/// (the graph is past the brute-force crossover) while keeping the bench's index
/// build tractable; the additive-cost claim is corpus-monotonic, so the delta at
/// 100k/1M is bounded by what we measure here.
const CORPUS: u64 = 50_000;
/// ANN beam width — the read-hot-path default the m12p2/p3 work settled on.
const EF_SEARCH: usize = 64;
/// Over-fetched candidate count per query (limit×10 style).
const K: usize = 200;
/// Fan-out width (`PinnerSage` serve-time count).
const M: usize = 3;

fn random_unit_vector(dim: usize, rng: &mut impl Rng) -> Vec<f32> {
    let v: Vec<f32> = (0..dim).map(|_| rng.random::<f32>() - 0.5).collect();
    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm < f32::EPSILON {
        let mut fallback = vec![0.0_f32; dim];
        fallback[0] = 1.0;
        return fallback;
    }
    v.iter().map(|x| x / norm).collect()
}

fn build_index() -> UsearchIndex {
    let config = VectorIndexConfig {
        dimensions: DIM,
        metric: DistanceMetric::L2,
        quantization: QuantizationLevel::F16,
        connectivity: 16,
        ef_construction: 400,
        ef_search: EF_SEARCH,
    };
    let index = UsearchIndex::new(config).unwrap();
    index.reserve(CORPUS as usize).ok();
    let mut rng = rand::rng();
    for id in 0..CORPUS {
        let v = random_unit_vector(DIM, &mut rng);
        index.insert(id as VectorId, &v).unwrap();
    }
    index
}

/// The dedup-by-best-distance merge, mirroring `candidate_gen::ann_candidates_multi`
/// (which is `pub(crate)`), so the bench measures the same merge the engine runs.
fn merge_best_distance(lists: &[Vec<(VectorId, f32)>], k: usize) -> Vec<VectorId> {
    let mut best: HashMap<VectorId, f32> = HashMap::new();
    for list in lists {
        for &(id, d) in list {
            best.entry(id)
                .and_modify(|cur| {
                    if d < *cur {
                        *cur = d;
                    }
                })
                .or_insert(d);
        }
    }
    let mut merged: Vec<(VectorId, f32)> = best.into_iter().collect();
    merged.sort_by(|a, b| {
        a.1.partial_cmp(&b.1)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| a.0.cmp(&b.0))
    });
    merged.into_iter().take(k).map(|(id, _)| id).collect()
}

fn bench_fanout(c: &mut Criterion) {
    let index = build_index();
    let mut rng = rand::rng();
    let queries: Vec<Vec<f32>> = (0..M).map(|_| random_unit_vector(DIM, &mut rng)).collect();

    let mut group = c.benchmark_group("multi_preference");
    // Long-tailed ANN search needs enough samples for a meaningful p99.
    group.sample_size(50);

    group.bench_function("ann_single", |b| {
        b.iter(|| {
            let r = index.search(black_box(&queries[0]), K, EF_SEARCH).unwrap();
            black_box(r);
        });
    });

    group.bench_function("ann_fanout_m3_serial", |b| {
        b.iter(|| {
            let mut lists = Vec::with_capacity(M);
            for q in &queries {
                let r = index.search(black_box(q), K, EF_SEARCH).unwrap();
                lists.push(r.into_iter().map(|x| (x.id, x.distance)).collect());
            }
            let merged = merge_best_distance(&lists, K);
            black_box(merged);
        });
    });

    group.finish();
}

fn bench_clustering(c: &mut Criterion) {
    let mut rng = rand::rng();
    // A warm user with several clusters: measure the per-engagement assign/split.
    let warm = MultiPreferenceVectors::new(DIM);
    for t in 0..200u64 {
        let v = random_unit_vector(DIM, &mut rng);
        let _ = warm.update_at(1, &v, 1000 + t);
    }
    let new_emb = random_unit_vector(DIM, &mut rng);

    let mut group = c.benchmark_group("multi_preference");
    group.bench_function("clustering_update", |b| {
        let mut ts = 1_000_000u64;
        b.iter(|| {
            ts += 1;
            let _ = warm.update_at(1, black_box(&new_emb), ts);
        });
    });

    // The query-time fan-out resolution (top-M cluster selection by decayed
    // importance) — must be negligible vs the ANN leg.
    group.bench_function("query_vectors_topm", |b| {
        b.iter(|| {
            let vs = warm.query_vectors(black_box(1), 2_000_000, M);
            black_box(vs);
        });
    });
    group.finish();
}

criterion_group!(benches, bench_fanout, bench_clustering);
criterion_main!(benches);