#include <gtest/gtest.h>

#include <cstdint>
#include <cstring>
#include <array>

#include "dpf.hpp"
#include "simde/simde/x86/avx2.h"

namespace
{

template <unsigned Bits>
unsigned lane_of(unsigned byte, unsigned shift)
{
    return (byte >> shift) & ((1u << Bits) - 1u);
}

template <typename Reg, unsigned Bits, typename Simd>
void expect_vv(Simd simd, unsigned (*scalar)(unsigned, unsigned))
{
    constexpr int nbytes = static_cast<int>(sizeof(Reg));
    constexpr unsigned per = 8u / Bits;
    for (int a0 = 0; a0 < 256; ++a0)
    {
        for (int b0 = 0; b0 < 256; ++b0)
        {
            alignas(32) unsigned char ab[32]{};
            alignas(32) unsigned char bb[32]{};
            alignas(32) unsigned char cb[32]{};
            for (int i = 0; i < nbytes; ++i)
            {
                ab[i] = static_cast<unsigned char>((a0 + i * 17) & 255);
                bb[i] = static_cast<unsigned char>((b0 + i * 3) & 255);
            }
            Reg a, b, c;
            std::memcpy(&a, ab, sizeof(Reg));
            std::memcpy(&b, bb, sizeof(Reg));
            c = simd(a, b);
            std::memcpy(cb, &c, sizeof(Reg));
            for (int i = 0; i < nbytes; ++i)
            {
                for (unsigned s = 0; s < per; ++s)
                {
                    const unsigned shift = s * Bits;
                    const unsigned expect = scalar(
                        lane_of<Bits>(ab[i], shift),
                        lane_of<Bits>(bb[i], shift)) & ((1u << Bits) - 1u);
                    const unsigned got = lane_of<Bits>(cb[i], shift);
                    if (got != expect)
                    {
                        ADD_FAILURE() << "byte " << i << " shift " << shift
                                      << " a0 " << a0 << " b0 " << b0
                                      << " got " << got << " expect " << expect;
                        return;
                    }
                }
            }
        }
    }
}

template <typename Reg, unsigned Bits, typename Simd>
void expect_sv(Simd simd)
{
    constexpr int nbytes = static_cast<int>(sizeof(Reg));
    constexpr unsigned mask = (1u << Bits) - 1u;
    constexpr unsigned per = 8u / Bits;
    for (unsigned k = 0; k <= mask; ++k)
    {
        alignas(32) unsigned char ab[32]{};
        alignas(32) unsigned char cb[32]{};
        for (int i = 0; i < nbytes; ++i)
            ab[i] = static_cast<unsigned char>(i * 13u + k);
        Reg a, c;
        std::memcpy(&a, ab, sizeof(Reg));
        if constexpr (Bits == 2)
            c = simd(a, static_cast<dpf::twobit>(k));
        else
            c = simd(a, static_cast<dpf::nyble>(k));
        std::memcpy(cb, &c, sizeof(Reg));
        for (int i = 0; i < nbytes; ++i)
        {
            for (unsigned s = 0; s < per; ++s)
            {
                const unsigned shift = s * Bits;
                const unsigned expect
                    = (lane_of<Bits>(ab[i], shift) * k) & mask;
                const unsigned got = lane_of<Bits>(cb[i], shift);
                if (got != expect)
                {
                    ADD_FAILURE() << "sv byte " << i << " k " << k
                                  << " got " << got << " expect " << expect;
                    return;
                }
            }
        }
    }
}

template <typename Lane>
void check_ring()
{
    constexpr unsigned n = 1u << dpf::utils::packed_lane_bits_v<Lane>;
    for (unsigned a = 0; a < n; ++a)
    {
        for (unsigned b = 0; b < n; ++b)
        {
            const auto x = static_cast<Lane>(a);
            const auto y = static_cast<Lane>(b);
            EXPECT_EQ(static_cast<unsigned>(x + y), (a + b) & (n - 1u));
            EXPECT_EQ(static_cast<unsigned>(x - y), (a - b) & (n - 1u));
            EXPECT_EQ(static_cast<unsigned>(x * y), (a * b) & (n - 1u));
            EXPECT_EQ(static_cast<unsigned>(x + (x - x)), a);
            EXPECT_EQ(static_cast<unsigned>(-x), (0u - a) & (n - 1u));
        }
    }
}

template <typename NodeT, typename Lane>
void check_deposit(std::size_t lane)
{
    const auto y = static_cast<Lane>(lane & ((1u << dpf::utils::packed_lane_bits_v<Lane>) - 1u));
    auto node = dpf::packed::make_lane_node<NodeT>(lane, y);
    EXPECT_EQ(dpf::packed::extract_lane<Lane>(node, lane), y);
    if (lane > 0)
    {
        EXPECT_EQ(dpf::packed::extract_lane<Lane>(node, lane - 1), Lane{});
    }
}

template <typename Lane>
void check_dpf_output(Lane y)
{
    using input_type = std::uint32_t;
    constexpr input_type x = 40;
    constexpr input_type n = 96;
    auto [k0, k1] = dpf::make_dpf(x, y);
    for (input_type i = 0; i < n; ++i)
    {
        const auto got = dpf::reconstruct(
            *dpf::eval_point(k0, i), *dpf::eval_point(k1, i));
        EXPECT_EQ(got, i == x ? y : Lane{}) << "point " << i;
    }

    auto [buf0, it0] = dpf::eval_interval(k0, input_type{0}, input_type{n - 1});
    auto [buf1, it1] = dpf::eval_interval(k1, input_type{0}, input_type{n - 1});
    auto p0 = it0.begin();
    auto p1 = it1.begin();
    for (input_type i = 0; i < n; ++i, ++p0, ++p1)
    {
        EXPECT_EQ(dpf::reconstruct(*p0, *p1), i == x ? y : Lane{})
            << "interval " << i;
    }
    EXPECT_EQ(p0, it0.end());
    EXPECT_EQ(p1, it1.end());
}

}  // namespace

