#include <gtest/gtest.h>

#include <cstdint>
#include <vector>

#include "dpf/flute.hpp"

TEST(Flute, EveryTwoBitFunction)
{
    const unsigned delta = 2;
    std::uint8_t rows[4];
    for (unsigned fn = 0; fn < 16; ++fn)
    {
        for (unsigned j = 0; j < 4; ++j)
            rows[j] = static_cast<std::uint8_t>((fn >> j) & 1u);
        for (unsigned x = 0; x < 4; ++x)
        {
            std::uint8_t bits[2] = {
                static_cast<std::uint8_t>(x & 1u),
                static_cast<std::uint8_t>((x >> 1) & 1u),
            };
            auto want = dpf::flute::eval_plain(delta, 1, rows, bits);
            auto pair = dpf::flute::eval_pair(delta, 1, rows, bits);
            auto trio = dpf::flute::eval_trio(delta, 1, rows, bits);
            EXPECT_EQ(pair.opened, want);
            EXPECT_EQ(trio.opened, want);
            EXPECT_EQ(pair.online_bits, 2u);
            EXPECT_EQ(trio.online_bits, 3u);
            EXPECT_EQ(pair.mask.size(), 2u);
            EXPECT_EQ(trio.mask.size(), 3u);
        }
    }
}

TEST(Flute, TwoOutputBits)
{
    // LSB of the 3-bit index, and its majority with the other two bits.
    const unsigned delta = 3;
    const unsigned rows = 8;
    std::vector<std::uint8_t> columns(2 * rows);
    for (unsigned j = 0; j < rows; ++j)
    {
        columns[j] = static_cast<std::uint8_t>(j & 1u);
        const unsigned ones = (j & 1u) + ((j >> 1) & 1u) + ((j >> 2) & 1u);
        columns[rows + j] = static_cast<std::uint8_t>(ones >= 2);
    }
    for (unsigned x = 0; x < rows; ++x)
    {
        std::uint8_t bits[3] = {
            static_cast<std::uint8_t>(x & 1u),
            static_cast<std::uint8_t>((x >> 1) & 1u),
            static_cast<std::uint8_t>((x >> 2) & 1u),
        };
        auto want = dpf::flute::eval_plain(delta, 2, columns.data(), bits);
        auto got = dpf::flute::eval_pair(delta, 2, columns.data(), bits);
        EXPECT_EQ(got.opened, want);
        EXPECT_EQ(got.online_bits, 4u);
        std::uint8_t lam0 = 0;
        std::uint8_t lam1 = 0;
        lam0 = static_cast<std::uint8_t>(got.mask[0][0] ^ got.mask[0][1]);
        lam1 = static_cast<std::uint8_t>(got.mask[1][0] ^ got.mask[1][1]);
        (void)lam0;
        (void)lam1;
        for (unsigned w = 0; w < 2; ++w)
        {
            const std::uint8_t lam = static_cast<std::uint8_t>(
                got.mask[0][w] ^ got.mask[1][w]);
            EXPECT_EQ(static_cast<std::uint8_t>(got.masked[w] ^ lam), want[w]);
        }
    }
}

namespace
{

void expect_lut(unsigned delta, unsigned n_out, const std::vector<std::uint8_t> & columns,
    const std::uint8_t * bits)
{
    auto want = dpf::flute::eval_plain(delta, n_out, columns.data(), bits);
    auto pair = dpf::flute::eval_pair(delta, n_out, columns.data(), bits);
    auto trio = dpf::flute::eval_trio(delta, n_out, columns.data(), bits);
    EXPECT_EQ(pair.opened, want);
    EXPECT_EQ(trio.opened, want);
    EXPECT_EQ(pair.online_bits, static_cast<std::size_t>(2 * n_out));
    EXPECT_EQ(trio.online_bits, static_cast<std::size_t>(3 * n_out));
    for (unsigned w = 0; w < n_out; ++w)
    {
        const std::uint8_t two = static_cast<std::uint8_t>(
            pair.masked[w] ^ pair.mask[0][w] ^ pair.mask[1][w]);
        const std::uint8_t three = static_cast<std::uint8_t>(
            trio.masked[w] ^ trio.mask[0][w] ^ trio.mask[1][w] ^ trio.mask[2][w]);
        EXPECT_EQ(two, want[w]);
        EXPECT_EQ(three, want[w]);
    }
}

std::uint32_t xorshift(std::uint32_t & s)
{
    s ^= s << 13;
    s ^= s >> 17;
    s ^= s << 5;
    return s;
}

}  // namespace

