#include <gtest/gtest.h>

#include "dpf.hpp"

#include <cstdint>
#include <cstring>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <vector>

namespace
{

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),
            static_cast<long long>(n * 5 + 1));
        n += 2;
        return v;
    }
    uint8_t bit() { return static_cast<uint8_t>(n++ & 1u); }
};

void reset_roots()
{
    g_ri = 0;
    for (int i = 0; i < 4; ++i)
        g_roots[i] = simde_mm_set_epi64x(0x1000 * (i + 1), 0xA000u + i);
}

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 Key, typename In>
auto ev(const Key & key, In x)
{
    return bare(*dpf::eval_point(key, x));
}

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

unsigned as_u(dpf::twobit v) { return static_cast<unsigned>(v); }
unsigned as_u(dpf::nyble v) { return static_cast<unsigned>(v); }
unsigned as_u(dpf::bit v) { return static_cast<unsigned>(static_cast<bool>(v)); }

template <typename Lane>
Lane lane_of(unsigned v)
{
    if constexpr (std::is_same_v<Lane, dpf::bit>)
        return static_cast<dpf::bit>(v & 1u);
    else if constexpr (std::is_same_v<Lane, dpf::twobit>)
        return dpf::to_twobit(v);
    else
        return dpf::to_nyble(v);
}

template <typename Lane>
void expect_share_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 Lane av = lane_of<Lane>(a);
            const Lane bv = lane_of<Lane>(b);
            auto s0 = dpf::subtractive_share<Lane, 0>::from_raw(av);
            auto s1 = dpf::subtractive_share<Lane, 1>::from_raw(bv);
            EXPECT_EQ(dpf::reconstruct(s0, s1), av - bv) << a << " - " << b;

            auto p0 = dpf::additive_share<Lane, 0>::from_raw(av);
            auto p1 = dpf::additive_share<Lane, 1>::from_raw(bv);
            EXPECT_EQ(dpf::reconstruct(p0, p1), av + bv);

            auto c0 = s0.as_additive();
            auto c1 = s1.as_additive();
            EXPECT_EQ(dpf::reconstruct(c0, c1), dpf::reconstruct(s0, s1));

            // Additive lhs + subtractive rhs. Party 1 subtracts the
            // foreign share: raw1 = b - b, opened sum is (a+a) + (b-b).
            auto x0 = p0 + s0;
            auto x1 = p1 + s1;
            EXPECT_EQ(x0.raw(), av + av);
            EXPECT_EQ(x1.raw(), bv - bv);
            EXPECT_EQ(dpf::reconstruct(x0, x1), (av + av) + (bv - bv));

            auto y0 = s0;
            auto y1 = s1;
            y0 += av;
            y1 += av;
            EXPECT_EQ(y0.raw(), av + av);
            EXPECT_EQ(y1.raw(), bv);

            if constexpr (!std::is_same_v<Lane, dpf::bit>)
            {
                auto z0 = s0 * bv;
                EXPECT_EQ(z0.raw(), av * bv);
            }
        }
    }

    auto [z0, z1] = dpf::make_subtractive_shares(lane_of<Lane>(0));
    EXPECT_EQ(dpf::reconstruct(z0, z1), lane_of<Lane>(0));
    auto [o0, o1] = dpf::make_subtractive_shares(lane_of<Lane>(1));
    EXPECT_EQ(dpf::reconstruct(o0, o1), lane_of<Lane>(1));
    EXPECT_EQ(as_u(o1.raw()), 0u);
}

template <typename In, typename Lane>
void expect_every_point(In alpha, Lane y)
{
    auto [k0, k1] = dpf::make_dpf(alpha, y);
    constexpr auto bits = dpf::utils::bitlength_of_v<In>;
    const uint64_t n = uint64_t{1} << bits;
    for (uint64_t i = 0; i < n; ++i)
    {
        In q;
        std::memcpy(&q, &i, sizeof(q));
        const Lane got = opened(ev(k0, q), ev(k1, q));
        if (q == alpha)
            EXPECT_EQ(got, y) << "hit " << i;
        else
            EXPECT_EQ(got, Lane{}) << "miss " << i;
    }
}

