#include <gtest/gtest.h>

#include <cstdint>
#include <cstring>

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

namespace
{

template <unsigned M, unsigned LaneBits = 8>
struct scalar_acc
{
    using mod = dpf::bitmore_mod<M, LaneBits>;
    unsigned acc = 0;
    unsigned bound = 0;

    void add(unsigned x, unsigned addend_max)
    {
        x &= mod::slot_max;
        if (addend_max > mod::slot_max)
            addend_max = mod::slot_max;
        for (;;)
        {
            if (bound + addend_max <= mod::slot_max)
                break;
            if (bound > mod::partial_bound)
            {
                acc = mod::partial_reduce_slot(acc);
                bound = mod::partial_bound;
                continue;
            }
            if (addend_max > mod::partial_bound)
            {
                x = mod::partial_reduce_slot(x);
                addend_max = mod::partial_bound;
                continue;
            }
            if (bound > mod::stable_bound)
            {
                acc = mod::partial_reduce_slot(acc);
                bound = mod::stable_bound;
                continue;
            }
            if (addend_max > mod::stable_bound)
            {
                x = mod::partial_reduce_slot(x);
                addend_max = mod::stable_bound;
                continue;
            }
            break;
        }
        acc += x;
        bound += addend_max;
    }

    void insert(unsigned bit)
    {
        while (bound > (mod::slot_max >> 1))
        {
            acc = mod::partial_reduce_slot(acc);
            bound = bound > mod::partial_bound ? mod::partial_bound : mod::stable_bound;
        }
        acc = ((acc << 1) | (bit & 1u)) & mod::slot_max;
        bound = bound * 2u + 1u;
    }
};

template <unsigned M, typename Reg>
void expect_bytes()
{
    using mod = dpf::bitmore_mod<M>;
    unsigned span = mod::resume_bound + 1u;
    unsigned k = 0;
    while ((span << 1) <= 256u)
    {
        span <<= 1;
        ++k;
    }
    EXPECT_EQ(mod::shift_budget, k);
    EXPECT_EQ(mod::add_budget, 255u - mod::resume_bound);
    EXPECT_LE((mod::resume_bound + 1u) * (1u << k) - 1u, 255u);
    EXPECT_GT(span << 1, 256u);

    constexpr int n = static_cast<int>(sizeof(Reg));
    for (unsigned base = 0; base < 256u; ++base)
    {
        alignas(32) unsigned char in[32]{};
        alignas(32) unsigned char out[32]{};
        for (int i = 0; i < n; ++i)
            in[i] = static_cast<unsigned char>((base + static_cast<unsigned>(i) * 13u) & 255u);
        Reg x;
        std::memcpy(&x, in, sizeof(Reg));

        auto partial = mod::partial_reduce(x);
        std::memcpy(out, &partial, sizeof(Reg));
        for (int i = 0; i < n; ++i)
        {
            const unsigned expect = mod::partial_reduce_byte(in[i]);
            if (out[i] != expect || expect > mod::partial_bound || expect % M != in[i] % M)
            {
                ADD_FAILURE() << "partial M=" << M << " in=" << static_cast<int>(in[i])
                              << " got=" << static_cast<int>(out[i]) << " expect=" << expect;
                return;
            }
        }

        auto full = mod::full_reduce(x);
        std::memcpy(out, &full, sizeof(Reg));
        for (int i = 0; i < n; ++i)
        {
            if (out[i] != in[i] % M)
            {
                ADD_FAILURE() << "full M=" << M << " in=" << static_cast<int>(in[i])
                              << " got=" << static_cast<int>(out[i]);
                return;
            }
        }
    }
}

template <unsigned M, typename Reg>
void expect_accumulator()
{
    constexpr int n = static_cast<int>(sizeof(Reg));
    dpf::bitmore_accumulator<M, Reg> acc;
    scalar_acc<M> slots[32]{};
    unsigned math[32]{};

    for (int step = 0; step < 96; ++step)
    {
        alignas(32) unsigned char bytes[32]{};
        Reg packed;
        if (step % 3 == 0)
        {
            for (int i = 0; i < n; ++i)
            {
                bytes[i] = static_cast<unsigned char>(((step * 17 + i * 3) & 1u) | (i & 0xf0));
                slots[i].insert(bytes[i]);
                math[i] = (math[i] * 2u + (bytes[i] & 1u)) % M;
            }
            std::memcpy(&packed, bytes, sizeof(Reg));
            acc.insert_bit(packed);
        }
        else if (step % 3 == 1)
        {
            for (int i = 0; i < n; ++i)
            {
                bytes[i] = static_cast<unsigned char>((step + i) % M);
                slots[i].add(bytes[i], M - 1u);
                math[i] = (math[i] + bytes[i]) % M;
            }
            std::memcpy(&packed, bytes, sizeof(Reg));
            acc.add(packed, M - 1u);
        }
        else
        {
            for (int i = 0; i < n; ++i)
            {
                bytes[i] = static_cast<unsigned char>(200u + (i & 15u));
                slots[i].add(bytes[i], 255u);
                math[i] = (math[i] + bytes[i]) % M;
            }
            std::memcpy(&packed, bytes, sizeof(Reg));
            acc.add(packed, 255u);
        }

        if (acc.bound() != slots[0].bound || acc.bound() > 255u)
        {
            ADD_FAILURE() << "bound M=" << M << " step=" << step
                          << " simd=" << acc.bound() << " scalar=" << slots[0].bound;
            return;
        }

        alignas(32) unsigned char raw[32]{};
        const auto value = acc.value();
        std::memcpy(raw, &value, sizeof(Reg));
        for (int i = 0; i < n; ++i)
        {
            if (raw[i] != slots[i].acc || raw[i] > acc.bound() || raw[i] % M != math[i])
            {
                ADD_FAILURE() << "lane M=" << M << " step=" << step << " i=" << i
                              << " got=" << static_cast<int>(raw[i])
                              << " expect=" << slots[i].acc
                              << " math=" << math[i] << " bound=" << acc.bound();
                return;
            }
        }
    }

    alignas(32) unsigned char red[32]{};
    const auto reduced = acc.reduced();
    std::memcpy(red, &reduced, sizeof(Reg));
    for (int i = 0; i < n; ++i)
    {
        if (red[i] != math[i])
        {
            ADD_FAILURE() << "reduced M=" << M << " i=" << i
                          << " got=" << static_cast<int>(red[i])
                          << " expect=" << math[i];
            return;
        }
    }
}

template <unsigned M>
void expect_all()
{
    expect_bytes<M, simde__m128i>();
    expect_bytes<M, simde__m256i>();
    if constexpr (M <= 128u)
    {
        expect_accumulator<M, simde__m128i>();
        expect_accumulator<M, simde__m256i>();
    }
    if constexpr (M < 255u)
        expect_all<M + 1u>();
}

}  // namespace

