edgehdf5-memory 1.93.0

use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use tempfile::TempDir;
use edgehdf5_memory::bm25::BM25Index;
use edgehdf5_memory::cache::MemoryCache;
use edgehdf5_memory::hybrid::hybrid_search;
use edgehdf5_memory::pq::ProductQuantizer;
use edgehdf5_memory::ivf::{IVFIndex, IVFPQIndex};
use edgehdf5_memory::strategy::{self, SearchStrategy, HardwareCapabilities};
use edgehdf5_memory::vector_search;
use edgehdf5_memory::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};

// ---------------------------------------------------------------------------
// Simple deterministic PRNG (LCG)
// ---------------------------------------------------------------------------

struct Rng(u32);

impl Rng {
    fn new(seed: u32) -> Self {
        Self(seed)
    }
    fn next_u32(&mut self) -> u32 {
        self.0 = self.0.wrapping_mul(1103515245).wrapping_add(12345);
        self.0 >> 16
    }
    fn next_f32(&mut self) -> f32 {
        self.next_u32() as f32 / 65536.0 - 0.5
    }
    fn next_usize(&mut self, max: usize) -> usize {
        self.next_u32() as usize % max
    }
}

// ---------------------------------------------------------------------------
// Data generation helpers
// ---------------------------------------------------------------------------

const WORDS: &[&str] = &[
    "the", "quick", "brown", "fox", "jumps", "over", "lazy", "dog", "rust",
    "programming", "memory", "vector", "search", "index", "data", "system",
    "agent", "knowledge", "graph", "neural", "network", "machine", "learning",
    "deep", "embedding", "cosine", "similarity", "token", "chunk", "session",
    "channel", "hybrid",
];

fn make_vec(rng: &mut Rng, dim: usize) -> Vec<f32> {
    (0..dim).map(|_| rng.next_f32()).collect()
}

fn make_vecs(n: usize, dim: usize, seed: u32) -> Vec<Vec<f32>> {
    let mut rng = Rng::new(seed);
    (0..n).map(|_| make_vec(&mut rng, dim)).collect()
}

fn make_text(rng: &mut Rng, word_count: usize) -> String {
    (0..word_count)
        .map(|_| WORDS[rng.next_usize(WORDS.len())])
        .collect::<Vec<_>>()
        .join(" ")
}

fn make_texts(n: usize, word_count: usize, seed: u32) -> Vec<String> {
    let mut rng = Rng::new(seed);
    (0..n).map(|_| make_text(&mut rng, word_count)).collect()
}

fn make_entry(rng: &mut Rng, dim: usize) -> MemoryEntry {
    MemoryEntry {
        chunk: make_text(rng, 20),
        embedding: make_vec(rng, dim),
        source_channel: "bench".into(),
        timestamp: 1_000_000.0,
        session_id: "bench-session".into(),
        tags: "bench".into(),
    }
}

fn make_entries(n: usize, dim: usize, seed: u32) -> Vec<MemoryEntry> {
    let mut rng = Rng::new(seed);
    (0..n).map(|_| make_entry(&mut rng, dim)).collect()
}

fn make_config(dir: &TempDir, dim: usize) -> MemoryConfig {
    MemoryConfig::new(dir.path().join("bench.h5"), "bench-agent", dim)
}

// ---------------------------------------------------------------------------
// Save benchmarks
// ---------------------------------------------------------------------------

fn save_benches(c: &mut Criterion) {
    let dim = 384;

    c.bench_function("save_single", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let config = make_config(&dir, dim);
                let mem = HDF5Memory::create(config).unwrap();
                let mut rng = Rng::new(42);
                let entry = make_entry(&mut rng, dim);
                (dir, mem, entry)
            },
            |(_dir, mut mem, entry)| {
                mem.save(entry).unwrap();
            },
            BatchSize::SmallInput,
        );
    });

    c.bench_function("save_batch_100", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let config = make_config(&dir, dim);
                let mem = HDF5Memory::create(config).unwrap();
                let entries = make_entries(100, dim, 42);
                (dir, mem, entries)
            },
            |(_dir, mut mem, entries)| {
                mem.save_batch(entries).unwrap();
            },
            BatchSize::SmallInput,
        );
    });

    c.bench_function("save_batch_1000", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let config = make_config(&dir, dim);
                let mem = HDF5Memory::create(config).unwrap();
                let entries = make_entries(1000, dim, 42);
                (dir, mem, entries)
            },
            |(_dir, mut mem, entries)| {
                mem.save_batch(entries).unwrap();
            },
            BatchSize::LargeInput,
        );
    });
}