template <typename In, typename Lane>
void expect_interval(In alpha, Lane y, In from, In to)
{
    auto [k0, k1] = dpf::make_dpf(alpha, y);
    auto a = dpf::eval_interval(k0, from, to);
    auto b = dpf::eval_interval(k1, from, to);
    auto p0 = std::begin(a.second);
    auto p1 = std::begin(b.second);
    for (In q = from; ; ++q)
    {
        const Lane got = opened(*p0, *p1);
        EXPECT_EQ(got, q == alpha ? y : Lane{}) << "interval q";
        EXPECT_EQ(got, opened(ev(k0, q), ev(k1, q)));
        ++p0;
        ++p1;
        if (q == to)
            break;
    }
    EXPECT_EQ(p0, std::end(a.second));
    EXPECT_EQ(p1, std::end(b.second));
}

template <typename Key>
void expect_live_words(const Key & key, const auto & g)
{
    ASSERT_LE(g.live_levels, g.correction_words.size());
    for (std::size_t i = 0; i < g.live_levels; ++i)
    {
        EXPECT_EQ(std::memcmp(&g.correction_words[i], &key.correction_word(i),
            sizeof(simde__m128i)), 0) << i;
        EXPECT_EQ(g.correction_advice[i], key.correction_advice(i)) << i;
    }
    if (g.leaf_live)
        EXPECT_EQ(std::memcmp(&g.leaf, &key.template leaf<0>(), sizeof(g.leaf)), 0);
    else if (g.live_levels < g.correction_words.size())
    {
        EXPECT_NE(std::memcmp(&g.correction_words[g.live_levels],
            &key.correction_word(g.live_levels), sizeof(simde__m128i)), 0);
    }
}

}  // namespace

TEST(LaneBlast, ParseStreamsAndLimits)
{
    using namespace dpf::literals::twobit;
    using namespace dpf::literals::nyble;

    EXPECT_EQ(std::numeric_limits<dpf::twobit>::digits, 2);
    EXPECT_EQ(std::numeric_limits<dpf::twobit>::min(), dpf::twobit::zero);
    EXPECT_EQ(std::numeric_limits<dpf::twobit>::max(), dpf::twobit::three);
    EXPECT_EQ(std::numeric_limits<dpf::nyble>::digits, 4);
    EXPECT_EQ(std::numeric_limits<dpf::nyble>::max(), dpf::nyble{15});
    EXPECT_EQ(std::numeric_limits<dpf::nyble>::lowest(), dpf::nyble{0});

    EXPECT_EQ(0_twobit, dpf::twobit::zero);
    EXPECT_EQ(3_twobit, dpf::twobit::three);
    EXPECT_EQ(4_twobit, dpf::twobit::zero);
    EXPECT_EQ(7_twobit, dpf::twobit::three);
    EXPECT_EQ(15_nyble, dpf::nyble{15});
    EXPECT_EQ(16_nyble, dpf::nyble{0});
    EXPECT_EQ(31_nyble, dpf::nyble{15});

    EXPECT_EQ(dpf::to_twobit('0'), dpf::twobit::zero);
    EXPECT_EQ(dpf::to_twobit('3'), dpf::twobit::three);
    EXPECT_THROW(dpf::to_twobit('4'), std::domain_error);
    EXPECT_THROW(dpf::to_twobit('a'), std::domain_error);
    EXPECT_EQ(dpf::to_nyble('0'), dpf::nyble{0});
    EXPECT_EQ(dpf::to_nyble('9'), dpf::nyble{9});
    EXPECT_EQ(dpf::to_nyble('a'), dpf::nyble{10});
    EXPECT_EQ(dpf::to_nyble('A'), dpf::nyble{10});
    EXPECT_EQ(dpf::to_nyble('f'), dpf::nyble{15});
    EXPECT_EQ(dpf::to_nyble('F'), dpf::nyble{15});
    EXPECT_THROW(dpf::to_nyble('g'), std::domain_error);
    EXPECT_THROW(dpf::to_nyble('G'), std::domain_error);
    EXPECT_THROW(dpf::to_nyble('/'), std::domain_error);

    {
        std::stringstream in("3");
        dpf::twobit v = dpf::twobit::zero;
        ASSERT_TRUE(in >> v);
        EXPECT_EQ(v, dpf::twobit::three);
        std::stringstream bad("4");
        ASSERT_FALSE(bad >> v);
        EXPECT_TRUE(bad.fail());
    }
    {
        std::stringstream in("a");
        dpf::nyble v{0};
        ASSERT_TRUE(in >> v);
        EXPECT_EQ(v, dpf::nyble{10});
        std::stringstream bad("g");
        ASSERT_FALSE(bad >> v);
        EXPECT_TRUE(bad.fail());
        std::stringstream empty;
        ASSERT_FALSE(empty >> v);
    }
    {
        std::stringstream out;
        out << dpf::twobit::two << dpf::nyble{12};
        EXPECT_EQ(out.str(), "2c");
    }
}

