#include <gtest/gtest.h>

#include "dpf.hpp"

#include <array>
#include <cstdint>
#include <cstring>
#include <type_traits>
#include <vector>

namespace
{

template <typename LaneT>
using storage = dpf::leaf_storage_t<LaneT>;

template <typename LaneT>
constexpr unsigned lane_bits()
{
    if constexpr (dpf::utils::is_packed_subbyte_v<LaneT>)
        return static_cast<unsigned>(dpf::utils::packed_lane_bits_v<LaneT>);
    else
        return static_cast<unsigned>(sizeof(LaneT) * 8u);
}

template <typename LaneT>
constexpr std::uint64_t lane_mask()
{
    if constexpr (dpf::utils::is_packed_subbyte_v<LaneT>)
        return (std::uint64_t{1} << dpf::utils::packed_lane_bits_v<LaneT>) - 1u;
    else if constexpr (sizeof(LaneT) >= 8)
        return ~std::uint64_t{0};
    else
        return (std::uint64_t{1} << (sizeof(LaneT) * 8u)) - 1u;
}

template <typename LaneT>
std::uint64_t get_lane(const storage<LaneT> * buf, std::size_t i)
{
    if constexpr (dpf::utils::is_packed_subbyte_v<LaneT>)
    {
        constexpr unsigned w = dpf::utils::packed_lane_bits_v<LaneT>;
        const std::size_t bit = i * w;
        const std::size_t word = bit / 64u;
        const unsigned shift = static_cast<unsigned>(bit % 64u);
        std::uint64_t v = buf[word] >> shift;
        if (shift + w > 64u)
            v |= buf[word + 1u] << (64u - shift);
        return v & lane_mask<LaneT>();
    }
    else if constexpr (std::is_class_v<LaneT>)
    {
        return static_cast<std::uint64_t>(
            static_cast<typename LaneT::integral_type>(buf[i])) & lane_mask<LaneT>();
    }
    else
    {
        return static_cast<std::uint64_t>(buf[i]) & lane_mask<LaneT>();
    }
}

template <typename LaneT>
void set_lane(storage<LaneT> * buf, std::size_t i, std::uint64_t val)
{
    val &= lane_mask<LaneT>();
    if constexpr (dpf::utils::is_packed_subbyte_v<LaneT>)
    {
        constexpr unsigned w = dpf::utils::packed_lane_bits_v<LaneT>;
        const std::size_t bit = i * w;
        const std::size_t word = bit / 64u;
        const unsigned shift = static_cast<unsigned>(bit % 64u);
        const std::uint64_t mask = lane_mask<LaneT>();
        buf[word] = (buf[word] & ~(mask << shift)) | (val << shift);
        if (shift + w > 64u)
        {
            const unsigned lo = 64u - shift;
            buf[word + 1u] = (buf[word + 1u] & ~(mask >> lo)) | (val >> lo);
        }
    }
    else
    {
        buf[i] = static_cast<storage<LaneT>>(val);
    }
}

template <typename LaneT>
std::vector<storage<LaneT>> make_key(std::size_t nleaves, std::size_t key_id)
{
    std::vector<storage<LaneT>> v(dpf::leaf_storage_words<LaneT>(nleaves),
        storage<LaneT>{0});
    for (std::size_t i = 0; i < nleaves; ++i)
    {
        // Distinct pattern per (key, leaf); fits every width's mask.
        const std::uint64_t val
            = (key_id * 131u + i * 17u + 3u) & lane_mask<LaneT>();
        set_lane<LaneT>(v.data(), i, val);
    }
    return v;
}

template <typename LaneT>
void check_roundtrip(std::size_t nkeys, std::size_t nleaves)
{
    std::vector<std::vector<storage<LaneT>>> owned;
    owned.reserve(nkeys);
    std::vector<const storage<LaneT> *> in_ptrs;
    std::vector<storage<LaneT> *> out_ptrs;
    in_ptrs.reserve(nkeys);
    out_ptrs.reserve(nkeys);
    for (std::size_t k = 0; k < nkeys; ++k)
    {
        owned.push_back(make_key<LaneT>(nleaves, k));
        in_ptrs.push_back(owned.back().data());
    }

    std::vector<storage<LaneT>> interleaved(
        dpf::leaf_storage_words<LaneT>(nkeys * nleaves), storage<LaneT>{0});
    dpf::interleave_leaves<LaneT>(interleaved.data(), in_ptrs.data(), nkeys,
        nleaves);

    for (std::size_t i = 0; i < nleaves; ++i)
    {
        for (std::size_t k = 0; k < nkeys; ++k)
        {
            const std::size_t g = dpf::cohort_index(i, k, nkeys);
            EXPECT_EQ(get_lane<LaneT>(interleaved.data(), g),
                get_lane<LaneT>(owned[k].data(), i))
                << "LaneT bits=" << lane_bits<LaneT>() << " nkeys=" << nkeys
                << " nleaves=" << nleaves << " i=" << i << " k=" << k;
        }
    }

    std::vector<std::vector<storage<LaneT>>> round;
    round.reserve(nkeys);
    for (std::size_t k = 0; k < nkeys; ++k)
    {
        round.emplace_back(dpf::leaf_storage_words<LaneT>(nleaves),
            storage<LaneT>{0});
        out_ptrs.push_back(round.back().data());
    }
    dpf::deinterleave_leaves<LaneT>(out_ptrs.data(), interleaved.data(), nkeys,
        nleaves);

    for (std::size_t k = 0; k < nkeys; ++k)
    {
        for (std::size_t i = 0; i < nleaves; ++i)
        {
            EXPECT_EQ(get_lane<LaneT>(round[k].data(), i),
                get_lane<LaneT>(owned[k].data(), i))
                << "deinterleave bits=" << lane_bits<LaneT>()
                << " nkeys=" << nkeys << " i=" << i << " k=" << k;
        }
    }
}

}  // namespace

