// Copyright (c) Microsoft Corporation.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <algorithm>
#include <array>
#include <benchmark/benchmark.h>
#include <bitset>
#include <cstddef>
#include <random>
#include <sstream>
#include <string>

using namespace std;

namespace {
    template <size_t N, class charT, size_t Min_length>
    auto random_digits_init() {
        mt19937_64 rnd{};
        uniform_int_distribution<> dis('0', '1');

        constexpr size_t number_of_bitsets = (Min_length + N - 1) / N;
        static_assert(number_of_bitsets != 0);

        constexpr size_t actual_size = number_of_bitsets * (N + 1); // +1 for \0

        array<charT, actual_size> result;

        for (size_t i = 0; i < actual_size; ++i) {
            if (i % (N + 1) == N) {
                result[i] = charT{'\0'}; // write null terminators
            } else {
                result[i] = static_cast<charT>(dis(rnd)); // fill random digits
            }
        }

        return result;
    }

    enum class length_type : bool { char_count, null_term };

    template <size_t N, class charT>
    const auto random_digits = random_digits_init<N, charT, 2048>();

    template <length_type Length, size_t N, class charT>
    void bitset_from_string(benchmark::State& state) {
        const auto& digit_array = random_digits<N, charT>;
        for (auto _ : state) {
            benchmark::DoNotOptimize(digit_array);
            const auto arr_data = digit_array.data();
            const auto arr_size = digit_array.size();
            for (size_t pos = 0; pos != arr_size; pos += N + 1) {
                if constexpr (Length == length_type::char_count) {
                    bitset<N> bs(arr_data + pos, N);
                    benchmark::DoNotOptimize(bs);
                } else {
                    bitset<N> bs(arr_data + pos);
                    benchmark::DoNotOptimize(bs);
                }
            }
        }
    }

    template <class charT, size_t Length>
    basic_string<charT> random_digits_contiguous_string_init() {
        mt19937_64 rnd{};
        uniform_int_distribution<> dis('0', '1');

        basic_string<charT> result;

        result.resize_and_overwrite(Length, [&](charT* ptr, size_t) {
            generate_n(ptr, Length, [&] { return static_cast<charT>(dis(rnd)); });
            return Length;
        });

        return result;
    }

    template <class charT, size_t Length>
    const auto random_digits_contiguous_string = random_digits_contiguous_string_init<charT, Length>();

    template <size_t N, class charT>
    void bitset_from_stream(benchmark::State& state) {
        constexpr size_t string_length = 2048;
        constexpr size_t count         = string_length / N;
        basic_istringstream<charT> stream(random_digits_contiguous_string<charT, string_length>);
        bitset<N> bs;
        for (auto _ : state) {
            benchmark::DoNotOptimize(stream);
            for (size_t i = 0; i != count; ++i) {
                stream >> bs;
            }
            benchmark::DoNotOptimize(bs);
            stream.seekg(0);
        }
    }

} // namespace

BENCHMARK(bitset_from_string<length_type::char_count, 15, char>);
BENCHMARK(bitset_from_string<length_type::char_count, 16, char>);
BENCHMARK(bitset_from_string<length_type::char_count, 36, char>);
BENCHMARK(bitset_from_string<length_type::char_count, 64, char>);
BENCHMARK(bitset_from_string<length_type::char_count, 512, char>);
BENCHMARK(bitset_from_string<length_type::char_count, 2048, char>);

BENCHMARK(bitset_from_string<length_type::char_count, 15, wchar_t>);
BENCHMARK(bitset_from_string<length_type::char_count, 16, wchar_t>);
BENCHMARK(bitset_from_string<length_type::char_count, 36, wchar_t>);
BENCHMARK(bitset_from_string<length_type::char_count, 64, wchar_t>);
BENCHMARK(bitset_from_string<length_type::char_count, 512, wchar_t>);
BENCHMARK(bitset_from_string<length_type::char_count, 2048, wchar_t>);

BENCHMARK(bitset_from_string<length_type::null_term, 15, char>);
BENCHMARK(bitset_from_string<length_type::null_term, 16, char>);
BENCHMARK(bitset_from_string<length_type::null_term, 36, char>);
BENCHMARK(bitset_from_string<length_type::null_term, 64, char>);
BENCHMARK(bitset_from_string<length_type::null_term, 512, char>);
BENCHMARK(bitset_from_string<length_type::null_term, 2048, char>);

BENCHMARK(bitset_from_string<length_type::null_term, 15, wchar_t>);
BENCHMARK(bitset_from_string<length_type::null_term, 16, wchar_t>);
BENCHMARK(bitset_from_string<length_type::null_term, 36, wchar_t>);
BENCHMARK(bitset_from_string<length_type::null_term, 64, wchar_t>);
BENCHMARK(bitset_from_string<length_type::null_term, 512, wchar_t>);
BENCHMARK(bitset_from_string<length_type::null_term, 2048, wchar_t>);

BENCHMARK(bitset_from_stream<15, char>);
BENCHMARK(bitset_from_stream<16, char>);
BENCHMARK(bitset_from_stream<36, char>);
BENCHMARK(bitset_from_stream<64, char>);
BENCHMARK(bitset_from_stream<512, char>);
BENCHMARK(bitset_from_stream<2048, char>);

BENCHMARK(bitset_from_stream<15, wchar_t>);
BENCHMARK(bitset_from_stream<16, wchar_t>);
BENCHMARK(bitset_from_stream<36, wchar_t>);
BENCHMARK(bitset_from_stream<64, wchar_t>);
BENCHMARK(bitset_from_stream<512, wchar_t>);
BENCHMARK(bitset_from_stream<2048, wchar_t>);

BENCHMARK_MAIN();