TEST(PackedLane, ScalarRing)
{
    using namespace dpf::literals::twobit;
    using namespace dpf::literals::nyble;

    check_ring<dpf::twobit>();
    check_ring<dpf::nyble>();
    EXPECT_EQ(dpf::to_string(dpf::twobit::three), "3");
    EXPECT_EQ(dpf::to_string(dpf::nyble{0x0a}), "a");
    EXPECT_EQ(dpf::to_nyble('F'), dpf::nyble{15});
    EXPECT_EQ(2_twobit, dpf::twobit::two);
    EXPECT_EQ(10_nyble, dpf::nyble{10});
    EXPECT_EQ(dpf::utils::bitlength_of_v<dpf::twobit>, 2u);
    EXPECT_EQ(dpf::utils::bitlength_of_v<dpf::nyble>, 4u);
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    EXPECT_EQ((dpf::utils::bitlength_of_output_v<dpf::twobit, simde__m256i>), 2u);
    EXPECT_EQ((dpf::utils::bitlength_of_output_v<dpf::nyble, simde__m128i>), 4u);
    EXPECT_EQ((dpf::outputs_per_leaf_v<dpf::twobit, simde__m128i>), 64u);
    EXPECT_EQ((dpf::outputs_per_leaf_v<dpf::twobit, simde__m256i>), 128u);
    EXPECT_EQ((dpf::outputs_per_leaf_v<dpf::nyble, simde__m128i>), 32u);
    EXPECT_EQ((dpf::outputs_per_leaf_v<dpf::nyble, simde__m256i>), 64u);
    EXPECT_EQ((dpf::lg_outputs_per_leaf_v<dpf::twobit, simde__m256i>), 7u);
    HEDLEY_PRAGMA(GCC diagnostic pop)
}

TEST(PackedLane, SimdMatchesScalar)
{
    auto add2 = [](unsigned a, unsigned b) { return (a + b) & 3u; };
    auto sub2 = [](unsigned a, unsigned b) { return (a - b) & 3u; };
    auto mul2 = [](unsigned a, unsigned b) { return (a * b) & 3u; };
    auto add4 = [](unsigned a, unsigned b) { return (a + b) & 15u; };
    auto sub4 = [](unsigned a, unsigned b) { return (a - b) & 15u; };
    auto mul4 = [](unsigned a, unsigned b) { return (a * b) & 15u; };

    expect_vv<simde__m128i, 2>([](auto a, auto b) { return dpf::lane_arith::add_epi2(a, b); }, add2);
    expect_vv<simde__m256i, 2>([](auto a, auto b) { return dpf::lane_arith::add_epi2(a, b); }, add2);
    expect_vv<simde__m128i, 2>([](auto a, auto b) { return dpf::lane_arith::sub_epi2(a, b); }, sub2);
    expect_vv<simde__m256i, 2>([](auto a, auto b) { return dpf::lane_arith::sub_epi2(a, b); }, sub2);
    expect_vv<simde__m128i, 2>([](auto a, auto b) { return dpf::lane_arith::mullo_epi2(a, b); }, mul2);
    expect_vv<simde__m256i, 2>([](auto a, auto b) { return dpf::lane_arith::mullo_epi2(a, b); }, mul2);
    expect_sv<simde__m128i, 2>([](auto a, auto k) { return dpf::lane_arith::mul_epi2(a, k); });
    expect_sv<simde__m256i, 2>([](auto a, auto k) { return dpf::lane_arith::mul_epi2(a, k); });

    expect_vv<simde__m128i, 4>([](auto a, auto b) { return dpf::lane_arith::add_epi4(a, b); }, add4);
    expect_vv<simde__m256i, 4>([](auto a, auto b) { return dpf::lane_arith::add_epi4(a, b); }, add4);
    expect_vv<simde__m128i, 4>([](auto a, auto b) { return dpf::lane_arith::sub_epi4(a, b); }, sub4);
    expect_vv<simde__m256i, 4>([](auto a, auto b) { return dpf::lane_arith::sub_epi4(a, b); }, sub4);
    expect_vv<simde__m128i, 4>([](auto a, auto b) { return dpf::lane_arith::mullo_epi4(a, b); }, mul4);
    expect_vv<simde__m256i, 4>([](auto a, auto b) { return dpf::lane_arith::mullo_epi4(a, b); }, mul4);
    expect_sv<simde__m128i, 4>([](auto a, auto k) { return dpf::lane_arith::mul_epi4(a, k); });
    expect_sv<simde__m256i, 4>([](auto a, auto k) { return dpf::lane_arith::mul_epi4(a, k); });
}

