#include <gtest/gtest.h>

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

#include <cstdint>
#include <cstring>
#include <string>
#include <vector>

namespace
{

template <typename T>
T bare(const T & v) { return v; }

template <typename T, std::size_t Party, dpf::sharing Scheme>
T bare(const dpf::secret_share<T, Party, Scheme> & s) { return s.raw(); }

template <typename A, typename B>
auto opened(const A & a, const B & b)
{
    if constexpr (dpf::is_secret_share_v<std::decay_t<A>>
        && dpf::is_secret_share_v<std::decay_t<B>>)
        return dpf::reconstruct(a, b);
    else
        return a - b;
}

template <typename Key, typename In>
auto ev(const Key & key, In x)
{
    return *dpf::eval_point(key, x);
}

template <typename In, typename Out>
void every_point(In alpha, Out y, uint64_t n)
{
    auto [k0, k1] = dpf::make_dpf(alpha, y);
    for (uint64_t i = 0; i < n; ++i)
    {
        In q{static_cast<typename In::integral_type>(i)};
        const Out got = opened(ev(k0, q), ev(k1, q));
        EXPECT_EQ(got, q == alpha ? y : Out{}) << i;
    }
}

simde__m128i g_roots[4];
int g_ri = 0;
simde__m128i take_root() { return g_roots[g_ri++]; }
struct Pad
{
    uint64_t n = 1;
    simde__m128i block()
    {
        auto v = simde_mm_set_epi64x(static_cast<long long>(n * 3),
            static_cast<long long>(n));
        n += 2;
        return v;
    }
    uint8_t bit() { return static_cast<uint8_t>((n++ >> 1) & 1u); }
};
void reset_roots()
{
    g_ri = 0;
    for (int i = 0; i < 4; ++i)
        g_roots[i] = simde_mm_set_epi64x(0x42 * (i + 3), 0x1000 + i * 9);
}

}  // namespace

TEST(ContextBlast, ModintEdgesAndFullDomain)
{
    using m4 = dpf::modint<4>;
    using m7 = dpf::modint<7>;
    using m9 = dpf::modint<9>;
    using m10 = dpf::modint<10>;
    using m15 = dpf::modint<15>;

    EXPECT_EQ(m4{15} + m4{1}, m4{0});
    EXPECT_EQ(m4{0} - m4{1}, m4{15});
    EXPECT_EQ(m7{127} + m7{1}, m7{0});
    EXPECT_EQ(m7{0} - m7{1}, m7{127});
    EXPECT_EQ(m9{511} + m9{1}, m9{0});
    EXPECT_EQ(m10{1023} + m10{1}, m10{0});
    EXPECT_EQ(m15{32767} + m15{1}, m15{0});

    // A uint8 leaf needs 4 index bits, so the domain must be wider than that.
    every_point(m4{0}, uint32_t{1}, 16);
    every_point(m4{15}, uint32_t{9}, 16);
    every_point(m7{0}, uint32_t{3}, 128);
    every_point(m7{1}, uint32_t{3}, 128);
    every_point(m7{127}, uint32_t{8}, 128);
    every_point(m7{64}, dpf::xor_wrapper<uint32_t>{0x00ff}, 128);

    for (unsigned a : {0u, 1u, 511u, 512u, 1023u})
        every_point(m10{static_cast<m10::integral_type>(a)}, uint32_t{0xabcdu}, 1024);

    every_point(m9{0}, uint16_t{2}, 512);
    every_point(m9{511}, uint16_t{2}, 512);
    every_point(m15{0}, uint32_t{1}, 32768);
    every_point(m15{32767}, uint32_t{7}, 32768);
}

