//! Search measurement harness: recall vs. speed for the HNSW index, and
//! end-to-end `hybrid_search` latency as the store grows.
//!
//! Every search-path change should be justified by a before/after run of this
//! binary. It reports, for deterministic synthetic data:
//!
//! * **ANN** — index build time, and for each `ef`: recall@10 against an exact
//!   brute-force scan, queries/second, and p50/p99 latency.
//! * **End to end** — `HDF5Memory`: ingest time, checkpoint time, `open()`
//!   time, the one-off cold index build (first query ever), the first query
//!   after a reopen, and steady-state `hybrid_search` p50/p99 at each size.
//!
//! Data is *clustered* (points = cluster centre + noise, unit-normalised), not
//! uniform: uniform random high-dimensional vectors are nearly equidistant,
//! which makes recall numbers meaningless and is nothing like embeddings.
//!
//! ```text
//! cargo run --release -p clawhdf5-bench --bin search_harness            # 1K, 10K
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --full  # + 100K
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --ann-only --uniform
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --options-study --full
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --signing-study --full
//! ```

use std::time::{Duration, Instant};

use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};

const DIM: usize = 384;
const K: usize = 10;
const N_QUERIES: usize = 200;
const HNSW_M: usize = 16;
const HNSW_EF_CONSTRUCTION: usize = 64;
const EF_VALUES: [usize; 5] = [16, 32, 64, 128, 256];

// ---------------------------------------------------------------------------
// Deterministic data
// ---------------------------------------------------------------------------

struct Rng(u64);

impl Rng {
    fn next_u64(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }

    /// Uniform in [0, 1).
    fn unit(&mut self) -> f32 {
        (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32
    }

    /// Approximately standard normal (sum of uniforms).
    fn gauss(&mut self) -> f32 {
        let sum: f32 = (0..6).map(|_| self.unit()).sum();
        (sum - 3.0) * std::f32::consts::SQRT_2
    }

    fn below(&mut self, n: usize) -> usize {
        (self.next_u64() % n as u64) as usize
    }
}

fn normalize(v: &mut [f32]) {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        v.iter_mut().for_each(|x| *x /= norm);
    }
}

struct Dataset {
    vectors: Vec<Vec<f32>>,
    queries: Vec<Vec<f32>>,
    /// Cluster id of each vector (used to give records topical text).
    cluster_of: Vec<usize>,
    query_cluster: Vec<usize>,
}

/// `--uniform`: isotropic random unit vectors instead of clusters. Not a
/// realistic workload, but a useful second distribution — a recall problem
/// that appears only on clustered data points at graph connectivity.
static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

/// `--int8`: build the HNSW index over int8-quantised vectors (a quarter of
/// the memory) instead of f32, to price the recall it costs.
static INT8: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

/// `--f16-first`: in `--float16-study`, run the float16 store first.
static F16_FIRST: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

/// `--rerank`: re-score the candidate pool against the exact vectors before
/// taking the top K.
static RERANK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

fn storage() -> Storage {
    if INT8.load(std::sync::atomic::Ordering::Relaxed) {
        Storage::Int8
    } else {
        Storage::Float32
    }
}

fn make_dataset(n: usize, seed: u64) -> Dataset {
    let mut rng = Rng(seed);
    if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) {
        let random_unit = |rng: &mut Rng| {
            let mut v: Vec<f32> = (0..DIM).map(|_| rng.gauss()).collect();
            normalize(&mut v);
            v
        };
        return Dataset {
            vectors: (0..n).map(|_| random_unit(&mut rng)).collect(),
            queries: (0..N_QUERIES).map(|_| random_unit(&mut rng)).collect(),
            cluster_of: vec![0; n],
            query_cluster: vec![0; N_QUERIES],
        };
    }
    let n_clusters = (n / 100).clamp(8, 512);
    let centres: Vec<Vec<f32>> = (0..n_clusters)
        .map(|_| {
            let mut c: Vec<f32> = (0..DIM).map(|_| rng.gauss()).collect();
            normalize(&mut c);
            c
        })
        .collect();
    let point = |rng: &mut Rng, cluster: usize| {
        // Noise comparable to the centre's per-dimension magnitude, so
        // clusters overlap and the nearest neighbours are non-trivial.
        let scale = 0.6 / (DIM as f32).sqrt();
        let mut v: Vec<f32> = centres[cluster]
            .iter()
            .map(|c| c + rng.gauss() * scale)
            .collect();
        normalize(&mut v);
        v
    };
    let mut vectors = Vec::with_capacity(n);
    let mut cluster_of = Vec::with_capacity(n);
    for _ in 0..n {
        let c = rng.below(n_clusters);
        vectors.push(point(&mut rng, c));
        cluster_of.push(c);
    }
    let mut queries = Vec::with_capacity(N_QUERIES);
    let mut query_cluster = Vec::with_capacity(N_QUERIES);
    for _ in 0..N_QUERIES {
        let c = rng.below(n_clusters);
        queries.push(point(&mut rng, c));
        query_cluster.push(c);
    }
    Dataset {
        vectors,
        queries,
        cluster_of,
        query_cluster,
    }
}