template <unsigned M, typename Reg>
void expect_epi16_slots()
{
    using mod = dpf::bitmore_mod<M, 16>;
    constexpr int lanes = static_cast<int>(sizeof(Reg) / 2u);
    unsigned span = mod::resume_bound + 1u;
    unsigned k = 0;
    const unsigned half = (mod::slot_max + 1u) >> 1;
    while (span <= half)
    {
        span *= 2u;
        ++k;
    }
    EXPECT_EQ(mod::shift_budget, k);
    EXPECT_EQ(mod::add_budget, mod::slot_max - mod::resume_bound);

    for (unsigned base = 0; base < 65536u; base += static_cast<unsigned>(lanes))
    {
        alignas(32) std::uint16_t in[16]{};
        alignas(32) std::uint16_t out[16]{};
        for (int i = 0; i < lanes; ++i)
            in[i] = static_cast<std::uint16_t>(base + static_cast<unsigned>(i));
        Reg x;
        std::memcpy(&x, in, sizeof(Reg));
        auto partial = mod::partial_reduce(x);
        std::memcpy(out, &partial, sizeof(Reg));
        for (int i = 0; i < lanes; ++i)
        {
            const unsigned expect = mod::partial_reduce_slot(in[i]);
            if (out[i] != expect || expect > mod::partial_bound || expect % M != in[i] % M)
            {
                ADD_FAILURE() << "partial16 M=" << M << " in=" << in[i]
                              << " got=" << out[i] << " expect=" << expect;
                return;
            }
        }
        auto full = mod::full_reduce(x);
        std::memcpy(out, &full, sizeof(Reg));
        for (int i = 0; i < lanes; ++i)
        {
            if (out[i] != in[i] % M)
            {
                ADD_FAILURE() << "full16 M=" << M << " in=" << in[i] << " got=" << out[i];
                return;
            }
        }
    }
}

