/// @file dpf/parallel_bit_iterable_helpers.hpp
/// @author Christopher Jiang <christopher.jiang@ucalgary.ca>
/// @brief
/// @details
/// @copyright Copyright (c) 2019-2024 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_PARALLEL_BIT_ITERABLE_HELPERS_HPP__
#define LIBDPF_INCLUDE_DPF_PARALLEL_BIT_ITERABLE_HELPERS_HPP__

#include "hedley/hedley.h"

#include <cstddef>
#include <cstring>
#include <limits>
#include <array>
#include <type_traits>

#include "simde/simde/x86/avx2.h"
#include "portable-snippets/exact-int/exact-int.h"

#include "dpf/bit_array.hpp"

namespace dpf
{

namespace
{

/// Unaligned 256-bit load of `words_per_vec` words starting at `offset`.
/// Words past `nwords` are zero so a short batch does not read off the end.
template <typename Word>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
simde__m256i loadu_word_vec(const Word * words, std::size_t nwords,
    std::size_t offset) noexcept
{
    constexpr std::size_t words_per_vec = sizeof(simde__m256i) / sizeof(Word);
    static_assert(sizeof(simde__m256i) % sizeof(Word) == 0,
        "word size must divide a 256-bit vector");
    if (offset >= nwords)
    {
        return simde_mm256_setzero_si256();
    }
    if (offset + words_per_vec <= nwords)
    {
        return simde_mm256_loadu_si256(static_cast<const void *>(words + offset));
    }
    alignas(simde__m256i) Word tmp[words_per_vec]{};
    std::memcpy(tmp, words + offset, (nwords - offset) * sizeof(Word));
    return simde_mm256_loadu_si256(static_cast<const void *>(tmp));
}

template <std::size_t batch_size_log_2, typename ChildT>
struct parallel_bit_iterable_helper;

/// @brief for batch_size in 1..4
template <typename ChildT>
struct parallel_bit_iterable_helper<2, ChildT>
{
  public:
    using word_type = typename dpf::bit_array_base<ChildT>::word_type;
    using word_pointer = typename dpf::bit_array_base<ChildT>::word_pointer;
    using simde_type = simde__m256i;
    using simde_ptr = simde_type *;
    using element_type = psnip_uint64_t;
    static constexpr auto bits_per_word = dpf::bit_array_base<ChildT>::bits_per_word;
    static constexpr auto bits_per_element = std::numeric_limits<element_type>::digits;
    static constexpr auto elements_per_word = bits_per_word / bits_per_element;
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    using simde_array = std::array<simde_type, elements_per_word>;
HEDLEY_PRAGMA(GCC diagnostic pop)
    static constexpr auto left_shift = simde_mm256_slli_epi64;
    static constexpr auto right_shift = simde_mm256_srli_epi64;
    static constexpr auto bit_and = simde_mm256_and_si256;
    HEDLEY_NO_THROW
    static auto get_mask() noexcept
    {
        return simde_mm256_set1_epi64x(1);
    }
    HEDLEY_NO_THROW
    static simde_array build_vecs(const word_type * cur_word, std::size_t nwords) noexcept
    {
        return { loadu_word_vec(cur_word, nwords, 0) };
    }
};  // struct parallel_bit_iterable_helper<2>

/// @brief for batch_size in 5..8
template <typename ChildT>
struct parallel_bit_iterable_helper<3, ChildT>
{
  public:
    using word_pointer = typename dpf::bit_array_base<ChildT>::word_pointer;
    using simde_type = simde__m256i;
    using simde_ptr = simde_type *;
    using element_type = psnip_uint32_t;
    static constexpr auto bits_per_word = dpf::bit_array_base<ChildT>::bits_per_word;
    static constexpr auto bits_per_element = std::numeric_limits<element_type>::digits;
    static constexpr auto elements_per_word = bits_per_word / bits_per_element;
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    using simde_array = std::array<simde_type, elements_per_word>;
HEDLEY_PRAGMA(GCC diagnostic pop)
    static constexpr auto left_shift = simde_mm256_slli_epi64;
    static constexpr auto right_shift = simde_mm256_srli_epi64;
    static constexpr auto bit_and = simde_mm256_and_si256;
    HEDLEY_NO_THROW
    static auto get_mask() noexcept
    {
        return simde_mm256_set1_epi32(1);
    }
    HEDLEY_NO_THROW
    static simde_array build_vecs(const typename dpf::bit_array_base<ChildT>::word_type * cur_word,
        std::size_t nwords) noexcept
    {
        using word_type_ = typename dpf::bit_array_base<ChildT>::word_type;
        simde_type vec1 = loadu_word_vec<word_type_>(cur_word, nwords, 0);
        simde_type vec2 = loadu_word_vec<word_type_>(cur_word, nwords, 4);

        vec1 = simde_mm256_permutevar8x32_epi32(vec1, mask1);
        vec2 = simde_mm256_permutevar8x32_epi32(vec2, mask1);

        return {
            simde_mm256_permute2x128_si256(vec1, vec2, 0b00100000),
            simde_mm256_permute2x128_si256(vec1, vec2, 0b00110001) };
    }
  private:
    static constexpr simde_type mask1 = {
        0b0000000000000000000000000000001000000000000000000000000000000000,
        0b0000000000000000000000000000011000000000000000000000000000000100,
        0b0000000000000000000000000000001100000000000000000000000000000001,
        0b0000000000000000000000000000011100000000000000000000000000000101};
};  // struct parallel_bit_iterable_helper<3>

/// @brief for batch_size in 9..16
template <typename ChildT>
struct parallel_bit_iterable_helper<4, ChildT>
{
  public:
    using word_pointer = typename dpf::bit_array_base<ChildT>::word_pointer;
    using simde_type = simde__m256i;
    using simde_ptr = simde_type *;
    using element_type = psnip_uint16_t;
    static constexpr auto bits_per_word = dpf::bit_array_base<ChildT>::bits_per_word;
    static constexpr auto bits_per_element = std::numeric_limits<element_type>::digits;
    static constexpr auto elements_per_word = bits_per_word / bits_per_element;
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    using simde_array = std::array<simde_type, elements_per_word>;
HEDLEY_PRAGMA(GCC diagnostic pop)
    static constexpr auto left_shift = simde_mm256_slli_epi64;
    static constexpr auto right_shift = simde_mm256_srli_epi64;
    static constexpr auto bit_and = simde_mm256_and_si256;
    HEDLEY_NO_THROW
    static auto get_mask() noexcept
    {
        return simde_mm256_set1_epi16(1);
    }
    HEDLEY_NO_THROW
    static simde_array build_vecs(const typename dpf::bit_array_base<ChildT>::word_type * cur_word,
        std::size_t nwords) noexcept
    {
        using word_type_ = typename dpf::bit_array_base<ChildT>::word_type;
        simde_type vec1 = loadu_word_vec<word_type_>(cur_word, nwords, 0);
        simde_type vec2 = loadu_word_vec<word_type_>(cur_word, nwords, 4);
        simde_type vec3 = loadu_word_vec<word_type_>(cur_word, nwords, 8);
        simde_type vec4 = loadu_word_vec<word_type_>(cur_word, nwords, 12);

        vec1 = simde_mm256_permutevar8x32_epi32(vec1, mask1);
        vec2 = simde_mm256_permutevar8x32_epi32(vec2, mask1);
        vec3 = simde_mm256_permutevar8x32_epi32(vec3, mask1);
        vec4 = simde_mm256_permutevar8x32_epi32(vec4, mask1);

        vec1 = simde_mm256_shuffle_epi8(vec1, mask2);
        vec2 = simde_mm256_shuffle_epi8(vec2, mask2);
        vec3 = simde_mm256_shuffle_epi8(vec3, mask2);
        vec4 = simde_mm256_shuffle_epi8(vec4, mask2);

        simde_type bld_0002 = simde_mm256_unpacklo_epi64(vec1, vec2);
        simde_type bld_0406 = simde_mm256_unpacklo_epi64(vec3, vec4);
        simde_type bld_0103 = simde_mm256_unpackhi_epi64(vec1, vec2);
        simde_type bld_0507 = simde_mm256_unpackhi_epi64(vec3, vec4);

        return {
            simde_mm256_permute2x128_si256(bld_0002, bld_0406, 0b00100000),
            simde_mm256_permute2x128_si256(bld_0103, bld_0507, 0b00100000),
            simde_mm256_permute2x128_si256(bld_0002, bld_0406, 0b00110001),
            simde_mm256_permute2x128_si256(bld_0103, bld_0507, 0b00110001) };
    }
  private:
    static constexpr simde_type mask1 = {
        0b0000000000000000000000000000001000000000000000000000000000000000,
        0b0000000000000000000000000000011000000000000000000000000000000100,
        0b0000000000000000000000000000001100000000000000000000000000000001,
        0b0000000000000000000000000000011100000000000000000000000000000101};
    static constexpr simde_type mask2 = {
        0b0000110100001100000010010000100000000101000001000000000100000000,
        0b0000111100001110000010110000101000000111000001100000001100000010,
        0b0000110100001100000010010000100000000101000001000000000100000000,
        0b0000111100001110000010110000101000000111000001100000001100000010};
};  // struct parallel_bit_iterable_helper<4>

/// @brief for batch_size in 17..32
template <typename ChildT>
struct parallel_bit_iterable_helper<5, ChildT>
{
  public:
    using word_pointer = typename dpf::bit_array_base<ChildT>::word_pointer;
    using simde_type = simde__m256i;
    using simde_ptr = simde_type *;
    using element_type = psnip_uint8_t;
    static constexpr auto bits_per_word = dpf::bit_array_base<ChildT>::bits_per_word;
    static constexpr auto bits_per_element = std::numeric_limits<element_type>::digits;
    static constexpr auto elements_per_word = bits_per_word / bits_per_element;
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    using simde_array = std::array<simde_type, elements_per_word>;
HEDLEY_PRAGMA(GCC diagnostic pop)
    static constexpr auto left_shift = simde_mm256_slli_epi64;
    static constexpr auto right_shift = simde_mm256_srli_epi64;
    static constexpr auto bit_and = simde_mm256_and_si256;
    HEDLEY_NO_THROW
    static auto get_mask() noexcept
    {
        return simde_mm256_set1_epi8(1);
    }
    HEDLEY_NO_THROW
    static simde_array build_vecs(const typename dpf::bit_array_base<ChildT>::word_type * cur_word,
        std::size_t nwords) noexcept
    {
        using word_type_ = typename dpf::bit_array_base<ChildT>::word_type;
        simde_type vec1 = loadu_word_vec<word_type_>(cur_word, nwords, 0);
        simde_type vec2 = loadu_word_vec<word_type_>(cur_word, nwords, 4);
        simde_type vec3 = loadu_word_vec<word_type_>(cur_word, nwords, 8);
        simde_type vec4 = loadu_word_vec<word_type_>(cur_word, nwords, 12);
        simde_type vec5 = loadu_word_vec<word_type_>(cur_word, nwords, 16);
        simde_type vec6 = loadu_word_vec<word_type_>(cur_word, nwords, 20);
        simde_type vec7 = loadu_word_vec<word_type_>(cur_word, nwords, 24);
        simde_type vec8 = loadu_word_vec<word_type_>(cur_word, nwords, 28);

        vec1 = simde_mm256_permutevar8x32_epi32(vec1, mask1);
        vec2 = simde_mm256_permutevar8x32_epi32(vec2, mask1);
        vec3 = simde_mm256_permutevar8x32_epi32(vec3, mask1);
        vec4 = simde_mm256_permutevar8x32_epi32(vec4, mask1);
        vec5 = simde_mm256_permutevar8x32_epi32(vec5, mask1);
        vec6 = simde_mm256_permutevar8x32_epi32(vec6, mask1);
        vec7 = simde_mm256_permutevar8x32_epi32(vec7, mask1);
        vec8 = simde_mm256_permutevar8x32_epi32(vec8, mask1);

        vec1 = simde_mm256_shuffle_epi8(vec1, mask2);
        vec2 = simde_mm256_shuffle_epi8(vec2, mask2);
        vec3 = simde_mm256_shuffle_epi8(vec3, mask2);
        vec4 = simde_mm256_shuffle_epi8(vec4, mask2);
        vec5 = simde_mm256_shuffle_epi8(vec5, mask2);
        vec6 = simde_mm256_shuffle_epi8(vec6, mask2);
        vec7 = simde_mm256_shuffle_epi8(vec7, mask2);
        vec8 = simde_mm256_shuffle_epi8(vec8, mask2);

        simde_type bld_00010405 = simde_mm256_unpacklo_epi32(vec1, vec2);
        simde_type bld_08091213 = simde_mm256_unpacklo_epi32(vec3, vec4);
        simde_type bld_16172021 = simde_mm256_unpacklo_epi32(vec5, vec6);
        simde_type bld_24252829 = simde_mm256_unpacklo_epi32(vec7, vec8);
        simde_type bld_02030607 = simde_mm256_unpackhi_epi32(vec1, vec2);
        simde_type bld_10111415 = simde_mm256_unpackhi_epi32(vec3, vec4);
        simde_type bld_18192223 = simde_mm256_unpackhi_epi32(vec5, vec6);
        simde_type bld_26273031 = simde_mm256_unpackhi_epi32(vec7, vec8);

        simde_type bld_0004 = simde_mm256_unpacklo_epi64(bld_00010405, bld_08091213);
        simde_type bld_1620 = simde_mm256_unpacklo_epi64(bld_16172021, bld_24252829);
        simde_type bld_0206 = simde_mm256_unpacklo_epi64(bld_02030607, bld_10111415);
        simde_type bld_1822 = simde_mm256_unpacklo_epi64(bld_18192223, bld_26273031);
        simde_type bld_0105 = simde_mm256_unpackhi_epi64(bld_00010405, bld_08091213);
        simde_type bld_1721 = simde_mm256_unpackhi_epi64(bld_16172021, bld_24252829);
        simde_type bld_0307 = simde_mm256_unpackhi_epi64(bld_02030607, bld_10111415);
        simde_type bld_1923 = simde_mm256_unpackhi_epi64(bld_18192223, bld_26273031);

        return {
            simde_mm256_permute2x128_si256(bld_0004, bld_1620, 0b00100000),
            simde_mm256_permute2x128_si256(bld_0105, bld_1721, 0b00100000),
            simde_mm256_permute2x128_si256(bld_0206, bld_1822, 0b00100000),
            simde_mm256_permute2x128_si256(bld_0307, bld_1923, 0b00100000),
            simde_mm256_permute2x128_si256(bld_0004, bld_1620, 0b00110001),
            simde_mm256_permute2x128_si256(bld_0105, bld_1721, 0b00110001),
            simde_mm256_permute2x128_si256(bld_0206, bld_1822, 0b00110001),
            simde_mm256_permute2x128_si256(bld_0307, bld_1923, 0b00110001) };
    }
  private:
    static constexpr simde_type mask1 = {
        0b0000000000000000000000000000001000000000000000000000000000000000,
        0b0000000000000000000000000000011000000000000000000000000000000100,
        0b0000000000000000000000000000001100000000000000000000000000000001,
        0b0000000000000000000000000000011100000000000000000000000000000101};
    static constexpr simde_type mask2 = {
        0b0000110100001001000001010000000100001100000010000000010000000000,
        0b0000111100001011000001110000001100001110000010100000011000000010,
        0b0000110100001001000001010000000100001100000010000000010000000000,
        0b0000111100001011000001110000001100001110000010100000011000000010};
};  // struct parallel_bit_iterable_helper<5>

}  // anonymous namespace

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PARALLEL_BIT_ITERABLE_HELPERS_HPP__