const WORDS: &[&str] = &[
    "deploy", "latency", "cache", "schema", "index", "vector", "memory", "agent", "kernel",
    "buffer", "socket", "thread", "tensor", "gradient", "ledger", "invoice", "meeting", "roadmap",
    "customer", "contract", "sensor", "orbit", "protein", "genome", "harbor", "bridge", "engine",
    "battery", "harvest", "weather", "museum", "recipe",
];

/// Text whose vocabulary is biased by cluster, so keyword and vector signals
/// agree the way they do for real embedded text.
fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String {
    let topic = [
        WORDS[cluster % WORDS.len()],
        WORDS[(cluster / 7 + 3) % WORDS.len()],
    ];
    let mut words = Vec::with_capacity(14);
    for j in 0..14 {
        if j % 3 == 0 {
            words.push(topic[j / 3 % 2]);
        } else {
            words.push(WORDS[rng.below(WORDS.len())]);
        }
    }
    format!("record {i}: {}", words.join(" "))
}

// ---------------------------------------------------------------------------
// Measurement helpers
// ---------------------------------------------------------------------------

/// Exact cosine distance between unit-length vectors.
fn exact_dist(a: &[f32], b: &[f32]) -> f32 {
    1.0 - a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>()
}

fn exact_top_k(vectors: &[Vec<f32>], query: &[f32], k: usize) -> Vec<usize> {
    // Vectors are unit length, so cosine order == dot-product order.
    let mut scored: Vec<(usize, f32)> = vectors
        .iter()
        .enumerate()
        .map(|(i, v)| (i, v.iter().zip(query).map(|(a, b)| a * b).sum()))
        .collect();
    scored.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
    scored.truncate(k);
    scored.into_iter().map(|(i, _)| i).collect()
}

struct Latency {
    p50: Duration,
    p99: Duration,
    qps: f64,
}

fn summarize(mut samples: Vec<Duration>) -> Latency {
    samples.sort();
    let total: Duration = samples.iter().sum();
    let at = |q: f64| samples[((samples.len() - 1) as f64 * q).round() as usize];
    Latency {
        p50: at(0.50),
        p99: at(0.99),
        qps: samples.len() as f64 / total.as_secs_f64(),
    }
}

/// Counts live heap bytes, so a structure's cost can be measured by
/// difference.
///
/// RSS cannot do this from inside one process: freeing a large structure
/// returns its pages to the allocator's pool rather than to the OS, so
/// allocating the next one shows no change. Measured that way, a store that
/// holds the corpus twice and one that holds it once look identical.
struct CountingAllocator;

static LIVE_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0);

/// High-water mark of [`LIVE_BYTES`] since it was last reset.
///
/// Live bytes at a checkpoint cannot see a buffer that was allocated and
/// freed in between, and that is exactly the shape of a transient copy —
/// which still has to fit in memory while it exists.
static PEAK_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0);

fn note_peak(live: i64) {
    PEAK_BYTES.fetch_max(live, std::sync::atomic::Ordering::Relaxed);
}

// SAFETY: every method forwards to the system allocator with the same layout
// it was given, and only adds bookkeeping around it.
unsafe impl std::alloc::GlobalAlloc for CountingAllocator {
    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
        let ptr = unsafe { std::alloc::System.alloc(layout) };
        if !ptr.is_null() {
            let live = LIVE_BYTES
                .fetch_add(layout.size() as i64, std::sync::atomic::Ordering::Relaxed)
                + layout.size() as i64;
            note_peak(live);
        }
        ptr
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        LIVE_BYTES.fetch_sub(layout.size() as i64, std::sync::atomic::Ordering::Relaxed);
        unsafe { std::alloc::System.dealloc(ptr, layout) }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
        let new_ptr = unsafe { std::alloc::System.realloc(ptr, layout, new_size) };
        if !new_ptr.is_null() {
            let delta = new_size as i64 - layout.size() as i64;
            let live = LIVE_BYTES.fetch_add(delta, std::sync::atomic::Ordering::Relaxed) + delta;
            note_peak(live);
        }
        new_ptr
    }
}

#[global_allocator]
static ALLOCATOR: CountingAllocator = CountingAllocator;

/// Live heap bytes right now.
fn heap_bytes() -> u64 {
    LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64
}

/// Start watching for a new high-water mark from the current live total.
fn reset_peak() {
    PEAK_BYTES.store(
        LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed),
        std::sync::atomic::Ordering::Relaxed,
    );
}

/// The highest live total seen since [`reset_peak`].
fn peak_bytes() -> u64 {
    PEAK_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64
}

fn mib(bytes: u64) -> f64 {
    bytes as f64 / (1 << 20) as f64
}

fn micros(d: Duration) -> f64 {
    d.as_secs_f64() * 1e6
}

fn millis(d: Duration) -> f64 {
    d.as_secs_f64() * 1e3
}