TEST(Flute, EveryFunctionThroughThreeBits)
{
    for (unsigned delta = 1; delta <= 3; ++delta)
    {
        const unsigned rows = 1u << delta;
        const unsigned nfn = 1u << rows;
        std::vector<std::uint8_t> column(rows);
        std::vector<std::uint8_t> bits(delta);
        for (unsigned fn = 0; fn < nfn; ++fn)
        {
            for (unsigned j = 0; j < rows; ++j)
                column[j] = static_cast<std::uint8_t>((fn >> j) & 1u);
            for (unsigned x = 0; x < rows; ++x)
            {
                for (unsigned i = 0; i < delta; ++i)
                    bits[i] = static_cast<std::uint8_t>((x >> i) & 1u);
                expect_lut(delta, 1, column, bits.data());
            }
        }
    }
}

TEST(Flute, FourBitSampleAndEightBitSpecials)
{
    const unsigned delta = 4;
    const unsigned rows = 16;
    std::uint32_t rng = 0xC0FFEEu;
    std::vector<std::uint8_t> columns(2 * rows);
    std::vector<std::uint8_t> bits(delta);
    for (unsigned sample = 0; sample < 24; ++sample)
    {
        for (auto & cell : columns)
            cell = static_cast<std::uint8_t>(xorshift(rng) & 1u);
        for (unsigned x = 0; x < rows; ++x)
        {
            for (unsigned i = 0; i < delta; ++i)
                bits[i] = static_cast<std::uint8_t>((x >> i) & 1u);
            expect_lut(delta, 2, columns, bits.data());
        }
    }

    const unsigned wide = 8;
    const unsigned wrows = 256;
    std::vector<std::uint8_t> table(4 * wrows, 0);
    for (unsigned j = 0; j < wrows; ++j)
    {
        unsigned ones = 0;
        for (unsigned i = 0; i < wide; ++i)
            ones += (j >> i) & 1u;
        table[j] = static_cast<std::uint8_t>(ones & 1u);
        table[wrows + j] = static_cast<std::uint8_t>(ones == wide);
        table[2 * wrows + j] = static_cast<std::uint8_t>((j >> 7) & 1u);
        table[3 * wrows + j] = static_cast<std::uint8_t>(j == 0 || j == 255);
    }
    const unsigned points[] = {0u, 1u, 2u, 17u, 128u, 170u, 254u, 255u};
    std::vector<std::uint8_t> wbits(wide);
    for (unsigned x : points)
    {
        for (unsigned i = 0; i < wide; ++i)
            wbits[i] = static_cast<std::uint8_t>((x >> i) & 1u);
        expect_lut(wide, 4, table, wbits.data());
    }
}

TEST(Flute, RejectsABadIndex)
{
    std::uint8_t column[2] = {0, 1};
    std::uint8_t bit = 1;
    EXPECT_THROW(dpf::flute::eval_plain(0, 1, column, &bit), std::invalid_argument);
    EXPECT_THROW(dpf::flute::eval_pair(9, 1, column, &bit), std::invalid_argument);
    bit = 2;
    EXPECT_THROW(dpf::flute::eval_trio(1, 1, column, &bit), std::invalid_argument);
    EXPECT_THROW(dpf::flute::eval_pair(1, 0, column, &bit), std::invalid_argument);
}