#include <cstddef>
#include <cstdint>
#include <iostream>
#include <vector>

#include "dpf.hpp"
#include "dpf/app_flow.hpp"
#include "dpf/app_plans.hpp"

// Splinter, the DPF query step (Wang, Yun, Goldwasser, Vaikuntanathan, and
// Zeldovich, NSDI 2017). Private queries on public data with two-server FSS.
// The client's private selector is a unit DPF at a secret attribute value.
// Each server dots that selector with a public aggregate column, so the
// answer is the SUM (or COUNT) for the private key without either server
// learning which key was asked.
//
//   c++ -std=c++17 -march=native -I include -I thirdparty \
//       examples/applications/splinter.cpp

int main()
{
    constexpr std::size_t domain = 256;  // attribute values

    // Public data, pre-aggregated by attribute: group_sum[v] is the SUM of a
    // value column over the rows whose attribute equals v.
    std::vector<std::uint64_t> group_sum(domain);
    std::vector<std::uint64_t> group_cnt(domain, 1);
    for (std::size_t v = 0; v < domain; ++v)
        group_sum[v] = (v * 37 + 11) % 1000;

    constexpr std::uint8_t secret_key = 88;  // the private WHERE value

    // One selector key per server. reconstruct = the two servers' shares.
    auto [k0, k1] = dpf::make_dpf(secret_key, std::uint64_t{1});

    // SELECT SUM(value) WHERE attribute = secret_key.
    const auto sum = dpf::reconstruct(
        dpf::eval_full_inner_product(dpf::paired, k0, group_sum),
        dpf::eval_full_inner_product(dpf::paired, k1, group_sum));
    if (sum != group_sum[secret_key])
    {
        std::cerr << "splinter sum\n";
        return 1;
    }

    // SELECT COUNT(*) WHERE attribute = secret_key is the same selector on an
    // all-ones column.
    const auto cnt = dpf::reconstruct(
        dpf::eval_full_inner_product(dpf::paired, k0, group_cnt),
        dpf::eval_full_inner_product(dpf::paired, k1, group_cnt));
    if (cnt != 1)
    {
        std::cerr << "splinter count\n";
        return 1;
    }

    {
        if (int rc = dpf::app::run_measured("splinter",
                dpf::protocol::fss_point_plan(0), 8))
            return rc;
    }

    std::cout << sum << "\n";
    return 0;
}