// ---------------------------------------------------------------------------
// ANN: recall vs speed
// ---------------------------------------------------------------------------

fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
    let data = make_dataset(n, 0xA11CE ^ n as u64);
    let truth: Vec<Vec<usize>> = data
        .queries
        .iter()
        .map(|q| exact_top_k(&data.vectors, q, K))
        .collect();

    let started = Instant::now();
    let index = HnswIndex::build_with(
        &data.vectors,
        HNSW_M,
        HNSW_EF_CONSTRUCTION,
        DistanceMetric::Cosine,
        storage(),
    );
    let build = started.elapsed();

    // Exact scan baseline, for scale.
    let exact = summarize(
        data.queries
            .iter()
            .map(|q| {
                let t = Instant::now();
                std::hint::black_box(exact_top_k(&data.vectors, q, K));
                t.elapsed()
            })
            .collect(),
    );

    println!(
        "\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}, storage = {:?}\n",
        index.storage()
    );
    println!(
        "build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n",
        millis(build),
        n as f64 / build.as_secs_f64(),
        exact.qps,
        micros(exact.p50)
    );
    println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |");
    println!("|---:|---:|---:|---:|---:|");
    // With a quantised index the distances it returns are approximate, so
    // the candidates are re-scored against the exact vectors the caller
    // already holds (in the agent, the embedding cache) before taking the
    // top K. `--rerank` prices that: it costs one exact distance per
    // candidate and is what decides whether int8 is usable.
    let rerank = RERANK.load(std::sync::atomic::Ordering::Relaxed);
    let pool = if rerank { K * 4 } else { K };
    for ef in EF_VALUES {
        let mut hits = 0usize;
        let mut samples = Vec::with_capacity(data.queries.len());
        for (q, want) in data.queries.iter().zip(&truth) {
            let t = Instant::now();
            let mut got = index.search(q, pool, ef.max(pool));
            if rerank {
                for cand in &mut got {
                    cand.1 = exact_dist(&data.vectors[cand.0], q);
                }
                got.select_nth_unstable_by(K - 1, |a, b| a.1.total_cmp(&b.1));
                got.truncate(K);
            }
            samples.push(t.elapsed());
            hits += got.iter().filter(|(id, _)| want.contains(id)).count();
        }
        let recall = hits as f64 / (K * data.queries.len()) as f64;
        let lat = summarize(samples);
        println!(
            "| {ef} | {recall:.4} | {:.0} | {:.0} | {:.0} |",
            lat.qps,
            micros(lat.p50),
            micros(lat.p99)
        );
        json.push(serde_json::json!({
            "bench": "hnsw", "n": n, "ef": ef, "recall_at_10": recall,
            "qps": lat.qps, "p50_us": micros(lat.p50), "p99_us": micros(lat.p99),
            "build_ms": millis(build),
        }));
    }
}

// ---------------------------------------------------------------------------
// End to end: HDF5Memory::hybrid_search
// ---------------------------------------------------------------------------

fn bench_end_to_end(n: usize, json: &mut Vec<serde_json::Value>) {
    let data = make_dataset(n, 0xE2E ^ n as u64);
    let dir = tempfile::TempDir::new().unwrap();
    let path = dir.path().join("store.h5");
    let mut rng = Rng(7);

    let entries: Vec<MemoryEntry> = data
        .vectors
        .iter()
        .enumerate()
        .map(|(i, v)| MemoryEntry {
            chunk: text_for(data.cluster_of[i], i, &mut rng),
            embedding: v.clone(),
            source_channel: "bench".into(),
            timestamp: i as f64,
            session_id: format!("s{}", i % 50),
            tags: format!("t{i}"),
        })
        .collect();
    let query_texts: Vec<String> = data
        .query_cluster
        .iter()
        .enumerate()
        .map(|(i, c)| text_for(*c, i, &mut rng))
        .collect();

    let mut mem = HDF5Memory::create(MemoryConfig::new(path.clone(), "bench", DIM)).unwrap();
    let t = Instant::now();
    mem.save_batch(entries).unwrap();
    let ingest = t.elapsed();
    // The very first query builds the vector and keyword indexes from
    // scratch. It happens once per store, not once per session: the checkpoint
    // below saves the vector index, so a later `open()` reloads it.
    let t = Instant::now();
    std::hint::black_box(mem.hybrid_search(&data.queries[1], &query_texts[1], 0.7, 0.3, K));
    let cold_build = t.elapsed();

    let t = Instant::now();
    mem.flush_wal().unwrap();
    let checkpoint = t.elapsed();
    drop(mem);

    let t = Instant::now();
    let mut mem = HDF5Memory::open(&path).unwrap();
    let open = t.elapsed();

    // The first query after open pays for whatever is rebuilt lazily.
    let t = Instant::now();
    std::hint::black_box(mem.hybrid_search(&data.queries[0], &query_texts[0], 0.7, 0.3, K));
    let first_query = t.elapsed();

    // Fewer steady-state samples at large N: each query is currently O(N).
    let samples_wanted = if n >= 100_000 { 20 } else { N_QUERIES.min(100) };
    let steady = summarize(
        (0..samples_wanted)
            .map(|i| {
                let t = Instant::now();
                std::hint::black_box(mem.hybrid_search(
                    &data.queries[i % N_QUERIES],
                    &query_texts[i % N_QUERIES],
                    0.7,
                    0.3,
                    K,
                ));
                t.elapsed()
            })
            .collect(),
    );

    println!(
        "| {n} | {:.0} | {:.0} | {:.1} | {:.1} | {:.1} | {:.2} | {:.2} | {:.1} |",
        millis(ingest),
        millis(cold_build),
        millis(checkpoint),
        millis(open),
        millis(first_query),
        millis(steady.p50),
        millis(steady.p99),
        steady.qps
    );
    json.push(serde_json::json!({
        "bench": "hybrid_search", "n": n,
        "ingest_ms": millis(ingest), "cold_index_build_ms": millis(cold_build),
        "checkpoint_ms": millis(checkpoint),
        "open_ms": millis(open), "first_query_ms": millis(first_query),
        "p50_ms": millis(steady.p50), "p99_ms": millis(steady.p99), "qps": steady.qps,
    }));
}