// ---------------------------------------------------------------------------
// Vector search benchmarks (SIMD vs baseline)
// ---------------------------------------------------------------------------

fn vector_search_benches(c: &mut Criterion) {
    let dim = 384;
    let query = make_vec(&mut Rng::new(99), dim);

    {
        let vectors = make_vecs(1_000, dim, 42);
        let tombstones = vec![0u8; 1_000];
        c.bench_function("vector_search_1k", |b| {
            b.iter(|| vector_search::cosine_similarity_batch(&query, &vectors, &tombstones));
        });
    }

    {
        let vectors = make_vecs(10_000, dim, 42);
        let tombstones = vec![0u8; 10_000];
        c.bench_function("simd_cosine_10k", |b| {
            b.iter(|| vector_search::cosine_similarity_batch(&query, &vectors, &tombstones));
        });
    }

    {
        let vectors = make_vecs(100_000, dim, 42);
        let tombstones = vec![0u8; 100_000];
        c.bench_function("simd_cosine_100k", |b| {
            b.iter(|| vector_search::cosine_similarity_batch(&query, &vectors, &tombstones));
        });
    }
}

// ---------------------------------------------------------------------------
// Pre-computed norms benchmark
// ---------------------------------------------------------------------------

fn prenorm_benches(c: &mut Criterion) {
    let dim = 384;
    let n = 10_000;
    let query = make_vec(&mut Rng::new(99), dim);
    let vectors = make_vecs(n, dim, 42);
    let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
    let tombstones = vec![0u8; n];

    c.bench_function("prenorm_search_10k", |b| {
        b.iter(|| {
            vector_search::cosine_similarity_batch_prenorm(&query, &vectors, &norms, &tombstones)
        });
    });

    c.bench_function("full_norm_search_10k", |b| {
        b.iter(|| vector_search::cosine_similarity_batch(&query, &vectors, &tombstones));
    });
}

// ---------------------------------------------------------------------------
// PQ search benchmark
// ---------------------------------------------------------------------------

fn pq_benches(c: &mut Criterion) {
    let dim = 384;
    let n = 100_000;
    let vectors = make_vecs(n, dim, 42);
    let query = make_vec(&mut Rng::new(99), dim);
    let tombstones = vec![0u8; n];

    // Train PQ on subset for speed
    let train_set: Vec<Vec<f32>> = vectors[..1000].to_vec();
    let pq = ProductQuantizer::train(&train_set, dim, 48, 256);
    let all_codes = pq.encode_all(&vectors);

    c.bench_function("pq_search_100k", |b| {
        b.iter(|| pq.search(&query, &all_codes, &tombstones, 10));
    });
}

// ---------------------------------------------------------------------------
// IVF-PQ search benchmark
// ---------------------------------------------------------------------------

fn ivf_pq_benches(c: &mut Criterion) {
    let dim = 384;
    let n = 100_000;
    let vectors = make_vecs(n, dim, 42);
    let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
    let query = make_vec(&mut Rng::new(99), dim);
    let tombstones = vec![0u8; n];

    // Train on subset
    let train_set: Vec<Vec<f32>> = vectors[..2000].to_vec();
    let ivf = IVFIndex::train(&train_set, dim, 100);
    // Re-assign all vectors
    let mut ivf_full = IVFIndex {
        centroids: ivf.centroids.clone(),
        num_clusters: ivf.num_clusters,
        dim,
        inverted_lists: vec![Vec::new(); ivf.num_clusters],
    };
    for (i, v) in vectors.iter().enumerate() {
        let c = ivf_full.assign(v);
        ivf_full.inverted_lists[c].push(i);
    }

    c.bench_function("ivf_search_100k_nprobe10", |b| {
        b.iter(|| ivf_full.search(&query, &vectors, &norms, &tombstones, 10, 10));
    });

    // IVF-PQ combined
    let pq = ProductQuantizer::train(&train_set, dim, 48, 256);
    let codes = pq.encode_all(&vectors);
    let ivfpq = IVFPQIndex {
        ivf: ivf_full,
        pq,
        codes,
    };

    c.bench_function("ivf_pq_search_100k", |b| {
        b.iter(|| ivfpq.search(&query, &vectors, &norms, &tombstones, 10, 100, 10));
    });
}

