#include <gtest/gtest.h>

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

#include "dpf.hpp"

// Walk helpers fold a small operation into an existing DPF walk. See
// dpf/eval_walk.hpp.

TEST(WalkHelpers, RotatePairedInnerProduct)
{
    constexpr std::size_t n = 256;
    constexpr std::uint8_t r = 50;
    constexpr std::uint8_t a = 42;
    const std::size_t s = (n - static_cast<std::size_t>(
        static_cast<std::uint8_t>(r - a))) % n;

    std::vector<std::uint64_t> table(n);
    for (std::size_t i = 0; i < n; ++i) table[i] = i * 7 + 1;

    auto [k0, k1] = dpf::make_dpf(r, std::uint64_t{1});

    std::vector<std::uint64_t> manual(n);
    for (std::size_t i = 0; i < n; ++i) manual[i] = table[(i + s) % n];
    const auto ref = dpf::reconstruct(
        dpf::eval_full_inner_product(dpf::paired, k0, manual),
        dpf::eval_full_inner_product(dpf::paired, k1, manual));

    const auto got = dpf::reconstruct(
        dpf::eval_full_inner_product(dpf::paired, k0, table, dpf::rotate{s}),
        dpf::eval_full_inner_product(dpf::paired, k1, table, dpf::rotate{s}));

    EXPECT_EQ(got, ref);
    EXPECT_EQ(got, table[(r + s) % n]);
}

TEST(WalkHelpers, FullAddIntoHistogram)
{
    constexpr std::size_t n = 256;
    std::vector<dpf::field64> h0(n), h1(n);
    for (auto bin : std::array<std::uint8_t, 3>{3, 3, 7})
    {
        auto [k0, k1] = dpf::make_dpf(bin, dpf::field64{1});
        dpf::eval_full_add_into(h0, k0);
        dpf::eval_full_add_into(h1, k1);
    }
    EXPECT_EQ((h0[3] - h1[3]).raw(), 2u);
    EXPECT_EQ((h0[7] - h1[7]).raw(), 1u);
    EXPECT_EQ((h0[0] - h1[0]).raw(), 0u);
}

TEST(WalkHelpers, FullAddIntoRotate)
{
    constexpr std::size_t n = 256;
    constexpr std::uint8_t r = 10;
    constexpr std::size_t shift = 32;
    std::vector<std::uint64_t> d0(n, 0), d1(n, 0);
    auto [w0, w1] = dpf::make_dpf(r, std::uint64_t{7});
    dpf::eval_full_add_into(d0, w0, dpf::rotate{shift});
    dpf::eval_full_add_into(d1, w1, dpf::rotate{shift});
    const std::size_t hot = (r + shift) % n;
    EXPECT_EQ(d0[hot] - d1[hot], 7u);
    EXPECT_EQ(d0[r] - d1[r], 0u);
}

TEST(WalkHelpers, FullAddIntoSketch)
{
    constexpr std::size_t n = 256;
    constexpr std::uint8_t address = 9;
    const dpf::fp61 message{42};
    auto [k0, k1] = dpf::make_dpf(address, message, dpf::extractable{});

    std::vector<dpf::fp61> box0(n), box1(n);
    std::vector<dpf::fp61> challenge(n);
    for (std::size_t i = 0; i < n; ++i)
        challenge[i] = dpf::fp61{static_cast<std::uint64_t>(i + 1)};

    dpf::sketch_share s0{}, s1{};
    auto sk0 = dpf::sketch(s0, challenge);
    auto sk1 = dpf::sketch(s1, challenge);
    dpf::eval_full_add_into(box0, k0, sk0);
    dpf::eval_full_add_into(box1, k1, sk1);

    EXPECT_EQ(box0[address] - box1[address], message);
    EXPECT_EQ((box0[0] - box1[0]).raw(), 0u);
    EXPECT_TRUE(dpf::sketch_verify(s0, s1));
}

TEST(WalkHelpers, CyclicShift)
{
    std::vector<int> v(8);
    for (int i = 0; i < 8; ++i) v[i] = i;
    auto sh = dpf::cyclic_shift(v, 3);
    for (std::size_t i = 0; i < v.size(); ++i)
        EXPECT_EQ(sh[i], v[(i + v.size() - 3) % v.size()]);
}

TEST(WalkHelpers, PackBitColumns)
{
    constexpr std::size_t n = 256;
    constexpr std::uint8_t alpha = 42;
    auto [a0, a1] = dpf::make_dpf(alpha, dpf::bit::one);
    auto [b0, b1] = dpf::make_dpf(alpha, dpf::bit::one);
    const auto p0 = dpf::pack_bit_columns(a0, b0);
    const auto p1 = dpf::pack_bit_columns(a1, b1);
    for (std::size_t row = 0; row < n; ++row)
    {
        const std::uint64_t opened = p0[row] ^ p1[row];
        EXPECT_EQ(opened, row == alpha ? 0b11u : 0u) << "row " << row;
    }
}