// ---------------------------------------------------------------------------
// Signing study: what does an Ed25519-signed checkpoint cost?
// ---------------------------------------------------------------------------

/// `--signing-study`: checkpoint time unsigned vs signed, `verify` time, and
/// the file-size cost of the stored per-record hashes. Default store
/// settings (float16, int8 index). Medians of five checkpoints / three
/// verifies.
fn signing_study(n: usize) {
    use clawhdf5_agent::signing::SigningKey;
    let data = make_dataset(n, 0x516 ^ n as u64);
    let mut rng = Rng(9);
    let entries: Vec<MemoryEntry> = data
        .vectors
        .iter()
        .enumerate()
        .map(|(i, v)| MemoryEntry {
            chunk: text_for(data.cluster_of[i], i, &mut rng),
            embedding: v.clone(),
            source_channel: "bench".into(),
            timestamp: i as f64,
            session_id: format!("s{}", i % 50),
            tags: format!("t{i}"),
        })
        .collect();
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("sign.h5");
    let mut mem = HDF5Memory::create(MemoryConfig::new(path.clone(), "bench", DIM)).unwrap();
    mem.save_batch(entries).unwrap();
    std::hint::black_box(mem.hybrid_search(&data.queries[0], "", 1.0, 0.0, K));

    let median = |mut v: Vec<Duration>| {
        v.sort();
        v[v.len() / 2]
    };
    let checkpoint = |mem: &mut HDF5Memory| {
        median(
            (0..5)
                .map(|_| {
                    let t = Instant::now();
                    mem.flush_wal().unwrap();
                    t.elapsed()
                })
                .collect(),
        )
    };
    let unsigned = checkpoint(&mut mem);
    let unsigned_bytes = std::fs::metadata(&path).unwrap().len();
    let key = SigningKey::from_bytes(&[7; 32]);
    mem.set_signing_key(key.clone());
    let signed = checkpoint(&mut mem);
    let signed_bytes = std::fs::metadata(&path).unwrap().len();
    drop(mem);
    let vk = key.verifying_key();
    let verify = median(
        (0..3)
            .map(|_| {
                let t = Instant::now();
                let r = HDF5Memory::verify(&path, &vk).unwrap();
                let d = t.elapsed();
                assert!(r.is_valid());
                d
            })
            .collect(),
    );
    println!(
        "| {n} | {:.1} | {:.1} | {:+.1} | {:.1} | {:+.2} |",
        millis(unsigned),
        millis(signed),
        millis(signed) - millis(unsigned),
        millis(verify),
        (signed_bytes as f64 - unsigned_bytes as f64) / (1024.0 * 1024.0),
    );
}

// ---------------------------------------------------------------------------
// Search options study: source filters, re-ranking, confidence rejection
// ---------------------------------------------------------------------------

