Types

enum class Location {
    kRedis,      // In-memory storage via Redis
    kMemory,     // Temporary in-memory storage
    kPostgres,   // Relational database storage
    kNone        // Undefined storage type
};

LocationConfig(Location location,
                const std::optional<std::string>& table_name,
                const std::optional<std::string>& db_connection_string);

cyborg::LocationConfig index_loc(Location::kRedis, std::nullopt, "redis://localhost");
cyborg::LocationConfig config_loc(Location::kRedis, std::nullopt, "redis://localhost");
cyborg::LocationConfig items_loc(Location::kPostgres, "items", "host=localhost dbname=postgres");

enum class DistanceMetric {
    Cosine,
    Euclidean,
    SquaredEuclidean};

IndexIVF(size_t dimension,
         size_t n_lists,
         DistanceMetric metric = DistanceMetric::Euclidean);

IndexIVFFlat(size_t dimension,
             size_t n_lists,
             DistanceMetric metric = DistanceMetric::Euclidean);

IndexIVFPQ(size_t dimension,
           size_t n_lists,
           size_t pq_dim,
           size_t pq_bits,
           DistanceMetric metric = DistanceMetric::Euclidean);

Array2D(size_t rows, size_t cols, const T& initial_value = T());
Array2D(std::vector<T>&& data, size_t cols);
Array2D(const std::vector<T>& data, size_t cols);

// Converting a vector to an array
std::vector<uint8_t> vec = {0, 1, 2, 3, 4, 5, 6, 7};
cyborg::Array2D<uint8_t> arr(vec, 2);
// arr is now a 2D array of 4 rows and 2 columns, with the contents from vec

// Creating a 2D array with 3 rows and 2 columns, initialized to zero
cyborg::Array2D<int> array(3, 2, 0);

// Access and modify elements
array(0, 0) = 1;
array(0, 1) = 2;

// Printing the array
for (size_t i = 0; i < array.rows(); ++i) {
    for (size_t j = 0; j < array.cols(); ++j) {
        std::cout << array(i, j) << " ";
    }
    std::cout << std::endl;
}

TrainingConfig(size_t batch_size = 0,
                size_t max_iters = 0,
                double tolerance = 1e-6,
                size_t max_memory = 0);

QueryParams(size_t top_k = 100,
            size_t n_probes = 1,
            bool return_distances = true,
            bool greedy = false);

QueryResults results(num_queries, top_k);

// Access the top-k results for each query
for (size_t i = 0; i < num_queries; ++i) {
    auto result = results[i];
    for (size_t j = 0; j < result.num_results; ++j) {
        std::cout << "ID: " << result.ids[j] << ", Distance: " << result.distances[j] << std::endl;
    }
}

// Get the IDs and distances for all queries
auto all_ids = results.ids();
auto all_distances = results.distances();