TEST(PackedLane, LeafFunctorsMatchSimd)
{
    alignas(32) unsigned char bytes[32];
    for (int i = 0; i < 32; ++i)
        bytes[i] = static_cast<unsigned char>(0x5a ^ i);
    simde__m128i a128, b128;
    simde__m256i a256, b256;
    std::memcpy(&a128, bytes, 16);
    std::memcpy(&b128, bytes + 8, 16);
    std::memcpy(&a256, bytes, 32);
    std::memcpy(&b256, bytes, 32);

    auto sum128 = dpf::add_leaf<dpf::twobit>(a128, b128);
    auto via = dpf::lane_arith::add_epi2(a128, b128);
    EXPECT_EQ(std::memcmp(&sum128, &via, sizeof(via)), 0);

    auto scaled = dpf::multiply_leaf(a256, dpf::nyble{7});
    auto via4 = dpf::lane_arith::mul_epi4(a256, dpf::nyble{7});
    EXPECT_EQ(std::memcmp(&scaled, &via4, sizeof(via4)), 0);

    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    std::array<simde__m128i, 2> aa{a128, b128};
    std::array<simde__m128i, 2> bb{b128, a128};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto both = dpf::add_leaf<dpf::nyble>(aa, bb);
    auto e0 = dpf::lane_arith::add_epi4(a128, b128);
    auto e1 = dpf::lane_arith::add_epi4(b128, a128);
    EXPECT_EQ(std::memcmp(&both[0], &e0, sizeof(e0)), 0);
    EXPECT_EQ(std::memcmp(&both[1], &e1, sizeof(e1)), 0);
}

TEST(PackedLane, DepositExtractAndBuffer)
{
    for (std::size_t lane : {0u, 1u, 63u, 64u, 127u})
    {
        check_deposit<simde__m128i, dpf::bit>(lane);
        check_deposit<simde__m128i, dpf::twobit>(lane < 64 ? lane : lane / 2);
        check_deposit<simde__m128i, dpf::nyble>(lane % 32);
    }
    for (std::size_t lane : {0u, 127u, 128u, 200u, 255u})
    {
        check_deposit<simde__m256i, dpf::bit>(lane);
        check_deposit<simde__m256i, dpf::twobit>(lane < 128 ? lane : 127);
        check_deposit<simde__m256i, dpf::nyble>(lane % 64);
    }

    simde__m128i leaf{};
    for (std::size_t i = 0; i < 64; ++i)
        dpf::packed::deposit_lane(leaf, i, dpf::to_twobit(static_cast<unsigned>(i)));
    dpf::output_buffer<dpf::twobit> buf(128);
    dpf::store_leaf_bytes(buf, 0, leaf);
    simde__m128i other{};
    dpf::packed::deposit_lane(other, 3, dpf::twobit::two);
    dpf::store_leaf_bytes(buf, 1, other);
    for (std::size_t i = 0; i < 64; ++i)
    {
        EXPECT_EQ(static_cast<dpf::twobit>(buf[i]),
            dpf::to_twobit(static_cast<unsigned>(i)));
    }
    EXPECT_EQ(static_cast<dpf::twobit>(buf[64 + 3]), dpf::twobit::two);
    EXPECT_EQ(static_cast<dpf::twobit>(buf[64]), dpf::twobit::zero);

    buf[5] = dpf::twobit::one;
    EXPECT_EQ(static_cast<dpf::twobit>(buf[5]), dpf::twobit::one);
}

TEST(PackedLane, DpfPointAndInterval)
{
    check_dpf_output(dpf::twobit::three);
    check_dpf_output(dpf::nyble{0x0d});
    check_dpf_output(dpf::twobit::zero);
    check_dpf_output(dpf::nyble{1});
}