TEST(LaneBlast, ShareWrapsInTheLaneRing)
{
    expect_share_ring<dpf::bit>();
    expect_share_ring<dpf::twobit>();
    expect_share_ring<dpf::nyble>();

    auto z0 = dpf::subtractive_share<dpf::twobit, 0>::from_raw(dpf::twobit::zero);
    auto z1 = dpf::subtractive_share<dpf::twobit, 1>::from_raw(dpf::twobit::one);
    EXPECT_EQ(dpf::reconstruct(z0, z1), dpf::twobit::three);

    auto n0 = dpf::subtractive_share<dpf::nyble, 0>::from_raw(dpf::nyble{0});
    auto n1 = dpf::subtractive_share<dpf::nyble, 1>::from_raw(dpf::nyble{1});
    EXPECT_EQ(dpf::reconstruct(n0, n1), dpf::nyble{15});

    auto b0 = dpf::subtractive_share<dpf::bit, 0>::from_raw(dpf::bit{false});
    auto b1 = dpf::subtractive_share<dpf::bit, 1>::from_raw(dpf::bit{true});
    EXPECT_EQ(static_cast<bool>(dpf::reconstruct(b0, b1)), true);
}

TEST(LaneBlast, DepositDoesNotBleedIntoTheNextLane)
{
    simde__m128i node{};
    dpf::packed::deposit_lane(node, 0, dpf::twobit::three);
    dpf::packed::deposit_lane(node, 1, dpf::twobit::one);
    dpf::packed::deposit_lane(node, 2, dpf::twobit::two);
    dpf::packed::deposit_lane(node, 3, dpf::twobit::zero);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(node, 0), dpf::twobit::three);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(node, 1), dpf::twobit::one);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(node, 2), dpf::twobit::two);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(node, 3), dpf::twobit::zero);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(node, 4), dpf::twobit::zero);

    dpf::packed::deposit_lane(node, 63, dpf::twobit::two);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(node, 63), dpf::twobit::two);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(node, 62), dpf::twobit::zero);

    simde__m256i twos{};
    dpf::packed::deposit_lane(twos, 127, dpf::twobit::three);
    dpf::packed::deposit_lane(twos, 126, dpf::twobit::one);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(twos, 127), dpf::twobit::three);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(twos, 126), dpf::twobit::one);
    EXPECT_EQ(dpf::packed::extract_lane<dpf::twobit>(twos, 125), dpf::twobit::zero);

    simde__m256i nibs{};
    dpf::packed::deposit_lane(nibs, 0, dpf::nyble{0x0f});
    dpf::packed::deposit_lane(nibs, 1, dpf::nyble{0x01});
    dpf::packed::deposit_lane(nibs, 63, dpf::nyble{0x0a});
    EXPECT_EQ(dpf::packed::extract_lane<dpf::nyble>(nibs, 0), dpf::nyble{0x0f});
    EXPECT_EQ(dpf::packed::extract_lane<dpf::nyble>(nibs, 1), dpf::nyble{0x01});
    EXPECT_EQ(dpf::packed::extract_lane<dpf::nyble>(nibs, 2), dpf::nyble{0});
    EXPECT_EQ(dpf::packed::extract_lane<dpf::nyble>(nibs, 63), dpf::nyble{0x0a});
    EXPECT_EQ(dpf::packed::extract_lane<dpf::nyble>(nibs, 62), dpf::nyble{0});
}