TEST(ContextBlast, KeywordBinaryAndOctal)
{
    using bin = dpf::keyword<4, dpf::alphabets::binary>;
    using oct = dpf::keyword<2, dpf::alphabets::octal>;
    const uint16_t y = 0x1234;
    const char *bins[] = {"0000", "0001", "0010", "0111", "1000", "1111"};
    for (const char *s : bins)
    {
        bin alpha{s};
        auto [k0, k1] = dpf::make_dpf(alpha, y);
        for (unsigned i = 0; i < 16; ++i)
        {
            char buf[5] = {'0', '0', '0', '0', 0};
            for (int b = 0; b < 4; ++b)
                buf[3 - b] = ((i >> b) & 1u) ? '1' : '0';
            bin q{buf};
            EXPECT_EQ(opened(ev(k0, q), ev(k1, q)), q == alpha ? y : uint16_t{0})
                << s << " vs " << buf;
        }
    }
    oct alpha{"07"};
    auto [k0, k1] = dpf::make_dpf(alpha, uint8_t{5});
    oct hit{"07"};
    oct miss{"10"};
    oct zero{"00"};
    EXPECT_EQ(opened(ev(k0, hit), ev(k1, hit)), uint8_t{5});
    EXPECT_EQ(opened(ev(k0, miss), ev(k1, miss)), uint8_t{0});
    EXPECT_EQ(opened(ev(k0, zero), ev(k1, zero)), uint8_t{0});
    EXPECT_THROW(oct{"89"}, std::domain_error);
    EXPECT_THROW(bin{"2"}, std::domain_error);
}

TEST(ContextBlast, BitstringPointIntervalSequence)
{
    using bs = dpf::bitstring<6>;
    const uint32_t y = 0x11111111u;
    for (unsigned a : {0u, 1u, 31u, 32u, 63u})
    {
        bs alpha{static_cast<bs::word_type>(a)};
        auto [k0, k1] = dpf::make_dpf(alpha, y);
        for (unsigned i = 0; i < 64; ++i)
        {
            bs q{static_cast<bs::word_type>(i)};
            EXPECT_EQ(opened(ev(k0, q), ev(k1, q)), q == alpha ? y : 0u) << i;
        }
        bs from{0};
        bs to{63};
        auto iv0 = dpf::eval_interval(k0, from, to);
        auto iv1 = dpf::eval_interval(k1, from, to);
        auto p0 = std::begin(iv0.second);
        auto p1 = std::begin(iv1.second);
        for (unsigned i = 0; i < 64; ++i, ++p0, ++p1)
            EXPECT_EQ(opened(*p0, *p1), i == a ? y : 0u);
        EXPECT_EQ(p0, std::end(iv0.second));

        std::vector<bs> seq{bs{63}, bs{0}, alpha, bs{1}, alpha};
        std::vector<bs> sorted = seq;
        std::sort(sorted.begin(), sorted.end());
        auto recipe = dpf::make_sequence_recipe(k0, sorted.begin(), sorted.end());
        auto s0 = dpf::eval_sequence(k0, recipe);
        auto s1 = dpf::eval_sequence(k1, recipe);
        auto q0 = std::begin(s0.second);
        auto q1 = std::begin(s1.second);
        for (bs x : sorted)
        {
            EXPECT_EQ(opened(*q0, *q1), x == alpha ? y : 0u);
            ++q0;
            ++q1;
        }
    }
}

