#include <benchmark/benchmark.h>
#include <x86intrin.h>
#include <cstring>

#define ARR_SIZE 10240000
#define VAL 50

static int16_t *getRandArr() {
    auto res = new int16_t[ARR_SIZE];
    for (int i = 0; i < ARR_SIZE; ++i) {
        res[i] = static_cast<int16_t>(rand() % (VAL * 2));
    }
    return res;
}
static auto arr = getRandArr();

static int16_t *getAllignedArr() {
    auto res = aligned_alloc(32, sizeof(int16_t) * ARR_SIZE);
    memcpy(res, arr, sizeof(int16_t) * ARR_SIZE);
    return static_cast<int16_t *>(res);
}
static auto allignedArr = getAllignedArr();

static void BM_Count(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;
        for (int i = 0; i < ARR_SIZE; ++i)
            if (arr[i] == VAL)
                ++cnt;
        benchmark::DoNotOptimize(cnt);
    }
}

BENCHMARK(BM_Count);

static void BM_SSE_COUNT_SET_EPI(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto sseVal = _mm_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 8) {
            cnt += _popcnt32(
                    _mm_movemask_epi8(
                            _mm_cmpeq_epi16(
                                    sseVal,
                                    _mm_set_epi16(arr[i + 7], arr[i + 6], arr[i + 5], arr[i + 4],
                                                  arr[i + 3], arr[i + 2], arr[i + 1], arr[i])
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_SSE_COUNT_SET_EPI);

static void BM_SSE_COUNT_LOADU(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto sseVal = _mm_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 8) {
            cnt += _popcnt32(
                    _mm_movemask_epi8(
                            _mm_cmpeq_epi16(
                                    sseVal,
                                    _mm_loadu_si128((__m128i *) &arr[i])
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_SSE_COUNT_LOADU);

static void BM_SSE_COUNT_DIRECT(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto sseVal = _mm_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 8) {
            cnt += _popcnt32(
                    _mm_movemask_epi8(
                            _mm_cmpeq_epi16(
                                    sseVal,
                                    *(__m128i *) &allignedArr[i]
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_SSE_COUNT_DIRECT);

#ifdef __AVX2__

static void BM_AVX_COUNT_LOADU(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto avxVal = _mm256_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 16) {
            cnt += _popcnt32(
                    _mm256_movemask_epi8(
                            _mm256_cmpeq_epi16(
                                    avxVal,
                                    _mm256_loadu_si256((__m256i *) &arr[i])
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_AVX_COUNT_LOADU);

static void BM_AVX_COUNT_LOAD(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto avxVal = _mm256_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 16) {
            cnt += _popcnt32(
                    _mm256_movemask_epi8(
                            _mm256_cmpeq_epi16(avxVal,
                                               _mm256_load_si256((__m256i *) &allignedArr[i])
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_AVX_COUNT_LOAD);

static void BM_AVX_COUNT_DIRECT(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto avxVal = _mm256_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 16) {
            cnt += _popcnt32(
                    _mm256_movemask_epi8(
                            _mm256_cmpeq_epi16(
                                    avxVal,
                                    *(__m256i *) &allignedArr[i]
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_AVX_COUNT_DIRECT);

#endif

BENCHMARK_MAIN();