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

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

// Prio's frequency count, with the one-hot vector replaced by a DPF, and
// the prefix walk Poplar uses for heavy hitters (Boneh, Boyle,
// Corrigan-Gibbs, Gilboa, Ishai). Classic Prio proves an encoding with a
// SNIP; this file is only the DPF-shaped encoding.
// field64 is libprio's Field64.
//
//   c++ -std=c++17 -march=native -I include -I thirdparty \
//       examples/applications/prio.cpp

namespace
{

constexpr int nbins = 256;

}  // namespace

int main()
{
    // Histogram. Each client sends one unit DPF at a secret bin.
    // Each server adds the expansion into its running share with
    // `eval_full_add_into` (no separate expansion buffer). The opened bin
    // is the count.
    const std::array<std::uint8_t, 4> bins{3, 3, 7, 3};
    std::vector<dpf::field64> h0(nbins);
    std::vector<dpf::field64> h1(nbins);
    for (std::uint8_t bin : bins)
    {
        auto [k0, k1] = dpf::make_dpf(bin, dpf::field64{1});
        dpf::eval_full_add_into(h0, k0);
        dpf::eval_full_add_into(h1, k1);
    }
    // Leaf shares are subtractive, so the opened bin is share0 - share1.
    const dpf::field64 c3 = h0[3] - h1[3];
    const dpf::field64 c7 = h0[7] - h1[7];
    const dpf::field64 c0 = h0[0] - h1[0];
    if (c3.raw() != 3 || c7.raw() != 1 || c0.raw() != 0)
    {
        std::cerr << "prio histogram\n";
        return 1;
    }

    // Heavy-hitter prefixes. idpf plants a 1 on each prefix length.
    // Length 1 is the high bit. 0xA0 and 0xB0 share 101; they split at bit 4.
    constexpr std::uint8_t left = 0xA0;
    constexpr std::uint8_t right = 0xB0;
    auto [a0, a1] = dpf::make_dpf(left,
        dpf::idpf(std::uint64_t{1}, std::uint64_t{1}, std::uint64_t{1}));
    auto [b0, b1] = dpf::make_dpf(right,
        dpf::idpf(std::uint64_t{1}, std::uint64_t{1}, std::uint64_t{1}));

    auto one = [](auto tag, auto k0, auto k1, std::uint8_t node) {
        return dpf::reconstruct(*dpf::eval_point(tag, k0, node),
            *dpf::eval_point(tag, k1, node));
    };
    auto count = [&](auto tag, std::uint8_t node) {
        return one(tag, a0, a1, node) + one(tag, b0, b1, node);
    };

    // out<0> is prefix length 1, out<1> length 2, out<2> length 3.
    const auto high = count(dpf::out<0, 1>, std::uint8_t{0x80});
    const auto low = count(dpf::out<0, 1>, std::uint8_t{0x00});
    const auto shared = count(dpf::out<2, 3>, std::uint8_t{0xA0});
    const auto split = count(dpf::out<2, 3>, std::uint8_t{0x80});
    if (high != 2 || low != 0 || shared != 2 || split != 0)
    {
        std::cerr << "prio prefixes " << high << " " << low << " " << shared
                  << " " << split << "\n";
        return 1;
    }

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

    std::cout << c3.raw() << "\n";
    return 0;
}