template <typename LaneT>
class InterleaveLeavesTest : public ::testing::Test
{
};

using InterleaveTypes = ::testing::Types<
    std::uint64_t,
    std::uint32_t,
    std::uint16_t,
    std::uint8_t,
    dpf::nyble,
    dpf::twobit,
    dpf::bit,
    dpf::gf2,
    dpf::gf22,
    dpf::gf24,
    dpf::gf28,
    dpf::gf216,
    dpf::gf232,
    dpf::gf264>;

TYPED_TEST_SUITE(InterleaveLeavesTest, InterleaveTypes);

TYPED_TEST(InterleaveLeavesTest, KeyCountsAndOddLength)
{
    // Length not a multiple of pack width (8 for bits-in-byte, 32 for
    // 2-bit-in-u64, 16 for 4-bit-in-u64, etc.): 13 is odd for all of those.
    constexpr std::size_t nleaves = 13;
    for (std::size_t nkeys : {std::size_t{1}, std::size_t{3}, std::size_t{5},
             std::size_t{16}})
        check_roundtrip<TypeParam>(nkeys, nleaves);
}

TYPED_TEST(InterleaveLeavesTest, LongerNotMultipleOfWord)
{
    // 70 leaves: not a multiple of 64 (bit word) or 32 (twobit word).
    check_roundtrip<TypeParam>(5, 70);
    check_roundtrip<TypeParam>(8, 70);
}

TEST(InterleaveLeaves, BitPackingIsKeyMajorWithinLeaf)
{
    // Three keys, five 1-bit leaves. Logical order of bits in the output
    // stream is (i0,k0),(i0,k1),(i0,k2),(i1,k0),... — not a byte index of
    // i*m+k into a uint8_t array.
    constexpr std::size_t nkeys = 3;
    constexpr std::size_t nleaves = 5;
    std::array<std::uint64_t, 1> k0{0}, k1{0}, k2{0};
    // key0: 1 0 1 1 0
    // key1: 0 1 1 0 1
    // key2: 1 1 0 0 1
    k0[0] = 0b01101ull;
    k1[0] = 0b10110ull;
    k2[0] = 0b10011ull;
    const std::uint64_t * keys[3] = {k0.data(), k1.data(), k2.data()};
    std::uint64_t out = 0;
    dpf::interleave_leaves<dpf::bit>(&out, keys, nkeys, nleaves);

    // Expected 15 bits, low-first: for i=0..4, bits of keys 0,1,2
    // i0: 1,0,1 | i1: 0,1,1 | i2: 1,1,0 | i3: 1,0,0 | i4: 0,1,1
    // bit0..14 = 1,0,1,0,1,1,1,1,0,1,0,0,0,1,1
    const std::uint64_t expect = 0b110001011110101ull;
    EXPECT_EQ(out & ((1ull << 15) - 1ull), expect);
}