// ---------------------------------------------------------------------------
// BM25 benchmarks
// ---------------------------------------------------------------------------

fn bm25_benches(c: &mut Criterion) {
    {
        let docs = make_texts(1_000, 50, 42);
        let tombstones = vec![0u8; 1_000];
        let index = BM25Index::build(&docs, &tombstones);
        c.bench_function("bm25_search_1k", |b| {
            b.iter(|| index.search("rust programming memory", 10));
        });
    }

    {
        let docs = make_texts(10_000, 50, 42);
        let tombstones = vec![0u8; 10_000];
        let index = BM25Index::build(&docs, &tombstones);
        c.bench_function("bm25_search_10k", |b| {
            b.iter(|| index.search("rust programming memory", 10));
        });
    }
}

// ---------------------------------------------------------------------------
// Hybrid search benchmark
// ---------------------------------------------------------------------------

fn hybrid_benches(c: &mut Criterion) {
    let dim = 384;
    let n = 10_000;
    let vectors = make_vecs(n, dim, 42);
    let docs = make_texts(n, 50, 77);
    let tombstones = vec![0u8; n];
    let bm25 = BM25Index::build(&docs, &tombstones);
    let query_vec = make_vec(&mut Rng::new(99), dim);

    c.bench_function("hybrid_search_10k", |b| {
        b.iter(|| {
            hybrid_search(
                &query_vec,
                "rust memory search",
                &vectors,
                &docs,
                &tombstones,
                &bm25,
                0.7,
                0.3,
                10,
            )
        });
    });
}

// ---------------------------------------------------------------------------
// Compact benchmark
// ---------------------------------------------------------------------------

fn compact_benches(c: &mut Criterion) {
    let dim = 384;
    let n = 10_000;
    let mut rng = Rng::new(42);

    let mut cache = MemoryCache::new(dim);
    for _ in 0..n {
        cache.push(
            make_text(&mut rng, 20),
            make_vec(&mut rng, dim),
            "bench".into(),
            1_000_000.0,
            "session".into(),
            "tag".into(),
        );
    }
    // Tombstone 10%
    for i in (0..n).step_by(10) {
        cache.mark_deleted(i);
    }

    c.bench_function("compact_10pct_tombstoned", |b| {
        b.iter_batched(
            || cache.clone(),
            |mut c| {
                c.compact();
            },
            BatchSize::LargeInput,
        );
    });
}

// ---------------------------------------------------------------------------
// I/O benchmarks (open, snapshot, mmap)
// ---------------------------------------------------------------------------

fn io_benches(_c: &mut Criterion) {
    // Skipped: pre-existing file reopen issue prevents HDF5Memory::open() in bench
}

// ---------------------------------------------------------------------------
// Rayon parallel search benchmarks
// ---------------------------------------------------------------------------