TEST(ContextBlast, ComparisonKindsSignedAndUnsigned)
{
    auto check = [](auto alpha, auto kind_keys, auto pred) {
        auto [k0, k1] = kind_keys;
        const uint64_t mask = k0.cmp().mask;
        using In = decltype(alpha);
        constexpr auto bits = dpf::utils::bitlength_of_v<In>;
        const uint64_t n = uint64_t{1} << (bits > 8 ? 8 : bits);
        for (uint64_t i = 0; i < n; ++i)
        {
            In q;
            if constexpr (std::is_integral_v<In>)
                q = static_cast<In>(i);
            else
                q = In{static_cast<typename In::integral_type>(i)};
            const auto got = opened(dpf::eval_point(dpf::cmp, k0, q),
                dpf::eval_point(dpf::cmp, k1, q)) & mask;
            const uint64_t want = pred(q, alpha) ? 4u : 1u;
            EXPECT_EQ(got, want) << i;
        }
    };
    check(uint8_t{10}, dpf::make_dpf(uint8_t{10}, dpf::lt(uint64_t{4}, uint64_t{1})),
        [](uint8_t q, uint8_t a) { return q < a; });
    check(uint8_t{10}, dpf::make_dpf(uint8_t{10}, dpf::leq(uint64_t{4}, uint64_t{1})),
        [](uint8_t q, uint8_t a) { return q <= a; });
    check(uint8_t{10}, dpf::make_dpf(uint8_t{10}, dpf::gt(uint64_t{4}, uint64_t{1})),
        [](uint8_t q, uint8_t a) { return q > a; });
    check(uint8_t{10}, dpf::make_dpf(uint8_t{10}, dpf::geq(uint64_t{4}, uint64_t{1})),
        [](uint8_t q, uint8_t a) { return q >= a; });
    check(uint8_t{0}, dpf::make_dpf(uint8_t{0}, dpf::lt(uint64_t{4}, uint64_t{1})),
        [](uint8_t q, uint8_t a) { return q < a; });
    check(uint8_t{255}, dpf::make_dpf(uint8_t{255}, dpf::leq(uint64_t{4}, uint64_t{1})),
        [](uint8_t q, uint8_t a) { return q <= a; });
    check(int8_t{-1}, dpf::make_dpf(int8_t{-1}, dpf::lt(uint64_t{4}, uint64_t{1})),
        [](int8_t q, int8_t a) { return q < a; });
    check(int8_t{-128}, dpf::make_dpf(int8_t{-128}, dpf::gt(uint64_t{4}, uint64_t{1})),
        [](int8_t q, int8_t a) { return q > a; });
    using m = dpf::modint<8>;
    check(m{200}, dpf::make_dpf(m{200}, dpf::lt(uint64_t{4}, uint64_t{1})),
        [](m q, m a) { return q < a; });
}

TEST(ContextBlast, IncrementalSlotsAndInnerProduct)
{
    const uint16_t x = 0x81fe;
    auto [k0, k1] = dpf::make_dpf(x, dpf::at<8>(uint8_t{7}), dpf::at<12>(uint16_t{99}),
        uint32_t{5});
    for (uint16_t q : {uint16_t{0}, uint16_t{1}, x, uint16_t{0x8000}, uint16_t{0xffff},
             static_cast<uint16_t>(x ^ 1u), static_cast<uint16_t>(x ^ 0x0100u)})
    {
        const auto lane8 = static_cast<uint16_t>(q >> 8);
        const auto alane = static_cast<uint16_t>(x >> 8);
        EXPECT_EQ(opened(*dpf::eval_point(dpf::out<0, 8>, k0, q),
                      *dpf::eval_point(dpf::out<0, 8>, k1, q)),
            lane8 == alane ? uint8_t{7} : uint8_t{0}) << q;
        const auto lane12 = static_cast<uint16_t>(q >> 4);
        const auto blane = static_cast<uint16_t>(x >> 4);
        EXPECT_EQ(opened(*dpf::eval_point(dpf::out<1, 12>, k0, q),
                      *dpf::eval_point(dpf::out<1, 12>, k1, q)),
            lane12 == blane ? uint16_t{99} : uint16_t{0}) << q;
        EXPECT_EQ(opened(*dpf::eval_point(dpf::out<2>, k0, q),
                      *dpf::eval_point(dpf::out<2>, k1, q)),
            q == x ? uint32_t{5} : uint32_t{0}) << q;
    }

    const uint8_t alpha = 20;
    const uint32_t beta = 9;
    auto [p0, p1] = dpf::make_dpf(alpha, beta);
    // Weights follow the leaf-aligned interval layout: a uint32 leaf is 4
    // inputs, so [10, 30] is emitted from input 8 through 31.
    const uint8_t from = 10, to = 30;
    const uint8_t leaf = 4;
    const uint8_t begin = static_cast<uint8_t>(from & ~(leaf - 1u));
    const uint8_t end = static_cast<uint8_t>((to | (leaf - 1u)));
    std::vector<uint32_t> w(end - begin + 1);
    uint32_t expect = 0;
    for (uint8_t q = begin; q <= end; ++q)
    {
        w[q - begin] = q * 3u + 1u;
        if (q == alpha)
            expect = beta * w[q - begin];
    }
    auto memo0 = dpf::make_basic_interval_memoizer(p0, from, to);
    auto memo1 = dpf::make_basic_interval_memoizer(p1, from, to);
    auto ip0 = dpf::eval_inner_product<0>(p0, from, to, w, memo0);
    auto ip1 = dpf::eval_inner_product<0>(p1, from, to, w, memo1);
    EXPECT_EQ(opened(ip0, ip1), expect);
}