/// `--options-study`: what `HDF5Memory::search`'s options cost and whether a
/// filtered search finds the right records. Filters keep 50%, 10% or 1% of
/// the store at random, or two whole clusters away from the query (the case
/// the index cannot serve, which falls back to an exact scan). Recall is
/// vector-only against an exact scan of the allowed records; latency is full
/// hybrid search. Hebbian boosting is off.
fn options_study(n: usize) {
    use clawhdf5_agent::SearchOptions;
    use clawhdf5_agent::confidence::ConfidenceConfig;
    use clawhdf5_agent::hybrid::Fusion;
    use clawhdf5_agent::reranker::ReRankConfig;

    let data = make_dataset(n, 0x0B7 ^ n as u64);
    let n_clusters = data.cluster_of.iter().max().map_or(1, |m| m + 1);
    let mut rng = Rng(5);
    let bucket_of: Vec<usize> = (0..n).map(|_| rng.below(100)).collect();
    let bucket = &bucket_of;
    let query_texts: Vec<String> = data
        .query_cluster
        .iter()
        .enumerate()
        .map(|(i, c)| text_for(*c, i, &mut rng))
        .collect();
    let exact_top = |q: &[f32], allowed: &dyn Fn(usize) -> bool| -> Vec<usize> {
        let mut s: Vec<(usize, f32)> = (0..n)
            .filter(|&i| allowed(i))
            .map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum()))
            .collect();
        s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
        s.into_iter().take(K).map(|(i, _)| i).collect()
    };

    // Two stores: channel = random bucket, and channel = cluster.
    let dir = tempfile::tempdir().unwrap();
    let mut stores = Vec::new();
    for by_cluster in [false, true] {
        let mut rng = Rng(3);
        let entries: Vec<MemoryEntry> = data
            .vectors
            .iter()
            .enumerate()
            .map(|(i, v)| MemoryEntry {
                chunk: text_for(data.cluster_of[i], i, &mut rng),
                embedding: v.clone(),
                source_channel: if by_cluster {
                    format!("c{}", data.cluster_of[i])
                } else {
                    format!("b{}", bucket[i])
                },
                timestamp: i as f64,
                session_id: format!("s{}", i % 50),
                tags: format!("t{i}"),
            })
            .collect();
        let mut config = MemoryConfig::new(
            dir.path().join(format!("opt_{by_cluster}.h5")),
            "bench",
            DIM,
        );
        config.hebbian_boost = 0.0;
        let mut mem = HDF5Memory::create(config).unwrap();
        mem.save_batch(entries).unwrap();
        std::hint::black_box(mem.search(&data.queries[0], "", &SearchOptions::new(K)));
        stores.push(mem);
    }

    let vector_only = SearchOptions::new(K).with_fusion(Fusion::Weighted {
        vector: 1.0,
        keyword: 0.0,
    });
    // (label, store, channels for query i, allowed(i, record))
    type Case<'a> = (
        String,
        usize,
        Box<dyn Fn(usize) -> Option<Vec<String>> + 'a>,
        Box<dyn Fn(usize, usize) -> bool + 'a>,
    );
    let mut cases: Vec<Case> = vec![(
        "no filter".into(),
        0,
        Box::new(|_| None),
        Box::new(|_, _| true),
    )];
    for pct in [50usize, 10, 1] {
        cases.push((
            format!("random {pct}%"),
            0,
            Box::new(move |_| Some((0..pct).map(|b| format!("b{b}")).collect())),
            Box::new(move |_, i| bucket[i] < pct),
        ));
    }
    let d = &data;
    let away = move |qi: usize| {
        let qc = d.query_cluster[qi];
        [
            (qc + n_clusters / 3) % n_clusters,
            (qc + 2 * n_clusters / 3) % n_clusters,
        ]
    };
    cases.push((
        "2 clusters away from the query".into(),
        1,
        Box::new(move |qi| Some(away(qi).iter().map(|c| format!("c{c}")).collect())),
        Box::new(move |qi, i| away(qi).contains(&d.cluster_of[i])),
    ));

    for (label, store, channels, allowed) in &cases {
        let mem = &mut stores[*store];
        let mut hits = 0;
        let mut kept = 0;
        for (qi, q) in data.queries.iter().enumerate() {
            let mut opts = vector_only.clone();
            opts.source_channels = channels(qi);
            let got = mem.search(q, "", &opts);
            let want = exact_top(q, &|i| allowed(qi, i));
            kept += want.len();
            hits += got.iter().filter(|r| want.contains(&r.index)).count();
        }
        let latency = summarize(
            (0..N_QUERIES)
                .map(|qi| {
                    let mut opts = SearchOptions::new(K);
                    opts.source_channels = channels(qi);
                    let t = Instant::now();
                    std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts));
                    t.elapsed()
                })
                .collect(),
        );
        println!(
            "| {n} | {label} | {:.4} | {:.3} | {:.3} |",
            hits as f64 / kept.max(1) as f64,
            millis(latency.p50),
            millis(latency.p99),
        );
    }

    let mem = &mut stores[0];
    for (label, opts) in [
        (
            "re-rank",
            SearchOptions::new(K).with_rerank(ReRankConfig::default()),
        ),
        (
            "re-rank + confidence",
            SearchOptions::new(K)
                .with_rerank(ReRankConfig::default())
                .with_confidence(ConfidenceConfig::default()),
        ),
    ] {
        let latency = summarize(
            (0..N_QUERIES)
                .map(|qi| {
                    let t = Instant::now();
                    std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts));
                    t.elapsed()
                })
                .collect(),
        );
        println!(
            "| {n} | {label} | — | {:.3} | {:.3} |",
            millis(latency.p50),
            millis(latency.p99)
        );
    }
}

// ---------------------------------------------------------------------------
// float16 study: what does half-precision embedding storage cost?
// ---------------------------------------------------------------------------