TEST(LaneBlast, Uint8TwobitEveryInputAndEveryShape)
{
    using in_t = uint8_t;
    using out_t = dpf::twobit;
    const out_t y = dpf::twobit::three;
    for (in_t alpha : {in_t{0}, in_t{1}, in_t{63}, in_t{64}, in_t{127}, in_t{128}, in_t{255}})
        expect_every_point(alpha, y);

    expect_interval<in_t, out_t>(40, y, 40, 40);
    expect_interval<in_t, out_t>(40, y, 0, 255);
    expect_interval<in_t, out_t>(0, y, 1, 63);
    expect_interval<in_t, out_t>(64, dpf::twobit::one, 65, 70);
    expect_interval<in_t, out_t>(255, dpf::twobit::two, 192, 255);

    auto [k0, k1] = dpf::make_dpf(in_t{40}, y);
    auto [fb0, fi0] = dpf::eval_full(k0);
    auto [fb1, fi1] = dpf::eval_full(k1);
    auto f0 = std::begin(fi0);
    auto f1 = std::begin(fi1);
    for (int q = 0; q < 256; ++q, ++f0, ++f1)
    {
        EXPECT_EQ(opened(*f0, *f1), q == 40 ? y : out_t{}) << q;
    }
    EXPECT_EQ(f0, std::end(fi0));

    const in_t unsorted[] = {255, 40, 0};
    EXPECT_THROW(dpf::make_sequence_recipe(k0, std::begin(unsorted), std::end(unsorted)),
        std::runtime_error);
    const in_t seq[] = {0, 40, 40, 41, 104, 255};
    auto recipe = dpf::make_sequence_recipe(k0, std::begin(seq), std::end(seq));
    auto s0 = dpf::eval_sequence(k0, recipe);
    auto s1 = dpf::eval_sequence(k1, recipe);
    auto p0 = std::begin(s0.second);
    auto p1 = std::begin(s1.second);
    for (in_t q : seq)
    {
        EXPECT_EQ(opened(*p0, *p1), q == 40 ? y : out_t{}) << int(q);
        ++p0;
        ++p1;
    }
}

TEST(LaneBlast, Uint8NybleAndBitSameShapes)
{
    using in_t = uint8_t;
    expect_every_point<in_t>(in_t{0}, dpf::nyble{0});
    expect_every_point<in_t>(in_t{31}, dpf::nyble{15});
    expect_every_point<in_t>(in_t{32}, dpf::nyble{1});
    expect_every_point<in_t>(in_t{255}, dpf::nyble{10});
    expect_interval<in_t>(in_t{31}, dpf::nyble{15}, in_t{0}, in_t{30});
    expect_interval<in_t>(in_t{31}, dpf::nyble{15}, in_t{32}, in_t{63});
    expect_every_point<in_t>(in_t{0}, dpf::bit::one);
    expect_every_point<in_t>(in_t{127}, dpf::bit::one);
    expect_every_point<in_t>(in_t{128}, dpf::bit{false});
    expect_interval<in_t>(in_t{128}, dpf::bit::one, in_t{129}, in_t{255});
}

TEST(LaneBlast, SignedInputPackedLeaf)
{
    using in_t = int8_t;
    const auto y = dpf::twobit::two;
    for (in_t alpha : {in_t{-128}, in_t{-1}, in_t{0}, in_t{1}, in_t{127}})
        expect_every_point(alpha, y);
    expect_interval<in_t>(in_t{-1}, dpf::nyble{15}, in_t{-3}, in_t{3});
    expect_interval<in_t>(in_t{-128}, dpf::nyble{1}, in_t{-128}, in_t{127});

    auto [k0, k1] = dpf::make_dpf(in_t{-5}, y);
    auto a = dpf::eval_full(k0);
    auto b = dpf::eval_full(k1);
    auto p0 = std::begin(a.second);
    auto p1 = std::begin(b.second);
    for (int q = -128; q <= 127; ++q, ++p0, ++p1)
    {
        EXPECT_EQ(opened(*p0, *p1),
            static_cast<in_t>(q) == in_t{-5} ? y : dpf::twobit{}) << q;
    }
}