TEST(ContextBlast, XorWrapperAndGenevalIntegers)
{
    using X = dpf::xor_wrapper<uint32_t>;
    const uint16_t alpha = 0x00ff;
    const X y{0xa5a5a5a5u};
    auto [k0, k1] = dpf::make_dpf(alpha, y);
    for (uint16_t q : {uint16_t{0}, uint16_t{1}, alpha, uint16_t{0x0100},
             uint16_t{0x8000}, uint16_t{0xffff}, static_cast<uint16_t>(alpha ^ 1u)})
        EXPECT_EQ(opened(ev(k0, q), ev(k1, q)), q == alpha ? y : X{0});

    reset_roots();
    auto keys = dpf::make_dpf(int32_t{-2},
        dpf::root_sampler_t<dpf::prg::aes128>{take_root}, X{7});
    const int32_t x0 = 5;
    const int32_t secret = -2;
    const int32_t x1 = secret ^ x0;
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    auto g = dpf::geneval_point(x0, x1, secret,
        dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, X{7});
    HEDLEY_PRAGMA(GCC diagnostic pop)
    EXPECT_TRUE(g.leaf_live);
    EXPECT_EQ(g.live_levels, std::decay_t<decltype(keys.first)>::depth);
    EXPECT_EQ(opened(g.party0[0], g.party1[0]), X{7});
    EXPECT_EQ(g.party0[0], bare(ev(keys.first, secret)));

    const int32_t far = 100;
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    auto off = dpf::geneval_point(x0, x1, far,
        dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, X{7});
    HEDLEY_PRAGMA(GCC diagnostic pop)
    EXPECT_FALSE(off.leaf_live);
    EXPECT_GT(off.live_levels, 0u);
    EXPECT_LT(off.live_levels, off.correction_words.size());
    for (std::size_t i = 0; i < off.live_levels; ++i)
        EXPECT_EQ(std::memcmp(&off.correction_words[i], &keys.first.correction_word(i),
            sizeof(simde__m128i)), 0);
    EXPECT_EQ(opened(off.party0[0], off.party1[0]), X{0});

    const uint32_t u = 0x80000001u;
    const uint32_t u0 = 0x11111111u;
    const uint32_t u1 = u ^ u0;
    reset_roots();
    auto uk = dpf::make_dpf(u, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        uint64_t{99});
    reset_roots();
    const uint32_t qs[] = {0u, 1u, u, u ^ 1u, 0xffffffffu};
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    auto sq = dpf::geneval_sequence(u0, u1, std::begin(qs), std::end(qs),
        dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, uint64_t{99});
    HEDLEY_PRAGMA(GCC diagnostic pop)
    EXPECT_TRUE(sq.leaf_live);
    for (std::size_t i = 0; i < sq.live_levels; ++i)
        EXPECT_EQ(std::memcmp(&sq.correction_words[i], &uk.first.correction_word(i),
            sizeof(simde__m128i)), 0);
    for (std::size_t i = 0; i < 5; ++i)
        EXPECT_EQ(opened(sq.party0[i], sq.party1[i]), qs[i] == u ? uint64_t{99} : 0u);
}

TEST(ContextBlast, PrefixParityAgreesAcrossParties)
{
    const uint8_t alpha = 0x3c;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::bit::one);
    const std::array<uint8_t, 4> ends{0, 1, alpha, 255};
    auto [p0, n0] = grotto::prefix_parities(k0, ends);
    auto [p1, n1] = grotto::prefix_parities(k1, ends);
    (void)n0;
    (void)n1;
    for (std::size_t i = 0; i < ends.size(); ++i)
    {
        const bool bit = p0[i] ^ p1[i];
        bool acc = false;
        for (uint8_t q = 0; q < ends[i]; ++q)
            acc ^= (opened(ev(k0, q), ev(k1, q)) == dpf::bit::one);
        EXPECT_EQ(bit, acc) << int(ends[i]);
    }
}