/// `--float16-study`: the same data in an `f32` store and a `float16` store.
/// Reports file size, checkpoint and open time, vector-search recall@10
/// against an exact scan of the *original* f32 vectors, how often the two
/// stores return the same top 10, and `hybrid_search` latency. Hebbian
/// boosting is off, so every query sees the same store.
fn float16_study(n: usize) {
    let data = make_dataset(n, 0xF16 ^ n as u64);
    let mut rng = Rng(11);
    let query_texts: Vec<String> = data
        .query_cluster
        .iter()
        .enumerate()
        .map(|(i, c)| text_for(*c, i, &mut rng))
        .collect();

    // Exact top K by cosine (the vectors are unit length) on the f32 inputs.
    let exact: Vec<Vec<usize>> = data
        .queries
        .iter()
        .map(|q| {
            let mut scored: Vec<(usize, f32)> = data
                .vectors
                .iter()
                .enumerate()
                .map(|(i, v)| (i, v.iter().zip(q).map(|(a, b)| a * b).sum()))
                .collect();
            scored.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
            scored.into_iter().take(K).map(|(i, _)| i).collect()
        })
        .collect();

    let dir = tempfile::tempdir().unwrap();
    let mut per_variant: Vec<(bool, Vec<Vec<usize>>)> = Vec::new();
    // `--f16-first` swaps the order, to check the numbers do not depend on
    // which store runs first (page cache, allocator, CPU frequency).
    let order = if F16_FIRST.load(std::sync::atomic::Ordering::Relaxed) {
        [true, false]
    } else {
        [false, true]
    };
    for float16 in order {
        let path = dir.path().join(format!("f16study_{float16}.h5"));
        let mut rng = Rng(3);
        let entries: Vec<MemoryEntry> = data
            .vectors
            .iter()
            .enumerate()
            .map(|(i, v)| MemoryEntry {
                chunk: text_for(data.cluster_of[i], i, &mut rng),
                embedding: v.clone(),
                source_channel: "bench".into(),
                timestamp: i as f64,
                session_id: format!("s{}", i % 50),
                tags: format!("t{i}"),
            })
            .collect();
        let mut config = MemoryConfig::new(path.clone(), "bench", DIM);
        config.float16 = float16;
        config.hebbian_boost = 0.0;
        let mut mem = HDF5Memory::create(config).unwrap();
        mem.save_batch(entries).unwrap();
        // Build the indexes, then time a checkpoint that writes everything.
        std::hint::black_box(mem.hybrid_search(&data.queries[0], "", 1.0, 0.0, K));
        let t = Instant::now();
        mem.flush_wal().unwrap();
        let checkpoint = t.elapsed();
        drop(mem);
        let file_bytes = std::fs::metadata(&path).unwrap().len();

        // Median of three opens.
        let mut opens: Vec<Duration> = (0..3)
            .map(|_| {
                let t = Instant::now();
                let m = HDF5Memory::open(&path).unwrap();
                let d = t.elapsed();
                drop(m);
                d
            })
            .collect();
        opens.sort();
        let mut mem = HDF5Memory::open(&path).unwrap();

        // Vector-only search: empty text, all weight on the vector stage.
        let results: Vec<Vec<usize>> = data
            .queries
            .iter()
            .map(|q| {
                mem.hybrid_search(q, "", 1.0, 0.0, K)
                    .iter()
                    .map(|r| r.index)
                    .collect()
            })
            .collect();
        let hits: usize = results
            .iter()
            .zip(&exact)
            .map(|(got, want)| got.iter().filter(|i| want.contains(i)).count())
            .sum();
        let recall = hits as f64 / (K * data.queries.len()) as f64;

        let latency = summarize(
            (0..N_QUERIES)
                .map(|i| {
                    let t = Instant::now();
                    std::hint::black_box(mem.hybrid_search(
                        &data.queries[i],
                        &query_texts[i],
                        0.4,
                        0.6,
                        K,
                    ));
                    t.elapsed()
                })
                .collect(),
        );
        let overlap = match per_variant.first() {
            Some((_, other)) => {
                let same: usize = results
                    .iter()
                    .zip(other)
                    .map(|(a, b)| a.iter().filter(|i| b.contains(i)).count())
                    .sum();
                format!("{:.4}", same as f64 / (K * data.queries.len()) as f64)
            }
            None => "—".into(),
        };
        println!(
            "| {n} | {} | {:.1} | {:.0} | {:.1} | {recall:.4} | {overlap} | {:.3} |",
            if float16 { "float16" } else { "f32" },
            mib(file_bytes),
            millis(checkpoint),
            millis(opens[1]),
            millis(latency.p50),
        );
        per_variant.push((float16, results));
    }
}

// ---------------------------------------------------------------------------
// Fusion study: does capping the keyword candidate pool change the ranking?
// ---------------------------------------------------------------------------