TEST(InterleaveLeaves, InnerProductModintFromBits)
{
    // Same three keys as above. Leaf integers (key 0 in the low bit):
    // i0 = 0b101 = 5, i1 = 0b110 = 6, i2 = 0b011 = 3, i3 = 0b001 = 1, i4 = 0b110 = 6.
    constexpr std::size_t nkeys = 3;
    constexpr std::size_t nleaves = 5;
    std::array<std::uint64_t, 1> k0{0b01101ull}, k1{0b10110ull}, k2{0b10011ull};
    const std::uint64_t * keys[3] = {k0.data(), k1.data(), k2.data()};
    std::uint64_t bits = 0;
    dpf::interleave_leaves<dpf::bit>(&bits, keys, nkeys, nleaves);

    const std::uint64_t w[5] = {1, 2, 3, 4, 5};
    const auto got = dpf::interleaved_bits_inner_product<3>(&bits, nleaves, nkeys, w);
    // 5*1 + 6*2 + 3*3 + 1*4 + 6*5 = 5+12+9+4+30 = 60 ≡ 60-7*8 = 4 (mod 8)
    EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<3>::integral_type>(got)), 4u);

    std::array<dpf::modint<3>, 5> mw{1, 2, 3, 4, 5};
    const auto got_m = dpf::interleaved_bits_inner_product<3>(&bits, nleaves, nkeys, mw);
    EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<3>::integral_type>(got_m)), 4u);
}

TEST(InterleaveLeaves, InnerProductModint128)
{
    constexpr std::size_t nkeys = 128;
    constexpr std::size_t nleaves = 3;
    std::vector<std::uint64_t> key_bits(nkeys * 2, 0);
    std::vector<const std::uint64_t *> ptrs(nkeys);
    for (std::size_t k = 0; k < nkeys; ++k)
    {
        // Leaf 0: bit k set. Leaf 1: only key 0. Leaf 2: zero.
        if (k == 0)
            key_bits[k * 2] = 0b011ull;
        else
            key_bits[k * 2] = 0b001ull;
        ptrs[k] = key_bits.data() + k * 2;
    }
    std::vector<std::uint64_t> bits(dpf::leaf_storage_words<dpf::bit>(nkeys * nleaves));
    dpf::interleave_leaves<dpf::bit>(bits.data(), ptrs.data(), nkeys, nleaves);

    const std::uint64_t w[3] = {1, 3, 9};
    const auto got = dpf::interleaved_bits_inner_product<128>(
        bits.data(), nleaves, nkeys, w);
    using limb = dpf::modint<128>::integral_type;
    const limb all = ~limb{0};
    const limb expect = all * limb{1} + limb{1} * limb{3};
    EXPECT_EQ(static_cast<limb>(got), expect);
}

TEST(InterleaveLeaves, InnerProductLowBitsOfWideGroup)
{
    // Eight 1-bit keys, read as modint<4>: only keys 0..3.
    constexpr std::size_t nkeys = 8;
    constexpr std::size_t nleaves = 2;
    std::array<std::uint64_t, 8> raw{};
    const std::uint64_t * ptrs[8];
    for (std::size_t k = 0; k < nkeys; ++k)
    {
        raw[k] = (k < 4) ? 0b11ull : 0b01ull;  // both leaves set for low keys
        ptrs[k] = &raw[k];
    }
    std::uint64_t bits[4]{};
    dpf::interleave_leaves<dpf::bit>(bits, ptrs, nkeys, nleaves);
    const std::uint64_t w[2] = {1, 1};
    const auto got = dpf::interleaved_bits_inner_product<4>(bits, nleaves, nkeys, w);
    // Each leaf's low 4 bits are 0b1111 = 15. 15+15 = 30 ≡ 14 (mod 16).
    EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<4>::integral_type>(got)), 14u);
}

TYPED_TEST(InterleaveLeavesTest, EmptyAndSingleLeaf)
{
    check_roundtrip<TypeParam>(0, 0);
    check_roundtrip<TypeParam>(3, 0);
    check_roundtrip<TypeParam>(0, 5);
    check_roundtrip<TypeParam>(1, 1);
    check_roundtrip<TypeParam>(4, 1);
}

TEST(InterleaveLeaves, InnerProductEmptyIsZero)
{
    const std::uint64_t w[1] = {9};
    EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<3>::integral_type>(
        dpf::interleaved_bits_inner_product<3>(
            static_cast<const std::uint64_t *>(nullptr), 0, 3, w))),
        0u);
    std::uint64_t bits = 0;
    EXPECT_EQ(static_cast<unsigned>(static_cast<dpf::modint<3>::integral_type>(
        dpf::interleaved_bits_inner_product<3>(&bits, 5, 0, w))),
        0u);
}