TEST(LaneBlast, GenevalMatchesKeyOnPackedLanes)
{
    using in_t = uint8_t;
    using out_t = dpf::twobit;
    const out_t y = dpf::twobit::three;
    const in_t alpha = 40;
    const in_t x0 = 0x11;
    const in_t x1 = static_cast<in_t>(alpha ^ x0);

    reset_roots();
    auto keys = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root}, y);

    reset_roots();
    auto on = dpf::geneval_point(x0, x1, alpha, dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
    expect_live_words(keys.first, on);
    EXPECT_TRUE(on.leaf_live);
    EXPECT_EQ(opened(on.party0[0], on.party1[0]), y);
    EXPECT_EQ(on.party0[0], ev(keys.first, alpha));
    EXPECT_EQ(on.party1[0], ev(keys.second, alpha));

    const in_t neighbor = static_cast<in_t>(alpha ^ 1u);
    reset_roots();
    auto lane = dpf::geneval_point(x0, x1, neighbor,
        dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
    EXPECT_TRUE(lane.leaf_live);
    expect_live_words(keys.first, lane);
    EXPECT_EQ(opened(lane.party0[0], lane.party1[0]), out_t{});
    EXPECT_EQ(lane.party0[0], ev(keys.first, neighbor));

    const in_t far = 200;
    reset_roots();
    auto off = dpf::geneval_point(x0, x1, far, dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
    EXPECT_FALSE(off.leaf_live);
    EXPECT_LT(off.live_levels, off.correction_words.size());
    expect_live_words(keys.first, off);
    EXPECT_EQ(opened(off.party0[0], off.party1[0]), out_t{});
    EXPECT_NE(off.party0[0], ev(keys.first, far));

    reset_roots();
    auto iv = dpf::geneval_interval(x0, x1, in_t{41}, in_t{50},
        dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
    EXPECT_TRUE(iv.leaf_live);
    expect_live_words(keys.first, iv);
    ASSERT_EQ(iv.party0.size(), 10u);
    for (std::size_t i = 0; i < iv.party0.size(); ++i)
        EXPECT_EQ(opened(iv.party0[i], iv.party1[i]), out_t{});

    reset_roots();
    auto full = dpf::geneval_full(x0, x1, dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
    EXPECT_EQ(full.party0.size(), 256u);
    EXPECT_TRUE(full.leaf_live);
    expect_live_words(keys.first, full);
    for (int q = 0; q < 256; ++q)
    {
        EXPECT_EQ(opened(full.party0[q], full.party1[q]),
            static_cast<in_t>(q) == alpha ? y : out_t{}) << q;
        EXPECT_EQ(full.party0[q], ev(keys.first, static_cast<in_t>(q)));
    }

    const in_t seq[] = {0, 255, alpha, neighbor, alpha};
    reset_roots();
    auto sq = dpf::geneval_sequence(x0, x1, std::begin(seq), std::end(seq),
        dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
    EXPECT_TRUE(sq.leaf_live);
    expect_live_words(keys.first, sq);
    for (std::size_t i = 0; i < 5; ++i)
        EXPECT_EQ(opened(sq.party0[i], sq.party1[i]), seq[i] == alpha ? y : out_t{});
}

TEST(LaneBlast, GenevalNybleArithAndSigned)
{
    using out_t = dpf::nyble;
    const out_t y{0x0c};

    {
        using in_t = uint8_t;
        const in_t secret = 10;
        const in_t a0 = 200;
        const in_t a1 = 66;
        ASSERT_EQ(static_cast<in_t>(a0 + a1), secret);
        reset_roots();
        auto keys = dpf::make_dpf(secret,
            dpf::root_sampler_t<dpf::prg::aes128>{take_root}, y);
        reset_roots();
        auto g = dpf::geneval_point(dpf::arith_input, a0, a1, secret,
            dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
        EXPECT_TRUE(g.leaf_live);
        expect_live_words(keys.first, g);
        EXPECT_EQ(opened(g.party0[0], g.party1[0]), y);

        reset_roots();
        auto miss = dpf::geneval_point(dpf::arith_input, a0, a1, in_t{250},
            dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
        EXPECT_EQ(opened(miss.party0[0], miss.party1[0]), out_t{});
        expect_live_words(keys.first, miss);
    }
    {
        using in_t = int8_t;
        const in_t alpha = -20;
        const in_t x0 = 3;
        const in_t x1 = static_cast<in_t>(alpha ^ x0);
        reset_roots();
        auto keys = dpf::make_dpf(alpha,
            dpf::root_sampler_t<dpf::prg::aes128>{take_root}, y);
        reset_roots();
        auto full = dpf::geneval_full(x0, x1, dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
        EXPECT_EQ(full.party0.size(), 256u);
        EXPECT_TRUE(full.leaf_live);
        expect_live_words(keys.first, full);
        constexpr auto to_int = dpf::utils::to_integral_type<in_t>{};
        for (int q = -128; q <= 127; ++q)
        {
            in_t v = static_cast<in_t>(q);
            const auto i = static_cast<std::size_t>(to_int(v));
            EXPECT_EQ(opened(full.party0[i], full.party1[i]), v == alpha ? y : out_t{}) << q;
            EXPECT_EQ(full.party0[i], ev(keys.first, v));
        }
        reset_roots();
        auto iv = dpf::geneval_interval(x0, x1, in_t{-2}, in_t{2},
            dpf::ds_randomness<simde__m128i (*)(), Pad>{take_root, Pad{}}, y);
        ASSERT_EQ(iv.party0.size(), 5u);
        for (int q = -2; q <= 2; ++q)
        {
            EXPECT_EQ(opened(iv.party0[static_cast<std::size_t>(q + 2)],
                iv.party1[static_cast<std::size_t>(q + 2)]), out_t{});
        }
    }
}