fn rayon_benches(c: &mut Criterion) {
    let dim = 384;
    let query = make_vec(&mut Rng::new(99), dim);

    {
        let n = 10_000;
        let vectors = make_vecs(n, dim, 42);
        let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
        let tombstones = vec![0u8; n];

        c.bench_function("rayon_cosine_10k", |b| {
            b.iter(|| {
                use rayon::prelude::*;
                let query_norm = vector_search::compute_norm(&query);
                let num_cores = rayon::current_num_threads().max(1);
                let chunk_size = (n + num_cores - 1) / num_cores;
                let mut results: Vec<(usize, f32)> = vectors
                    .par_chunks(chunk_size)
                    .enumerate()
                    .flat_map(|(ci, chunk)| {
                        let base = ci * chunk_size;
                        chunk.iter().enumerate().filter_map(|(j, v)| {
                            let i = base + j;
                            if tombstones[i] != 0 { return None; }
                            let score = edgehdf5_memory::cosine_similarity_prenorm(
                                &query, query_norm, v, norms[i],
                            );
                            Some((i, score))
                        }).collect::<Vec<_>>()
                    })
                    .collect();
                results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
                results.truncate(10);
                results
            });
        });

        c.bench_function("sequential_cosine_10k", |b| {
            b.iter(|| {
                vector_search::cosine_similarity_batch_prenorm(&query, &vectors, &norms, &tombstones)
            });
        });
    }

    {
        let n = 100_000;
        let vectors = make_vecs(n, dim, 42);
        let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
        let tombstones = vec![0u8; n];

        c.bench_function("rayon_cosine_100k", |b| {
            b.iter(|| {
                use rayon::prelude::*;
                let query_norm = vector_search::compute_norm(&query);
                let num_cores = rayon::current_num_threads().max(1);
                let chunk_size = (n + num_cores - 1) / num_cores;
                let mut results: Vec<(usize, f32)> = vectors
                    .par_chunks(chunk_size)
                    .enumerate()
                    .flat_map(|(ci, chunk)| {
                        let base = ci * chunk_size;
                        chunk.iter().enumerate().filter_map(|(j, v)| {
                            let i = base + j;
                            if tombstones[i] != 0 { return None; }
                            let score = edgehdf5_memory::cosine_similarity_prenorm(
                                &query, query_norm, v, norms[i],
                            );
                            Some((i, score))
                        }).collect::<Vec<_>>()
                    })
                    .collect();
                results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
                results.truncate(10);
                results
            });
        });
    }
}

// ---------------------------------------------------------------------------
// BLAS search benchmarks
// ---------------------------------------------------------------------------

#[cfg(feature = "fast-math")]
fn blas_benches(c: &mut Criterion) {
    let dim = 384;
    let query = make_vec(&mut Rng::new(99), dim);

    {
        let n = 10_000;
        let vectors = make_vecs(n, dim, 42);
        let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
        let tombstones = vec![0u8; n];

        c.bench_function("blas_cosine_10k", |b| {
            b.iter(|| {
                vector_search::cosine_similarity_batch_blas(&query, &vectors, &norms, &tombstones, 10)
            });
        });
    }

    {
        let n = 100_000;
        let vectors = make_vecs(n, dim, 42);
        let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
        let tombstones = vec![0u8; n];

        c.bench_function("blas_cosine_100k", |b| {
            b.iter(|| {
                vector_search::cosine_similarity_batch_blas(&query, &vectors, &norms, &tombstones, 10)
            });
        });
    }

    // Batch norms benchmark
    {
        let n = 10_000;
        let vectors = make_vecs(n, dim, 42);
        let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();

        c.bench_function("blas_batch_norms_10k", |b| {
            b.iter(|| edgehdf5_memory::blas_search::blas_batch_norms(&flat, dim));
        });
    }
}

// ---------------------------------------------------------------------------
// Accelerate (cblas_sgemv / AMX) search benchmarks
// ---------------------------------------------------------------------------

#[cfg(any(feature = "accelerate", feature = "openblas"))]
fn accelerate_benches(c: &mut Criterion) {
    let dim = 384;
    let query = make_vec(&mut Rng::new(99), dim);

    {
        let n = 10_000;
        let vectors = make_vecs(n, dim, 42);
        let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
        let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
        let tombstones = vec![0u8; n];

        c.bench_function("accelerate_cosine_10k", |b| {
            b.iter(|| {
                edgehdf5_memory::accelerate_search::accelerate_cosine_batch(
                    &query, &flat, &norms, &tombstones, dim, 10,
                )
            });
        });

        // Compare: SIMD prenorm on same data
        c.bench_function("simd_prenorm_cosine_10k", |b| {
            b.iter(|| {
                vector_search::cosine_similarity_batch_prenorm(
                    &query, &vectors, &norms, &tombstones,
                )
            });
        });

        // Compare: Accelerate via Vec<Vec> wrapper
        c.bench_function("accelerate_vecs_cosine_10k", |b| {
            b.iter(|| {
                edgehdf5_memory::accelerate_search::accelerate_cosine_batch_vecs(
                    &query, &vectors, &norms, &tombstones, 10,
                )
            });
        });
    }

    {
        let n = 100_000;
        let vectors = make_vecs(n, dim, 42);
        let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
        let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
        let tombstones = vec![0u8; n];

        c.bench_function("accelerate_cosine_100k", |b| {
            b.iter(|| {
                edgehdf5_memory::accelerate_search::accelerate_cosine_batch(
                    &query, &flat, &norms, &tombstones, dim, 10,
                )
            });
        });
    }

    // Batch norms via cblas_snrm2
    {
        let n = 10_000;
        let vectors = make_vecs(n, dim, 42);
        let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();

        c.bench_function("accelerate_batch_norms_10k", |b| {
            b.iter(|| edgehdf5_memory::accelerate_search::accelerate_batch_norms(&flat, dim));
        });
    }

    // vDSP comparison (macOS only)
    #[cfg(target_os = "macos")]
    {
        let n = 10_000;
        let vectors = make_vecs(n, dim, 42);
        let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
        let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
        let tombstones = vec![0u8; n];

        c.bench_function("vdsp_cosine_10k", |b| {
            b.iter(|| {
                edgehdf5_memory::accelerate_search::vdsp_cosine_batch(
                    &query, &flat, &norms, &tombstones, dim, 10,
                )
            });
        });
    }
}

