#include <benchmark/benchmark.h>
#include <random>
#include <stdint.h>

#include <roaring/roaring.hh>

constexpr size_t BENCHMARK_NUM_ROWS = 8192;
constexpr size_t BENCHMARK_NUM_ITERATIONS = 100;

using namespace roaring;

static Roaring createBitmapWithFilterRate(double filter_rate,
                                          uint32_t starting_row,
                                          size_t num_rows,
                                          size_t num_iterations,
                                          unsigned seed) {
    std::mt19937_64 rng(seed);
    std::uniform_real_distribution<double> dist(0.0, 1.0);

    Roaring bitmap;

    std::vector<uint32_t> values;
    values.reserve(static_cast<size_t>(num_rows * num_iterations * filter_rate *
                                       1.2));  // Reserve slightly more

    /// Generate values based on filter rate for multiple iterations
    for (size_t iter = 0; iter < num_iterations; ++iter) {
        uint32_t iter_offset = static_cast<uint32_t>(iter * num_rows);
        for (size_t i = 0; i < num_rows; ++i) {
            if (dist(rng) < filter_rate) {
                values.push_back(
                    static_cast<uint32_t>(starting_row + iter_offset + i));
            }
        }
    }

    if (!values.empty()) bitmap.addMany(values.size(), values.data());
    return bitmap;
}

void fillBoolArrayByIterator(const api::roaring_bitmap_t* bitmap, bool* array,
                             uint32_t starting_row, size_t ending_row) {
    api::roaring_uint32_iterator_t it;
    roaring_iterator_init(bitmap, &it);
    if (!roaring_uint32_iterator_move_equalorlarger(&it, starting_row)) return;
    while (it.current_value < ending_row) {
        array[it.current_value - starting_row] = true;
        if (!roaring_uint32_iterator_advance(&it)) break;
    }
}

// Helper macro to generate benchmark functions
// Each pair of benchmarks shares the same bitmap instance
#define BENCHMARK_FILTER_RATE(name, rate_value, label)                      \
    static Roaring g_bitmap_##name = createBitmapWithFilterRate(            \
        rate_value, 0, BENCHMARK_NUM_ROWS, BENCHMARK_NUM_ITERATIONS, 42);   \
                                                                            \
    static void BM_ToBoolArray_FilterRate_##name(benchmark::State& state) { \
        const size_t total_rows =                                           \
            BENCHMARK_NUM_ITERATIONS * BENCHMARK_NUM_ROWS;                  \
                                                                            \
        bool* array = new bool[total_rows];                                 \
        for (auto _ : state) {                                              \
            for (size_t i = 0; i < BENCHMARK_NUM_ITERATIONS; ++i) {         \
                uint32_t iter_start =                                       \
                    static_cast<uint32_t>(i * BENCHMARK_NUM_ROWS);          \
                api::roaring_bitmap_range_bool_array(                       \
                    &g_bitmap_##name.roaring, iter_start,                   \
                    iter_start + BENCHMARK_NUM_ROWS, array);                \
                benchmark::DoNotOptimize(array);                            \
            }                                                               \
            benchmark::DoNotOptimize(array);                                \
        }                                                                   \
        state.SetItemsProcessed(state.iterations() * total_rows);           \
        state.SetLabel("filter_rate=" label ",iterations=100");             \
        delete[] array;                                                     \
    }                                                                       \
                                                                            \
    static void BM_ToBoolArrayByIterator_FilterRate_##name(                 \
        benchmark::State& state) {                                          \
        const size_t total_rows =                                           \
            BENCHMARK_NUM_ITERATIONS * BENCHMARK_NUM_ROWS;                  \
                                                                            \
        bool* array = new bool[total_rows];                                 \
        for (auto _ : state) {                                              \
            for (size_t i = 0; i < BENCHMARK_NUM_ITERATIONS; ++i) {         \
                uint32_t iter_start =                                       \
                    static_cast<uint32_t>(i * BENCHMARK_NUM_ROWS);          \
                fillBoolArrayByIterator(&g_bitmap_##name.roaring, array,    \
                                        iter_start,                         \
                                        iter_start + BENCHMARK_NUM_ROWS);   \
                benchmark::DoNotOptimize(array);                            \
            }                                                               \
            benchmark::DoNotOptimize(array);                                \
        }                                                                   \
        state.SetItemsProcessed(state.iterations() * total_rows);           \
        state.SetLabel("filter_rate=" label ",iterations=100");             \
        delete[] array;                                                     \
    }                                                                       \
                                                                            \
    BENCHMARK(BM_ToBoolArray_FilterRate_##name);                            \
    BENCHMARK(BM_ToBoolArrayByIterator_FilterRate_##name)

// Generate benchmarks for different filter rates
BENCHMARK_FILTER_RATE(0_1, 0.001, "0.1%");
BENCHMARK_FILTER_RATE(1, 0.01, "1%");
BENCHMARK_FILTER_RATE(10, 0.10, "10%");
BENCHMARK_FILTER_RATE(50, 0.50, "50%");
BENCHMARK_FILTER_RATE(90, 0.90, "90%");
BENCHMARK_FILTER_RATE(99, 0.99, "99%");
BENCHMARK_MAIN();