TEST(WalkHelpers, ModBitColumns)
{
    constexpr std::size_t n = 256;
    constexpr std::uint8_t alpha = 42;
    auto [a0, a1] = dpf::make_dpf(alpha, dpf::bit::one);
    auto [b0, b1] = dpf::make_dpf(alpha, dpf::bit::one);
    auto [c0, c1] = dpf::make_dpf(alpha, dpf::bit::one);

    auto expect_mod = [&](auto modulus, const auto & k0, const auto & k1, const auto & k2)
    {
        constexpr unsigned m = decltype(modulus)::value;
        const auto packed = dpf::pack_bit_columns(k0, k1, k2);
        const auto got = dpf::mod_bit_columns<m>(k0, k1, k2);
        ASSERT_EQ(got.size(), n);
        for (std::size_t row = 0; row < n; ++row)
            EXPECT_EQ(static_cast<unsigned>(got[row]), packed[row] % m) << "row " << row;
    };
    expect_mod(std::integral_constant<unsigned, 5>{}, a0, b0, c0);
    expect_mod(std::integral_constant<unsigned, 128>{}, a0, b0, c0);
    expect_mod(std::integral_constant<unsigned, 200>{}, a1, b1, c1);
    expect_mod(std::integral_constant<unsigned, 1000>{}, a0, b1, c0);
    expect_mod(std::integral_constant<unsigned, 32768>{}, a0, b0, c0);

    const auto one = dpf::mod_bit_columns<7>(a0);
    const auto bits = dpf::pack_bit_columns(a0);
    for (std::size_t row = 0; row < n; ++row)
        EXPECT_EQ(static_cast<unsigned>(one[row]), bits[row] % 7u);
}

TEST(WalkHelpers, UnitSignBitLeaf)
{
    constexpr std::uint8_t r = 50;
    int w0 = 0, w1 = 0;
    auto [k0, k1] = dpf::make_dpf(r, dpf::bit::one, dpf::unit_sign{w0, w1});
    const int sign = w0 - w1;
    EXPECT_TRUE(sign == 1 || sign == -1);
    EXPECT_EQ(w0, static_cast<int>(static_cast<bool>((*dpf::eval_point(k0, r)).raw())));
    EXPECT_EQ(w1, static_cast<int>(static_cast<bool>((*dpf::eval_point(k1, r)).raw())));
    // The key still opens as an ordinary unit bit DPF.
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(k0, r), *dpf::eval_point(k1, r)),
        dpf::bit::one);
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(k0, std::uint8_t{0}),
        *dpf::eval_point(k1, std::uint8_t{0})), dpf::bit::zero);
}

TEST(WalkHelpers, CmpFullInnerProduct)
{
    constexpr std::size_t n = 256;
    constexpr std::uint8_t threshold = 50;
    auto [c0, c1] = dpf::make_dpf(threshold, dpf::gt(std::uint64_t{1}));
    std::vector<std::uint64_t> w(n, 0);
    w[90] = 8; w[200] = 3; w[30] = 5;  // 90,200 are > 50
    const auto h0 = dpf::eval_full_inner_product(dpf::cmp, c0, w);
    const auto h1 = dpf::eval_full_inner_product(dpf::cmp, c1, w);
    EXPECT_EQ(dpf::reconstruct_cmp_halves(h0, h1).raw(), 11u);
}

TEST(WalkHelpers, Dpf3FullAddInto)
{
    constexpr std::size_t n = 256;
    std::vector<dpf::fp61> l1(n), l2(n), l3(n);
    auto [k1, k2, k3] = dpf::make_dpf3(std::uint8_t{5}, dpf::fp61{100});
    dpf::eval_full_add_into(l1, k1);
    dpf::eval_full_add_into(l2, k2);
    dpf::eval_full_add_into(l3, k3);
    auto open2 = [](dpf::fp61 a, dpf::fp61 b) {
        return dpf::shamir3::reconstruct(dpf::shamir3::share{1, a},
            dpf::shamir3::share{2, b});
    };
    EXPECT_EQ(open2(l1[5], l2[5]), dpf::fp61{100});
    EXPECT_EQ(open2(l1[0], l2[0]).raw(), 0u);
}

TEST(WalkHelpers, PrefixesAndPrefixInnerProduct)
{
    constexpr std::uint8_t left = 0xA0;  // prefix "101"
    auto [k0, k1] = dpf::make_dpf(left,
        dpf::idpf(std::uint64_t{1}, std::uint64_t{1}, std::uint64_t{1}));

    auto [buf0, it0] = dpf::eval_prefixes(dpf::out<0, 1>, k0);
    auto [buf1, it1] = dpf::eval_prefixes(dpf::out<0, 1>, k1);
    EXPECT_EQ(dpf::reconstruct(buf0[0], buf1[0]), 0u);
    EXPECT_EQ(dpf::reconstruct(buf0[1], buf1[1]), 1u);

    std::vector<std::uint64_t> vals1{11, 23};
    const auto dot1 =
        dpf::eval_prefix_inner_product(dpf::out<0, 1>, k0, vals1)
        - dpf::eval_prefix_inner_product(dpf::out<0, 1>, k1, vals1);
    EXPECT_EQ(dot1, vals1[1]);

    std::vector<std::uint64_t> vals3(8);
    for (std::size_t i = 0; i < 8; ++i) vals3[i] = 100 + i;
    const auto dot3 =
        dpf::eval_prefix_inner_product(dpf::out<2, 3>, k0, vals3)
        - dpf::eval_prefix_inner_product(dpf::out<2, 3>, k1, vals3);
    EXPECT_EQ(dot3, vals3[0b101]);
}