#include <gtest/gtest.h>

#include "dpf.hpp"
#include "grotto/constant_lut.hpp"
#include "grotto/prefix_parity.hpp"

#include <array>
#include <cstdint>
#include <limits>
#include <vector>

namespace
{

uint64_t opened(uint64_t a, uint64_t b, uint64_t mask)
{
    return (a + b) & mask;
}

template <typename T, std::size_t N>
int piece_containing(T alpha, const std::array<T, N> & ends)
{
    for (std::size_t i = 0; i + 1 < N; ++i)
    {
        if (alpha >= ends[i] && alpha < ends[i + 1])
            return static_cast<int>(i);
    }
    return static_cast<int>(N - 1);
}

}  // namespace

TEST(SignedPrefix, MatchesComparisonEvalOnEveryUint8Point)
{
    const uint8_t alpha = 0x3c;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(uint64_t{1}));
    const uint64_t mask = k0.cmp().mask;
    std::array<uint8_t, 8> ends{0, 1, 10, alpha, 100, 200, 254, 255};
    const auto p0 = grotto::signed_prefix_parities(k0, ends);
    const auto p1 = grotto::signed_prefix_parities(k1, ends);
    for (std::size_t i = 0; i < ends.size(); ++i)
    {
        const auto e0 = dpf::eval_point(dpf::cmp, k0, ends[i]);
        const auto e1 = dpf::eval_point(dpf::cmp, k1, ends[i]);
        EXPECT_EQ(p0[i], e0.raw()) << int(ends[i]);
        EXPECT_EQ(p1[i], e1.raw()) << int(ends[i]);
        EXPECT_EQ(opened(p0[i], p1[i], mask), ends[i] > alpha ? 1u : 0u);
    }
}

TEST(SignedPrefix, SegmentsAreOneHot)
{
    const uint8_t alpha = 40;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(uint64_t{1}));
    const uint64_t mask = k0.cmp().mask;
    const std::array<uint8_t, 5> ends{0, 1, 10, 40, 200};
    const auto s0 = grotto::signed_segment_parities(k0, ends);
    const auto s1 = grotto::signed_segment_parities(k1, ends);
    const int hot = piece_containing(alpha, ends);
    uint64_t sum = 0;
    for (std::size_t i = 0; i < ends.size(); ++i)
    {
        const uint64_t bit = opened(s0[i], s1[i], mask);
        sum = (sum + bit) & mask;
        EXPECT_EQ(bit, i == static_cast<std::size_t>(hot) ? 1u : 0u) << i;
    }
    EXPECT_EQ(sum, 1u);

    for (uint8_t a : {uint8_t{0}, uint8_t{1}, uint8_t{9}, uint8_t{200}, uint8_t{255}})
    {
        auto [a0, a1] = dpf::make_dpf(a, dpf::gt(uint64_t{1}));
        const auto t0 = grotto::signed_segment_parities(a0, ends);
        const auto t1 = grotto::signed_segment_parities(a1, ends);
        const int where = piece_containing(a, ends);
        for (std::size_t i = 0; i < ends.size(); ++i)
            EXPECT_EQ(opened(t0[i], t1[i], mask), i == static_cast<std::size_t>(where) ? 1u : 0u)
                << "alpha=" << int(a) << " piece=" << i;
    }
}

TEST(SignedPrefix, WholeDomainIsThePublicOne)
{
    auto [k0, k1] = dpf::make_dpf(uint8_t{7}, dpf::gt(uint64_t{1}));
    const std::array<uint8_t, 1> ends{0};
    const auto s0 = grotto::signed_segment_parities(k0, ends);
    const auto s1 = grotto::signed_segment_parities(k1, ends);
    EXPECT_EQ(opened(s0[0], s1[0], k0.cmp().mask), 1u);
}