/// `hybrid_search` min-max normalises each signal over the candidates it is
/// given. The vector stage supplies a pool of `max(8k, 64)`; the keyword stage
/// supplies *every* matching record, which is what now dominates query time.
/// This compares the current fusion with one whose keyword stage is capped to
/// a pool, reporting how often the final top-k agree and what each costs.
fn fusion_study(n: usize) {
    use clawhdf5_agent::bm25::BM25Index;
    use clawhdf5_agent::hybrid::merge_vector_keyword;

    let data = make_dataset(n, 0xE2E ^ n as u64);
    let mut rng = Rng(7);
    let texts: Vec<String> = (0..n)
        .map(|i| text_for(data.cluster_of[i], i, &mut rng))
        .collect();
    let query_texts: Vec<String> = data
        .query_cluster
        .iter()
        .enumerate()
        .map(|(i, c)| text_for(*c, i, &mut rng))
        .collect();
    let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
    let index = HnswIndex::build_with(
        &data.vectors,
        HNSW_M,
        HNSW_EF_CONSTRUCTION,
        DistanceMetric::Cosine,
        storage(),
    );

    let vec_pool = (K * 8).max(64);
    println!("\n### Fusion study, N = {n} (k = {K}, weights 0.7 / 0.3, vector pool {vec_pool})\n");
    println!(
        "| keyword pool | top-{K} overlap vs full | identical top-{K} | same #1 | keyword+merge µs |"
    );
    println!("|---:|---:|---:|---:|---:|");

    let fuse = |q: usize, kw_pool: usize| -> (Vec<usize>, Duration) {
        let vec_scores: Vec<(usize, f32)> = index
            .search(&data.queries[q], vec_pool, vec_pool)
            .into_iter()
            .map(|(id, d)| (id, 1.0 - d))
            .collect();
        let t = Instant::now();
        let kw = bm25.search(&query_texts[q], kw_pool);
        let merged = merge_vector_keyword(vec_scores, kw, 0.7, 0.3, K);
        let took = t.elapsed();
        (merged.into_iter().map(|(id, _)| id).collect(), took)
    };

    let full: Vec<(Vec<usize>, Duration)> = (0..N_QUERIES).map(|q| fuse(q, n)).collect();
    let full_time: Duration = full.iter().map(|f| f.1).sum();
    println!(
        "| all ({n}) | 1.0000 | 100.0% | 100.0% | {:.0} |",
        micros(full_time) / N_QUERIES as f64
    );
    for pool in [vec_pool, vec_pool * 4, 1000] {
        if pool >= n {
            continue;
        }
        let (mut overlap, mut identical, mut same_first) = (0usize, 0usize, 0usize);
        let mut time = Duration::ZERO;
        for (q, (want, _)) in full.iter().enumerate() {
            let (got, took) = fuse(q, pool);
            time += took;
            overlap += got.iter().filter(|id| want.contains(id)).count();
            identical += usize::from(&got == want);
            same_first += usize::from(got.first() == want.first());
        }
        println!(
            "| {pool} | {:.4} | {:.1}% | {:.1}% | {:.0} |",
            overlap as f64 / (K * N_QUERIES) as f64,
            100.0 * identical as f64 / N_QUERIES as f64,
            100.0 * same_first as f64 / N_QUERIES as f64,
            micros(time) / N_QUERIES as f64
        );
    }
}