template <unsigned M, typename Reg>
void expect_epi16_accumulator()
{
    constexpr int lanes = static_cast<int>(sizeof(Reg) / 2u);
    dpf::bitmore_accumulator<M, Reg, 16> acc;
    scalar_acc<M, 16> slots[16]{};
    unsigned math[16]{};

    for (int step = 0; step < 64; ++step)
    {
        alignas(32) std::uint16_t words[16]{};
        Reg packed;
        if (step % 3 == 0)
        {
            for (int i = 0; i < lanes; ++i)
            {
                words[i] = static_cast<std::uint16_t>(((step * 17 + i * 3) & 1) | ((i * 0x1111) & 0xfff0));
                slots[i].insert(words[i]);
                math[i] = (math[i] * 2u + (words[i] & 1u)) % M;
            }
            std::memcpy(&packed, words, sizeof(Reg));
            acc.insert_bit(packed);
        }
        else if (step % 3 == 1)
        {
            for (int i = 0; i < lanes; ++i)
            {
                words[i] = static_cast<std::uint16_t>((step * 100 + i * 17) % M);
                slots[i].add(words[i], M - 1u);
                math[i] = (math[i] + words[i]) % M;
            }
            std::memcpy(&packed, words, sizeof(Reg));
            acc.add(packed, M - 1u);
        }
        else
        {
            for (int i = 0; i < lanes; ++i)
            {
                words[i] = static_cast<std::uint16_t>(40000u + static_cast<unsigned>(i) * 97u);
                slots[i].add(words[i], 65535u);
                math[i] = (math[i] + words[i]) % M;
            }
            std::memcpy(&packed, words, sizeof(Reg));
            acc.add(packed, 65535u);
        }

        if (acc.bound() != slots[0].bound || acc.bound() > 65535u)
        {
            ADD_FAILURE() << "bound16 M=" << M << " step=" << step
                          << " simd=" << acc.bound() << " scalar=" << slots[0].bound;
            return;
        }

        alignas(32) std::uint16_t raw[16]{};
        const auto value = acc.value();
        std::memcpy(raw, &value, sizeof(Reg));
        for (int i = 0; i < lanes; ++i)
        {
            if (raw[i] != slots[i].acc || raw[i] > acc.bound() || raw[i] % M != math[i])
            {
                ADD_FAILURE() << "lane16 M=" << M << " step=" << step << " i=" << i
                              << " got=" << raw[i] << " expect=" << slots[i].acc
                              << " math=" << math[i] << " bound=" << acc.bound();
                return;
            }
        }
    }

    alignas(32) std::uint16_t red[16]{};
    const auto reduced = acc.reduced();
    std::memcpy(red, &reduced, sizeof(Reg));
    for (int i = 0; i < lanes; ++i)
    {
        if (red[i] != math[i])
        {
            ADD_FAILURE() << "reduced16 M=" << M << " i=" << i
                          << " got=" << red[i] << " expect=" << math[i];
            return;
        }
    }
}

template <unsigned M>
void expect_wide()
{
    expect_epi16_slots<M, simde__m128i>();
    expect_epi16_slots<M, simde__m256i>();
    if constexpr (M <= 32768u)
    {
        expect_epi16_accumulator<M, simde__m128i>();
        expect_epi16_accumulator<M, simde__m256i>();
    }
}

TEST(BitmoreMod, BytesAndBudget)
{
    expect_all<2>();
}

TEST(BitmoreMod, Epi16)
{
    expect_wide<2>();
    expect_wide<3>();
    expect_wide<15>();
    expect_wide<127>();
    expect_wide<128>();
    expect_wide<255>();
    expect_wide<256>();
    expect_wide<257>();
    expect_wide<4095>();
    expect_wide<4096>();
    expect_wide<4097>();
    expect_wide<16384>();
    expect_wide<16385>();
    expect_wide<32765>();
    expect_wide<32767>();
    expect_wide<32768>();
    expect_wide<32769>();
    expect_wide<40000>();
    expect_wide<61440>();
    expect_wide<65535>();
}