TEST(SignedPrefix, SignumLutIsSignCorrect)
{
    for (int alpha = -128; alpha <= 127; ++alpha)
    {
        const auto a = static_cast<int8_t>(alpha);
        auto [k0, k1] = dpf::make_dpf(a, dpf::gt(uint64_t{1}));
        const auto lut = grotto::make_exact_constant_lut<int8_t>(
            grotto::exact_constant::signum, 0);
        ASSERT_EQ(lut.bounds.size(), 3u);
        std::array<int8_t, 3> ends{};
        for (std::size_t i = 0; i < 3; ++i)
            ends[i] = lut.bounds[i];
        const auto s0 = grotto::signed_segment_parities(k0, ends);
        const auto s1 = grotto::signed_segment_parities(k1, ends);
        uint64_t acc0 = 0;
        uint64_t acc1 = 0;
        for (std::size_t i = 0; i < 3; ++i)
        {
            const auto coef = static_cast<uint64_t>(lut.values[i]);
            acc0 += s0[i] * coef;
            acc1 += s1[i] * coef;
        }
        const auto got = static_cast<int64_t>(acc0 + acc1);
        const int64_t want = alpha < 0 ? -1 : (alpha > 0 ? 1 : 0);
        EXPECT_EQ(got, want) << alpha;
        // The advice-bit lift is ±want. The DCF shares must not come out negated.
        if (want != 0)
            EXPECT_NE(got, -want) << alpha;
    }
}

TEST(SignedPrefix, SharedPrefixAgreesWithAFreshWalk)
{
    const uint16_t alpha = 0x0100;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(uint64_t{1}));
    // Consecutive knots share a long prefix; the cached sum must match a
    // one-knot walk.
    const std::array<uint16_t, 4> ends{0, 1, 0x0100, 0x0101};
    const auto many = grotto::signed_prefix_parities(k0, ends);
    for (std::size_t i = 0; i < ends.size(); ++i)
    {
        const std::array<uint16_t, 1> one{ends[i]};
        const auto alone = grotto::signed_prefix_parities(k0, one);
        EXPECT_EQ(many[i], alone[0]) << i;
    }
    (void)k1;
}

TEST(SignedPrefix, RejectsAKeyWithoutAComparison)
{
    auto [k0, k1] = dpf::make_dpf(uint8_t{1}, dpf::bit::one);
    const std::array<uint8_t, 2> ends{0, 1};
    EXPECT_THROW(grotto::signed_prefix_parities(k0, ends), std::invalid_argument);
    EXPECT_THROW(grotto::signed_segment_parities(k1, ends), std::invalid_argument);
}

template <std::size_t N>
void expect_ilogb_segments(const grotto::constant_lut<int8_t> & lut, int8_t alpha)
{
    if (lut.bounds.size() != N)
        return;
    std::array<int8_t, N> ends{};
    for (std::size_t i = 0; i < N; ++i)
        ends[i] = lut.bounds[i];
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(uint64_t{1}));
    const uint64_t mask = k0.cmp().mask;
    const auto s0 = grotto::signed_segment_parities(k0, ends);
    const auto s1 = grotto::signed_segment_parities(k1, ends);
    uint64_t acc = 0;
    for (std::size_t i = 0; i < N; ++i)
    {
        const uint64_t bit = opened(s0[i], s1[i], mask);
        acc += bit * static_cast<uint64_t>(lut.values[i]);
    }
    EXPECT_EQ(acc, static_cast<uint64_t>(lut(alpha))) << int(alpha);
}

template <std::size_t N>
void dispatch_ilogb(const grotto::constant_lut<int8_t> & lut, int8_t alpha, bool & matched)
{
    if (lut.bounds.size() == N)
    {
        expect_ilogb_segments<N>(lut, alpha);
        matched = true;
        return;
    }
    if constexpr (N > 1)
        dispatch_ilogb<N - 1>(lut, alpha, matched);
}

TEST(SignedPrefix, SegmentsRecoverIlogbInt8)
{
    for (unsigned frac : {0u, 4u})
    {
        const auto lut = grotto::make_exact_constant_lut<int8_t>(
            grotto::exact_constant::ilogb, frac);
        ASSERT_GE(lut.bounds.size(), 3u);
        ASSERT_LE(lut.bounds.size(), 40u);
        for (int v = -128; v <= 127; ++v)
        {
            bool matched = false;
            dispatch_ilogb<40>(lut, static_cast<int8_t>(v), matched);
            ASSERT_TRUE(matched) << lut.bounds.size();
        }
    }
}