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

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

// BitMore, the DPF query only (Hafiz and Henry, PoPETs 2019 §5.2).
// ell = 2^L servers. The client samples L independent 1-bit DPFs at the
// secret row. Server j, whose label bits are j_{L-1} ... j_0, receives
// key j_e of DPF e and expands it. Concatenating those bits per row is
// the query string the information-theoretic response then consumes.
//
// `dpf::pack_bit_columns(keys...)` runs the full-domain bit walk once per
// key and writes lane e = key e into one integer per row, so the server
// loop reads `symbol[row]` instead of unpacking one int per bit.
//
//   c++ -std=c++17 -march=native -I include -I thirdparty \
//       examples/applications/bitmore.cpp

namespace
{

constexpr int bits_l = 2;
constexpr int nservers = 1 << bits_l;
constexpr std::size_t nrows = 256;

}  // namespace

int main()
{
    constexpr std::uint8_t alpha = 42;

    // L independent 1-bit DPFs at the secret row; keep both parties' keys.
    auto [e0k0, e0k1] = dpf::make_dpf(alpha, dpf::bit::one);
    auto [e1k0, e1k1] = dpf::make_dpf(alpha, dpf::bit::one);

    // Server j reads key (j>>e)&1 of DPF e. Pack those L bits per row into
    // one digit with pack_bit_columns; lane e is DPF e.
    std::array<std::vector<std::uint64_t>, nservers> symbol{};
    symbol[0] = dpf::pack_bit_columns(e0k0, e1k0);  // parties (0,0)
    symbol[1] = dpf::pack_bit_columns(e0k1, e1k0);  // parties (1,0)
    symbol[2] = dpf::pack_bit_columns(e0k0, e1k1);  // parties (0,1)
    symbol[3] = dpf::pack_bit_columns(e0k1, e1k1);  // parties (1,1)

    std::array<std::uint64_t, nservers> at_alpha{};
    for (std::size_t row = 0; row < nrows; ++row)
    {
        if (row == alpha)
        {
            for (int j = 0; j < nservers; ++j)
                at_alpha[static_cast<std::size_t>(j)] =
                    symbol[static_cast<std::size_t>(j)][row];
            continue;
        }
        for (int j = 1; j < nservers; ++j)
        {
            if (symbol[static_cast<std::size_t>(j)][row]
                != symbol[0][row])
            {
                std::cerr << "bitmore off-row\n";
                return 1;
            }
        }
    }

    // On the secret row the server digits are a translate of the server
    // ids: symbol(j) = symbol(0) XOR j. That is a permutation of 0 .. ell-1.
    for (int j = 0; j < nservers; ++j)
    {
        const std::uint64_t expect = at_alpha[0] ^ static_cast<std::uint64_t>(j);
        if (at_alpha[static_cast<std::size_t>(j)] != expect)
        {
            std::cerr << "bitmore secret row\n";
            return 1;
        }
    }

    // L = 1 is the 2-server member of the same family: XOR the rows each
    // party's bit selects, read off the packed digit's low bit.
    std::vector<std::uint64_t> records(nrows);
    records[alpha] = 99;
    records[7] = 3;
    auto [q0, q1] = dpf::make_dpf(alpha, dpf::bit::one);
    const auto s0 = dpf::pack_bit_columns(q0);
    const auto s1 = dpf::pack_bit_columns(q1);
    std::uint64_t a0 = 0;
    std::uint64_t a1 = 0;
    for (std::size_t i = 0; i < nrows; ++i)
    {
        if (s0[i] & 1u)
            a0 ^= records[i];
        if (s1[i] & 1u)
            a1 ^= records[i];
    }
    if ((a0 ^ a1) != records[alpha])
    {
        std::cerr << "bitmore two-server\n";
        return 1;
    }

    {
        if (int rc = dpf::app::run_measured("bitmore",
                dpf::protocol::bitmore_fan_plan(0, 4, 6), 6))
            return rc;
    }

    std::cout << (a0 ^ a1) << "\n";
    return 0;
}