#include "../tests/testutil.hpp"

#include <trie.hpp>
#include <radix_tree.hpp>

#include <benchmark/benchmark.h>
#include <chrono>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <random>
#include <vector>

const std::size_t nkeys = 10000;
const std::size_t nqueries = 100000;
const std::size_t warmup_runs = 10;
const std::size_t samples = 10;
const std::size_t key_length = 20;
const char* chars = "abcdefghijklmnopqrstuvwxyz0123456789";
const int chars_len = 36;

static void BM_TrieLookup(benchmark::State& state)
{
    state.PauseTiming();
    std::minstd_rand rng(123456789);
    std::vector<unsigned char*> input_set;
    std::vector<unsigned char*> queries;
    input_set.reserve(nkeys);
    queries.reserve(nqueries);

    for (std::size_t i = 0; i < nkeys; ++i) {
        unsigned char* key = new unsigned char[key_length];
        for (std::size_t j = 0; j < key_length; j++)
            key[j] = static_cast<unsigned char>(chars[rng() % chars_len]);
        input_set.emplace_back(key);
    }
    for (std::size_t i = 0; i < nqueries; ++i)
        queries.push_back(input_set[rng() % nkeys]);

    zmq::trie_t trie;
    for (auto key : input_set) {
        trie.add(key, key_length);
    }

    state.ResumeTiming();
    for (auto _ : state) {
        for (auto query : queries)
            benchmark::DoNotOptimize(trie.check(query, key_length));
    }
}

static void BM_RadixTreeLookup(benchmark::State& state)
{
    state.PauseTiming();
    std::minstd_rand rng(123456789);
    std::vector<unsigned char*> input_set;
    std::vector<unsigned char*> queries;
    input_set.reserve(nkeys);
    queries.reserve(nqueries);

    for (std::size_t i = 0; i < nkeys; ++i) {
        unsigned char* key = new unsigned char[key_length];
        for (std::size_t j = 0; j < key_length; j++)
            key[j] = static_cast<unsigned char>(chars[rng() % chars_len]);
        input_set.emplace_back(key);
    }
    for (std::size_t i = 0; i < nqueries; ++i)
        queries.push_back(input_set[rng() % nkeys]);

    zmq::radix_tree tree;
    for (auto key : input_set) {
        tree.add(key, key_length);
    }

    state.ResumeTiming();
    for (auto _ : state) {
        for (auto query : queries)
            benchmark::DoNotOptimize(tree.check(query, key_length));
    }
}

// BENCHMARK(BM_TrieLookup)->Unit(benchmark::kMillisecond)->UseRealTime();
BENCHMARK(BM_RadixTreeLookup)->Unit(benchmark::kMillisecond);

BENCHMARK_MAIN();