#include <gtest/gtest.h>

#include <algorithm>
#include <cstdint>
#include <functional>
#include <utility>
#include <vector>

#include "dpf.hpp"

namespace
{

using key0_t = decltype(dpf::make_dpf(std::uint16_t{0},
    dpf::idpf_ones<16>()).first);
using key1_t = decltype(dpf::make_dpf(std::uint16_t{0},
    dpf::idpf_ones<16>()).second);

void split_keys(const std::vector<std::uint16_t> & values,
    std::vector<key0_t> & k0, std::vector<key1_t> & k1)
{
    k0.clear();
    k1.clear();
    for (auto v : values)
    {
        auto [a, b] = dpf::make_dpf(v, dpf::idpf_ones<16>());
        k0.push_back(std::move(a));
        k1.push_back(std::move(b));
    }
}

}  // namespace

TEST(IdpfAgg, MaxAndKthAgainstPlaintext)
{
    const std::vector<std::uint16_t> values{3, 100, 7, 100, 42};
    std::vector<key0_t> k0;
    std::vector<key1_t> k1;
    split_keys(values, k0, k1);

    EXPECT_EQ(dpf::idpf_agg_max(k0, k1), 100u);
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, 1), 100u);
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, values.size()), 3u);

    auto sorted = values;
    std::sort(sorted.begin(), sorted.end(), std::greater<>{});
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, 3), sorted[2]);
}

TEST(IdpfAgg, AllEqual)
{
    const std::vector<std::uint16_t> values{0x00AA, 0x00AA, 0x00AA};
    std::vector<key0_t> k0;
    std::vector<key1_t> k1;
    split_keys(values, k0, k1);
    EXPECT_EQ(dpf::idpf_agg_max(k0, k1), 0x00AAu);
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, 1), 0x00AAu);
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, 3), 0x00AAu);
}

TEST(IdpfAgg, KEqualsOneAndN)
{
    const std::vector<std::uint16_t> values{1, 2, 3, 4};
    std::vector<key0_t> k0;
    std::vector<key1_t> k1;
    split_keys(values, k0, k1);
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, 1), 4u);
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, 4), 1u);
}

TEST(IdpfAgg, Duplicates)
{
    const std::vector<std::uint16_t> values{9, 9, 1, 9, 5};
    std::vector<key0_t> k0;
    std::vector<key1_t> k1;
    split_keys(values, k0, k1);
    EXPECT_EQ(dpf::idpf_agg_max(k0, k1), 9u);
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, 2), 9u);
    EXPECT_EQ(dpf::idpf_agg_kth(k0, k1, 4), 5u);
}

// Seam: max equals kth(k=1); both agree with a plaintext sort.
TEST(IdpfAgg, MaxEqualsKthOne)
{
    const std::vector<std::uint16_t> values{4, 90, 4, 15};
    std::vector<key0_t> k0;
    std::vector<key1_t> k1;
    split_keys(values, k0, k1);
    EXPECT_EQ(dpf::idpf_agg_max(k0, k1), dpf::idpf_agg_kth(k0, k1, 1));
}