/// What an in-memory store costs, stage by stage. The vectors are the floor:
/// everything above it is bookkeeping that could in principle be shared.
fn bench_footprint(n: usize) {
    let data = make_dataset(n, 0xF007 ^ n as u64);
    let mut rng = Rng(11);
    let dir = tempfile::TempDir::new().unwrap();
    let path = dir.path().join("footprint.h5");

    let base = heap_bytes();
    let entries: Vec<MemoryEntry> = data
        .vectors
        .iter()
        .enumerate()
        .map(|(i, v)| MemoryEntry {
            chunk: text_for(data.cluster_of[i], i, &mut rng),
            embedding: v.clone(),
            source_channel: "bench".into(),
            timestamp: i as f64,
            session_id: format!("s{}", i % 50),
            tags: format!("t{i}"),
        })
        .collect();
    let after_entries = heap_bytes();

    let mut config = MemoryConfig::new(path, "bench", DIM);
    config.quantized_index = INT8.load(std::sync::atomic::Ordering::Relaxed);
    let mut mem = HDF5Memory::create(config).unwrap();
    mem.save_batch(entries).unwrap();
    let after_store = heap_bytes();

    // First query builds the vector and keyword indexes.
    std::hint::black_box(mem.hybrid_search(&data.queries[0], "record", 0.7, 0.3, K));
    let after_indexes = heap_bytes();

    // Reopening is the figure that matters for a long-lived process, and the
    // only one RSS reports honestly: memory freed when the ingest buffers went
    // away stays in the allocator's pool, so the stage deltas above understate
    // what was given back.
    let path = mem.config().path.clone();
    drop(mem);
    let before_open = heap_bytes();
    reset_peak();
    let reopened = HDF5Memory::open(&path).unwrap();
    let after_open = heap_bytes();
    let loaded = after_open.saturating_sub(before_open);
    // Peak over the open, not just what it leaves behind: a buffer allocated
    // and freed during the parse never shows up in the live total.
    let peak = peak_bytes().saturating_sub(before_open);
    drop(reopened);

    let raw = (n * DIM * 4) as u64;
    println!(
        "| {n} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.2}x |",
        mib(raw),
        mib(after_entries.saturating_sub(base)),
        mib(after_store.saturating_sub(after_entries)),
        mib(after_indexes.saturating_sub(after_store)),
        mib(loaded),
        mib(peak),
        loaded as f64 / raw as f64,
    );
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let full = args.iter().any(|a| a == "--full");
    let ann_only = args.iter().any(|a| a == "--ann-only");
    if args.iter().any(|a| a == "--fusion-study") {
        for &n in if full {
            &[10_000, 100_000][..]
        } else {
            &[10_000][..]
        } {
            fusion_study(n);
        }
        return;
    }
    if args.iter().any(|a| a == "--signing-study") {
        println!("## Signed checkpoints ({DIM}-dim, float16, int8 index)\n");
        println!(
            "| N | checkpoint ms, unsigned | checkpoint ms, signed | signing adds ms | verify ms | file MiB added |"
        );
        println!("|---:|---:|---:|---:|---:|---:|");
        for &n in if full {
            &[1_000, 10_000, 100_000][..]
        } else {
            &[1_000, 10_000][..]
        } {
            signing_study(n);
        }
        return;
    }
    if args.iter().any(|a| a == "--options-study") {
        println!("## Search options ({DIM}-dim, k = {K}, Hebbian boost off)\n");
        println!("| N | options | filtered recall@10 | p50 ms | p99 ms |");
        println!("|---:|---|---:|---:|---:|");
        for &n in if full {
            &[10_000, 100_000][..]
        } else {
            &[10_000][..]
        } {
            options_study(n);
        }
        return;
    }
    if args.iter().any(|a| a == "--f16-first") {
        F16_FIRST.store(true, std::sync::atomic::Ordering::Relaxed);
    }
    if args.iter().any(|a| a == "--float16-study") {
        println!("## float16 embedding storage ({DIM}-dim, int8 index, Hebbian boost off)\n");
        println!(
            "| N | embeddings | file MiB | checkpoint ms | open ms | recall@10 | top-10 overlap with the other | hybrid p50 ms |"
        );
        println!("|---:|---|---:|---:|---:|---:|---:|---:|");
        for &n in if full {
            &[1_000, 10_000, 100_000][..]
        } else {
            &[1_000, 10_000][..]
        } {
            float16_study(n);
        }
        return;
    }
    if args.iter().any(|a| a == "--int8") {
        INT8.store(true, std::sync::atomic::Ordering::Relaxed);
        println!("(int8-quantised index vectors)");
    }
    if args.iter().any(|a| a == "--rerank") {
        RERANK.store(true, std::sync::atomic::Ordering::Relaxed);
        println!("(candidates re-scored against exact vectors)");
    }
    if args.iter().any(|a| a == "--uniform") {
        UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed);
        println!("(uniform random data)");
    }
    let json_path = args
        .iter()
        .position(|a| a == "--json")
        .and_then(|i| args.get(i + 1))
        .cloned();
    let sizes: &[usize] = if full {
        &[1_000, 10_000, 100_000]
    } else {
        &[1_000, 10_000]
    };

    if cfg!(debug_assertions) {
        eprintln!("warning: debug build — numbers are meaningless. Use --release.");
    }

    let mut json = Vec::new();
    println!("## Search harness");

    if args.iter().any(|a| a == "--footprint") {
        println!("\n### Resident memory, {DIM}-dim f32\n");
        println!(
            "| N | vectors (raw) | entries MiB | store MiB | indexes MiB | reopened MiB | peak during open MiB | reopened / raw |"
        );
        println!("|---:|---:|---:|---:|---:|---:|---:|---:|");
        for &n in sizes {
            bench_footprint(n);
        }
        return;
    }
    // `--e2e-only` skips the index benchmarks, so the end-to-end section runs
    // in a process that has not already spun up a thread pool.
    if !args.iter().any(|a| a == "--e2e-only") {
        for &n in sizes {
            bench_ann(n, &mut json);
        }
    }

    if ann_only {
        return;
    }
    println!("\n### End to end: `HDF5Memory::hybrid_search` (k = {K}, weights 0.7 / 0.3)\n");
    println!(
        "| N | ingest ms | cold index build ms | checkpoint ms | open ms | first query after open ms | p50 ms | p99 ms | QPS |"
    );
    println!("|---:|---:|---:|---:|---:|---:|---:|---:|---:|");
    for &n in sizes {
        bench_end_to_end(n, &mut json);
    }

    if let Some(path) = json_path {
        std::fs::write(&path, serde_json::to_string_pretty(&json).unwrap()).unwrap();
        eprintln!("wrote {path}");
    }
}