// ---------------------------------------------------------------------------
// Adaptive strategy benchmarks
// ---------------------------------------------------------------------------

fn adaptive_benches(c: &mut Criterion) {
    let dim = 384;
    let n = 10_000;
    let query = make_vec(&mut Rng::new(99), dim);
    let vectors = make_vecs(n, dim, 42);
    let norms: Vec<f32> = vectors.iter().map(|v| vector_search::compute_norm(v)).collect();
    let tombstones = vec![0u8; n];

    c.bench_function("adaptive_search_10k", |b| {
        let hw = HardwareCapabilities::detect();
        let strat = strategy::auto_select_strategy(n, &hw);
        b.iter(|| {
            strategy::search_with_metrics(
                &query, &vectors, &norms, &tombstones, 10, strat, None,
            )
        });
    });

    // Strategy comparison: run all CPU strategies on same dataset
    c.bench_function("strategy_scalar_10k", |b| {
        b.iter(|| {
            strategy::search_with_metrics(
                &query, &vectors, &norms, &tombstones, 10,
                SearchStrategy::Scalar, None,
            )
        });
    });

    c.bench_function("strategy_simd_10k", |b| {
        b.iter(|| {
            strategy::search_with_metrics(
                &query, &vectors, &norms, &tombstones, 10,
                SearchStrategy::SimdBruteForce, None,
            )
        });
    });

    c.bench_function("strategy_rayon_10k", |b| {
        b.iter(|| {
            strategy::search_with_metrics(
                &query, &vectors, &norms, &tombstones, 10,
                SearchStrategy::RayonParallel, None,
            )
        });
    });
}

// ---------------------------------------------------------------------------
// Criterion group & main
// ---------------------------------------------------------------------------

// ---------------------------------------------------------------------------
// WAL benchmarks — save latency with/without WAL
// ---------------------------------------------------------------------------

fn wal_benches(c: &mut Criterion) {
    let dim = 384;

    c.bench_function("save_with_wal_single", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let mut config = make_config(&dir, dim);
                config.wal_enabled = true;
                let mem = HDF5Memory::create(config).unwrap();
                let mut rng = Rng::new(42);
                let entry = make_entry(&mut rng, dim);
                (dir, mem, entry)
            },
            |(_dir, mut mem, entry)| {
                mem.save(entry).unwrap();
            },
            BatchSize::SmallInput,
        );
    });

    c.bench_function("save_without_wal_single", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let mut config = make_config(&dir, dim);
                config.wal_enabled = false;
                let mem = HDF5Memory::create(config).unwrap();
                let mut rng = Rng::new(42);
                let entry = make_entry(&mut rng, dim);
                (dir, mem, entry)
            },
            |(_dir, mut mem, entry)| {
                mem.save(entry).unwrap();
            },
            BatchSize::SmallInput,
        );
    });

    // WAL save into a store with 1K existing entries
    c.bench_function("save_wal_1k_existing", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let mut config = make_config(&dir, dim);
                config.wal_enabled = true;
                config.wal_max_entries = 5000;
                let mut mem = HDF5Memory::create(config).unwrap();
                let entries = make_entries(1000, dim, 99);
                mem.save_batch(entries).unwrap();
                let mut rng = Rng::new(42);
                let entry = make_entry(&mut rng, dim);
                (dir, mem, entry)
            },
            |(_dir, mut mem, entry)| {
                mem.save(entry).unwrap();
            },
            BatchSize::LargeInput,
        );
    });

    c.bench_function("wal_flush_100_entries", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let mut config = make_config(&dir, dim);
                config.wal_enabled = true;
                config.wal_max_entries = 5000;
                let mut mem = HDF5Memory::create(config).unwrap();
                let entries = make_entries(100, dim, 77);
                for e in entries { mem.save(e).unwrap(); }
                (dir, mem)
            },
            |(_dir, mut mem)| {
                mem.flush_wal().unwrap();
            },
            BatchSize::LargeInput,
        );
    });
}

