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

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

// Mastic, the DPF step (private weighted heavy-hitters / attribute-based
// metrics). Each client keys an incremental DPF whose payload is its weight
// instead of a plain 1. Servers sum the weighted prefix shares at each depth
// and keep the heavy prefixes. This is Poplar's prefix walk with a weight
// payload; VIDPF path-consistency is `verify_idpf_path` below.
//
//   c++ -std=c++17 -march=native -I include -I thirdparty \
//       examples/applications/mastic.cpp

namespace
{

// Weighted counts of all 2^length prefixes, summed over the two clients.
template <typename Tag, typename A0, typename A1, typename Kb0, typename Kb1>
std::vector<std::uint64_t> weighted_level(Tag tag, std::size_t nprefix,
    const A0 & a0, const A1 & a1, const Kb0 & b0, const Kb1 & b1)
{
    auto [ba0, ia0] = dpf::eval_prefixes(tag, a0);
    auto [ba1, ia1] = dpf::eval_prefixes(tag, a1);
    auto [bb0, ib0] = dpf::eval_prefixes(tag, b0);
    auto [bb1, ib1] = dpf::eval_prefixes(tag, b1);
    std::vector<std::uint64_t> w(nprefix);
    for (std::size_t p = 0; p < nprefix; ++p)
        w[p] = dpf::reconstruct(ba0[p], ba1[p])
             + dpf::reconstruct(bb0[p], bb1[p]);
    return w;
}

std::vector<dpf::fp61> path_challenges(std::size_t n)
{
    std::vector<dpf::fp61> rs(n);
    for (auto & r : rs)
        r = dpf::uniform_sample<dpf::fp61>();
    return rs;
}

}  // namespace

int main()
{
    // Two clients report strings 0xA0 and 0xB0 (both begin "101"), with
    // weights 5 and 3. A heavy-hitter threshold of 6 should keep prefix 101.
    auto [a0, a1] = dpf::make_dpf(std::uint8_t{0xA0},
        dpf::idpf(std::uint64_t{5}, std::uint64_t{5}, std::uint64_t{5}));
    auto [b0, b1] = dpf::make_dpf(std::uint8_t{0xB0},
        dpf::idpf(std::uint64_t{3}, std::uint64_t{3}, std::uint64_t{3}));

    // One-time VIDPF path check per client (weight-1 / parent consistency).
    const auto rs = path_challenges(dpf::path_sketch_challenge_count(3));
    if (!dpf::verify_idpf_path<3>(a0, a1, rs)
        || !dpf::verify_idpf_path<3>(b0, b1, rs))
    {
        std::cerr << "mastic path sketch\n";
        return 1;
    }

    // Length 1: prefix "1" carries the full weight 8; "0" carries 0.
    const auto lvl1 = weighted_level(dpf::out<0, 1>, 2, a0, a1, b0, b1);
    if (lvl1[1] != 8 || lvl1[0] != 0)
    {
        std::cerr << "mastic level1\n";
        return 1;
    }

    // Length 3: prefix 101 (=5) is the heavy hitter with weight 8.
    const auto lvl3 = weighted_level(dpf::out<2, 3>, 8, a0, a1, b0, b1);
    constexpr std::uint64_t threshold = 6;
    std::size_t heavy = 0, nheavy = 0;
    for (std::size_t p = 0; p < lvl3.size(); ++p)
        if (lvl3[p] >= threshold) { heavy = p; ++nheavy; }
    if (nheavy != 1 || heavy != 0b101 || lvl3[0b101] != 8)
    {
        std::cerr << "mastic heavy\n";
        return 1;
    }

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

    std::cout << lvl3[0b101] << "\n";
    return 0;
}