qdrant-edge 0.8.0

use std::fmt::Debug;
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;

use atomic_refcell::AtomicRefCell;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::types::PointOffsetType;
use crate::common::universal_io::{MmapFile, MmapFs};
use rand::Rng;
use crate::sparse::common::sparse_vector::SparseVector;
use crate::sparse::common::sparse_vector_fixture::random_sparse_vector;
use crate::sparse::index::inverted_index::InvertedIndexReadWrite;

use crate::segment::common::operation_error::OperationResult;
use crate::segment::fixtures::payload_context_fixture::create_id_tracker_fixture;
use crate::segment::index::VectorIndexRead;
use crate::segment::index::sparse_index::sparse_index_config::{SparseIndexConfig, SparseIndexType};
use crate::segment::index::sparse_index::sparse_vector_index::{
    SparseVectorIndex, SparseVectorIndexOpenArgs,
};
use crate::segment::index::struct_payload_index::{IndexLoadMode, StorageType, StructPayloadIndex};
use crate::segment::payload_storage::in_memory_payload_storage::InMemoryPayloadStorage;
use crate::segment::vector_storage::sparse::mmap_sparse_vector_storage::MmapSparseVectorStorage;
use crate::segment::vector_storage::{VectorStorage, VectorStorageEnum, VectorStorageRead};

/// Prepares a sparse vector index with a given iterator of sparse vectors
pub fn fixture_sparse_index_from_iter<I: InvertedIndexReadWrite<MmapFile>>(
    data_dir: &Path,
    vectors: impl ExactSizeIterator<Item = SparseVector>,
    full_scan_threshold: usize,
    index_type: SparseIndexType,
) -> OperationResult<SparseVectorIndex<I>> {
    let stopped = AtomicBool::new(false);

    // directories
    let index_dir = &data_dir.join("index");
    let payload_dir = &data_dir.join("payload");
    let storage_dir = &data_dir.join("storage");

    // setup
    let id_tracker = Arc::new(AtomicRefCell::new(create_id_tracker_fixture(vectors.len())));
    let payload_storage = InMemoryPayloadStorage::default();
    let wrapped_payload_storage = Arc::new(AtomicRefCell::new(payload_storage.into()));
    let payload_index = StructPayloadIndex::open(
        wrapped_payload_storage,
        id_tracker.clone(),
        std::collections::HashMap::new(),
        payload_dir,
        StorageType::Appendable,
        IndexLoadMode::CreateIfMissing,
    )?;
    let wrapped_payload_index = Arc::new(AtomicRefCell::new(payload_index));

    let vector_storage = Arc::new(AtomicRefCell::new(VectorStorageEnum::SparseMmap(
        MmapSparseVectorStorage::open_or_create(storage_dir)?,
    )));
    let mut borrowed_storage = vector_storage.borrow_mut();

    let num_vectors = vectors.len();
    let mut num_vectors_not_empty = 0;
    let hw_counter = HardwareCounterCell::new();
    for (idx, vec) in vectors.enumerate() {
        borrowed_storage
            .insert_vector(idx as PointOffsetType, (&vec).into(), &hw_counter)
            .unwrap();
        num_vectors_not_empty += usize::from(!vec.is_empty());
    }
    drop(borrowed_storage);

    // assert all empty points are in storage
    assert_eq!(
        vector_storage.borrow().available_vector_count(),
        num_vectors,
    );

    let sparse_index_config =
        SparseIndexConfig::new(Some(full_scan_threshold), index_type, None, None);
    let sparse_vector_index: SparseVectorIndex<I> =
        SparseVectorIndex::open(SparseVectorIndexOpenArgs {
            fs: &MmapFs,
            config: sparse_index_config,
            id_tracker,
            vector_storage: vector_storage.clone(),
            payload_index: wrapped_payload_index,
            path: index_dir,
            stopped: &stopped,
            tick_progress: || (),
        })?;

    assert_eq!(
        sparse_vector_index.indexed_vector_count(),
        num_vectors_not_empty
    );

    Ok(sparse_vector_index)
}

/// Prepares a sparse vector index with random sparse vectors
pub fn fixture_sparse_index<I: InvertedIndexReadWrite<MmapFile> + Debug, R: Rng + ?Sized>(
    rnd: &mut R,
    num_vectors: usize,
    max_dim: usize,
    full_scan_threshold: usize,
    data_dir: &Path,
) -> SparseVectorIndex<I> {
    fixture_sparse_index_from_iter(
        data_dir,
        (0..num_vectors).map(|_| random_sparse_vector(rnd, max_dim)),
        full_scan_threshold,
        SparseIndexType::ImmutableRam,
    )
    .unwrap()
}

macro_rules! fixture_for_all_indices {
    ($test:ident::<_>($($args:tt)*)) => {
        eprintln!("InvertedIndexCompressedImmutableRam<f32>");
        $test::<
            ::sparse::index::inverted_index::inverted_index_compressed_immutable_ram::InvertedIndexCompressedImmutableRam<f32>
        >($($args)*);

        eprintln!("InvertedIndexCompressedMmap<f32, MmapFile>");
        $test::<
            ::sparse::index::inverted_index::inverted_index_compressed_mmap::InvertedIndexCompressedMmap<f32, MmapFile>
        >($($args)*);
    };
}