// ---------------------------------------------------------------------------
// Decision Gate benchmarks
// ---------------------------------------------------------------------------

fn gate_benches(c: &mut Criterion) {
    use edgehdf5_memory::decision_gate::{DecisionGate, GateConfig};

    let gate = DecisionGate::new(GateConfig::default());

    c.bench_function("gate_trivial_skip", |b| {
        b.iter(|| gate.should_save("ok"));
    });

    c.bench_function("gate_nontrivial_pass", |b| {
        b.iter(|| gate.should_save("Tell me about the deployment architecture for the new system"));
    });

    c.bench_function("gate_short_phrase_skip", |b| {
        b.iter(|| gate.should_save("got it"));
    });

    c.bench_function("gate_ratio_check", |b| {
        b.iter(|| gate.should_save("yes yes definitely sure absolutely right"));
    });
}

// ---------------------------------------------------------------------------
// Hebbian activation benchmarks
// ---------------------------------------------------------------------------

fn hebbian_benches(c: &mut Criterion) {
    let dim = 384;

    c.bench_function("tick_session_1k", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let config = make_config(&dir, dim);
                let mut mem = HDF5Memory::create(config).unwrap();
                let entries = make_entries(1000, dim, 55);
                mem.save_batch(entries).unwrap();
                (dir, mem)
            },
            |(_dir, mut mem)| {
                mem.tick_session().unwrap();
            },
            BatchSize::LargeInput,
        );
    });

    c.bench_function("tick_session_10k", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let config = make_config(&dir, dim);
                let mut mem = HDF5Memory::create(config).unwrap();
                let entries = make_entries(10000, dim, 55);
                mem.save_batch(entries).unwrap();
                (dir, mem)
            },
            |(_dir, mut mem)| {
                mem.tick_session().unwrap();
            },
            BatchSize::LargeInput,
        );
    });
}

// ---------------------------------------------------------------------------
// Entity alias resolution benchmarks
// ---------------------------------------------------------------------------

fn alias_benches(c: &mut Criterion) {
    use edgehdf5_memory::knowledge::KnowledgeCache;

    // Build a knowledge cache with 50 entities and 100 aliases
    let mut kg = KnowledgeCache::new();
    let names = ["Alice", "Bob", "Charlie", "David", "Eve", "Frank", "Grace",
                 "Henry", "Irene", "Jack", "Karen", "Leo", "Mona", "Nick", "Olivia",
                 "Peter", "Quinn", "Rachel", "Sam", "Tina", "Uma", "Victor", "Wendy",
                 "Xander", "Yara", "Zack", "Acme Corp", "BigCo", "CloudNet", "DataSys",
                 "EdgeAI", "FinTech", "GlobalX", "HyperNet", "InfoSec", "JetOps",
                 "KernelDev", "LogicAI", "MeshNet", "NanoTech", "OpenStack",
                 "PlatformX", "QuantumAI", "RoboSys", "SkyNet", "TechFlow",
                 "UniCloud", "VirtNet", "WaveAI", "XenoData"];
    for (i, name) in names.iter().enumerate() {
        let id = kg.add_entity(name, "entity", -1);
        kg.add_alias(&format!("alias_{i}"), id as i64);
        kg.add_alias(&name.to_lowercase(), id as i64);
    }

    c.bench_function("alias_resolve_short_query", |b| {
        b.iter(|| kg.resolve_aliases("what does alias_7 do at alias_30?"));
    });

    c.bench_function("alias_resolve_no_match", |b| {
        b.iter(|| kg.resolve_aliases("tell me about the weather forecast for tomorrow"));
    });

    c.bench_function("alias_resolve_long_query", |b| {
        b.iter(|| kg.resolve_aliases("I need alias_0 to talk to alias_15 about the alias_27 project and also check with alias_42"));
    });
}

// ---------------------------------------------------------------------------
// MemoryStrategy benchmarks
// ---------------------------------------------------------------------------

fn strategy_benches(c: &mut Criterion) {
    use edgehdf5_memory::memory_strategy::*;
    use edgehdf5_memory::SearchResult;

    struct EmptyStore;
    impl MemoryStoreView for EmptyStore {
        fn search(&self, _: &[f32], _: usize) -> Vec<SearchResult> { vec![] }
        fn memory_count(&self) -> usize { 0 }
        fn entity_count(&self) -> usize { 0 }
    }

    let save_every = SaveEveryExchange::default();
    let shift = SaveOnSemanticShift {
        gate: edgehdf5_memory::decision_gate::DecisionGate::new(
            edgehdf5_memory::decision_gate::GateConfig::default()
        ),
        shift_threshold: 0.25,
        lookback_k: 5,
    };

    let exchange = Exchange {
        user_turn: "What is the deployment architecture for the new microservices?".into(),
        agent_turn: "The deployment uses Kubernetes with a service mesh for inter-service communication.".into(),
        session_id: "bench-session".into(),
        turn_number: 1,
        timestamp: 1_000_000.0,
        user_embedding: Some(vec![0.1; 384]),
        agent_embedding: Some(vec![0.2; 384]),
    };

    let trivial_exchange = Exchange {
        user_turn: "ok".into(),
        agent_turn: "Got it!".into(),
        session_id: "bench-session".into(),
        turn_number: 2,
        timestamp: 1_000_001.0,
        user_embedding: None,
        agent_embedding: None,
    };

    let store = EmptyStore;

    c.bench_function("strategy_save_every_substantive", |b| {
        b.iter(|| save_every.evaluate(&exchange, &store));
    });

    c.bench_function("strategy_save_every_trivial", |b| {
        b.iter(|| save_every.evaluate(&trivial_exchange, &store));
    });

    c.bench_function("strategy_semantic_shift_empty_store", |b| {
        b.iter(|| shift.evaluate(&exchange, &store));
    });
}

// ---------------------------------------------------------------------------
// AGENTS.md generation benchmark
// ---------------------------------------------------------------------------

fn agents_md_benches(c: &mut Criterion) {
    let dim = 384;

    c.bench_function("agents_md_generate_1k", |b| {
        b.iter_batched(
            || {
                let dir = TempDir::new().unwrap();
                let config = make_config(&dir, dim);
                let mut mem = HDF5Memory::create(config).unwrap();
                let entries = make_entries(1000, dim, 42);
                mem.save_batch(entries).unwrap();
                mem.add_entity("Alice", "person", -1).unwrap();
                mem.add_entity("ProjectX", "project", -1).unwrap();
                (dir, mem)
            },
            |(_dir, mem)| {
                mem.generate_agents_md();
            },
            BatchSize::LargeInput,
        );
    });
}

criterion_group!(
    benches,
    save_benches,
    vector_search_benches,
    prenorm_benches,
    pq_benches,
    ivf_pq_benches,
    bm25_benches,
    hybrid_benches,
    compact_benches,
    io_benches,
    rayon_benches,
    adaptive_benches,
    wal_benches,
    gate_benches,
    hebbian_benches,
    alias_benches,
    strategy_benches,
    agents_md_benches,
);

#[cfg(feature = "fast-math")]
criterion_group!(blas_bench_group, blas_benches);

#[cfg(any(feature = "accelerate", feature = "openblas"))]
criterion_group!(accelerate_bench_group, accelerate_benches);

// Criterion main: include all available bench groups
#[cfg(all(feature = "fast-math", any(feature = "accelerate", feature = "openblas")))]
criterion_main!(benches, blas_bench_group, accelerate_bench_group);

#[cfg(all(feature = "fast-math", not(any(feature = "accelerate", feature = "openblas"))))]
criterion_main!(benches, blas_bench_group);

#[cfg(all(not(feature = "fast-math"), any(feature = "accelerate", feature = "openblas")))]
criterion_main!(benches, accelerate_bench_group);

#[cfg(all(not(feature = "fast-math"), not(any(feature = "accelerate", feature = "openblas"))))]
criterion_main!(benches);