/// @file dpf/bench_cells.hpp
/// @brief Bench cells for work that is not a DPF walk.
/// @details A secret index stays a key. These cells time what happens after
///          the parties already hold shares: a constant-round word gadget, a
///          stacked branch of a leaf netlist, a short public table, and a
///          hidden reorder of an RSS column. Each cell is a compose plan. The
///          harness drives it with `run_parties`, so the payload crosses the
///          same transport as a DPF plan (`DPF_TRANSPORT`).
///
///          Two-party cells put the garble or the table on party 0 and the
///          second evaluation on party 1, with one peer payload between them.
///          The shuffle is three ring passes. The left-out party's slot is
///          empty; the other two carry that pass's array.
#ifndef LIBDPF_INCLUDE_DPF_BENCH_CELLS_HPP__
#define LIBDPF_INCLUDE_DPF_BENCH_CELLS_HPP__

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <numeric>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "dpf/arith_garble.hpp"
#include "dpf/compose.hpp"
#include "dpf/flute.hpp"
#include "dpf/shuffle.hpp"
#include "dpf/yao.hpp"
#include "dpf/yao_stack.hpp"

namespace dpf
{
namespace bench
{

/// @brief One measured configuration. The low 16 bits ride in the node aux.
enum class kind : std::uint32_t
{
    proj5 = 1,
    proj17,
    proj64,
    mul5,
    mul7,
    mul11,
    thresh8,
    thresh16,
    chain7,
    yao_if4,
    yao_if16,
    yao_hot4,
    yao_hot8,
    flute2,
    flute4,
    flute8,
    flute4x8,
    shuf16,
    shuf64,
    shuf256
};

inline constexpr std::uint32_t aux_pass(kind k, unsigned pass) noexcept
{
    return static_cast<std::uint32_t>(k) | (static_cast<std::uint32_t>(pass) << 16);
}

inline kind kind_of(std::uint32_t aux) noexcept
{
    return static_cast<kind>(aux & 0xffffu);
}

inline unsigned pass_of(std::uint32_t aux) noexcept
{
    return aux >> 16;
}

inline bool is_shuffle(kind k) noexcept
{
    return k == kind::shuf16 || k == kind::shuf64 || k == kind::shuf256;
}

inline std::size_t shuffle_n(kind k)
{
    if (k == kind::shuf16)
        return 16;
    if (k == kind::shuf64)
        return 64;
    if (k == kind::shuf256)
        return 256;
    throw std::invalid_argument("bench: shuffle cell");
}

namespace detail
{

inline dpf::yao::netlist and_n(unsigned n)
{
    dpf::yao::netlist nl;
    std::vector<dpf::yao::bit> in;
    for (unsigned i = 0; i < n; ++i)
        in.push_back(nl.shared_in());
    auto acc = in[0];
    for (unsigned i = 1; i < n; ++i)
        acc = nl.and_(acc, in[i]);
    nl.out(acc);
    return nl;
}

inline dpf::yao::netlist xor2()
{
    dpf::yao::netlist nl;
    auto a = nl.shared_in();
    auto b = nl.shared_in();
    nl.out(nl.xor_(a, b));
    return nl;
}

inline std::uint64_t mix_bytes(const std::uint8_t * p, std::size_t n)
{
    std::uint64_t h = 14695981039346656037ull;
    for (std::size_t i = 0; i < n; ++i)
    {
        h ^= p[i];
        h *= 1099511628211ull;
    }
    return h;
}

inline void paint(std::uint8_t * out, std::size_t n, std::uint64_t mix)
{
    for (std::size_t i = 0; i < n; i += 8)
    {
        const std::uint64_t w = mix + static_cast<std::uint64_t>(i);
        const std::size_t k = std::min<std::size_t>(8, n - i);
        std::memcpy(out + i, &w, k);
    }
}

inline std::size_t proj_bytes(std::uint16_t mod)
{
    arith_garble::circuit c;
    auto x = c.input(mod);
    std::vector<std::uint16_t> id(mod);
    for (std::uint16_t i = 0; i < mod; ++i)
        id[i] = i;
    c.out(c.project(x, mod, std::move(id)));
    const std::uint16_t in = 1;
    return std::max<std::size_t>(
        32, dpf::arith_garble::eval_pair(c, &in).ciphertext_rows * 32u);
}

inline std::size_t mul_bytes(std::uint16_t p)
{
    arith_garble::circuit c;
    auto x = c.input(p);
    auto y = c.input(p);
    c.out(c.mul(x, y));
    const std::uint16_t in[2] = {1, 1};
    return std::max<std::size_t>(
        32, dpf::arith_garble::eval_pair(c, in).ciphertext_rows * 32u);
}

inline std::size_t thresh_bytes(unsigned b)
{
    arith_garble::circuit c;
    std::vector<arith_garble::wire> bits;
    const auto mod = static_cast<std::uint16_t>(b + 1);
    for (unsigned i = 0; i < b; ++i)
        bits.push_back(c.input(mod));
    c.out(c.threshold(bits, static_cast<std::uint16_t>(b)));
    std::vector<std::uint16_t> in(b, 1);
    return std::max<std::size_t>(
        32, dpf::arith_garble::eval_pair(c, in.data()).ciphertext_rows * 32u);
}

inline std::size_t chain_bytes()
{
    arith_garble::circuit c;
    auto x = c.input(7);
    auto y = c.input(7);
    auto z = c.mul(x, y);
    for (int i = 0; i < 3; ++i)
        z = c.mul(z, x);
    c.out(z);
    const std::uint16_t in[2] = {2, 3};
    return std::max<std::size_t>(
        32, dpf::arith_garble::eval_pair(c, in).ciphertext_rows * 32u);
}

inline std::uint64_t run_proj(std::uint16_t mod)
{
    arith_garble::circuit c;
    auto x = c.input(mod);
    std::vector<std::uint16_t> id(mod);
    for (std::uint16_t i = 0; i < mod; ++i)
        id[i] = static_cast<std::uint16_t>((i * 3) % mod);
    c.out(c.project(x, mod, std::move(id)));
    const std::uint16_t in = static_cast<std::uint16_t>(mod / 2);
    auto got = dpf::arith_garble::eval_pair(c, &in);
    return got.opened.empty() ? 0 : got.opened[0];
}

inline std::uint64_t run_mul(std::uint16_t p)
{
    arith_garble::circuit c;
    auto x = c.input(p);
    auto y = c.input(p);
    c.out(c.mul(x, y));
    const std::uint16_t in[2] = {static_cast<std::uint16_t>(p - 1), 2};
    auto got = dpf::arith_garble::eval_pair(c, in);
    return got.opened.empty() ? 0 : got.opened[0];
}

inline std::uint64_t run_thresh(unsigned b)
{
    arith_garble::circuit c;
    std::vector<arith_garble::wire> bits;
    const auto mod = static_cast<std::uint16_t>(b + 1);
    for (unsigned i = 0; i < b; ++i)
        bits.push_back(c.input(mod));
    c.out(c.threshold(bits, static_cast<std::uint16_t>(b / 2)));
    std::vector<std::uint16_t> in(b, 0);
    for (unsigned i = 0; i < b; i += 2)
        in[i] = 1;
    auto got = dpf::arith_garble::eval_pair(c, in.data());
    return got.opened.empty() ? 0 : got.opened[0];
}

inline std::uint64_t run_chain()
{
    arith_garble::circuit c;
    auto x = c.input(7);
    auto y = c.input(7);
    auto z = c.mul(x, y);
    for (int i = 0; i < 3; ++i)
        z = c.mul(z, x);
    c.out(z);
    const std::uint16_t in[2] = {2, 3};
    auto got = dpf::arith_garble::eval_pair(c, in);
    return got.opened.empty() ? 0 : got.opened[0];
}

inline std::size_t yao_if_bytes(unsigned heavy, unsigned light)
{
    std::vector<std::uint8_t> h(heavy, 1), l(light, 1);
    auto got = dpf::yao::eval_if(and_n(heavy), and_n(light), 0, 0, h.data(),
        h.data(), l.data(), l.data());
    return std::max<std::size_t>(
        16, (got.stack_blocks + got.extra_blocks) * 16u);
}

inline std::uint64_t run_if(unsigned heavy, unsigned light)
{
    std::vector<std::uint8_t> h0(heavy, 1), h1(heavy, 0), l0(light, 1), l1(light, 1);
    auto got = dpf::yao::eval_if(and_n(heavy), and_n(light), 1, 0, h0.data(),
        h1.data(), l0.data(), l1.data());
    return got.share0.empty()
        ? 0
        : static_cast<std::uint64_t>(got.share0[0] ^ got.share1[0]);
}

inline std::size_t yao_hot_bytes(unsigned k)
{
    std::vector<dpf::yao::netlist> br;
    std::vector<std::vector<std::uint8_t>> p0, p1;
    for (unsigned i = 0; i < k; ++i)
    {
        br.push_back(i % 2 == 0 ? and_n(2) : xor2());
        p0.push_back({1, 1});
        p1.push_back({0, 0});
    }
    auto got = dpf::yao::eval_one_hot(br, 1, 0, p0, p1);
    return std::max<std::size_t>(
        16, (got.stack_blocks + got.extra_blocks) * 16u);
}

inline std::uint64_t run_hot(unsigned k)
{
    std::vector<dpf::yao::netlist> br;
    std::vector<std::vector<std::uint8_t>> p0, p1;
    for (unsigned i = 0; i < k; ++i)
    {
        br.push_back(i % 2 == 0 ? and_n(2) : xor2());
        p0.push_back({1, static_cast<std::uint8_t>(i & 1u)});
        p1.push_back({0, 1});
    }
    auto got = dpf::yao::eval_one_hot(br, 1, 2, p0, p1);
    return got.share0.empty()
        ? 0
        : static_cast<std::uint64_t>(got.share0[0] ^ got.share1[0]);
}

inline std::uint64_t run_flute(unsigned delta, unsigned n_out)
{
    const unsigned rows = 1u << delta;
    std::vector<std::uint8_t> columns(static_cast<std::size_t>(n_out) * rows);
    for (unsigned w = 0; w < n_out; ++w)
        for (unsigned j = 0; j < rows; ++j)
            columns[static_cast<std::size_t>(w) * rows + j] =
                static_cast<std::uint8_t>(((j >> (w % delta)) ^ w) & 1u);
    std::vector<std::uint8_t> bits(delta, 0);
    bits[0] = 1;
    if (delta > 2)
        bits[2] = 1;
    auto got = dpf::flute::eval_pair(delta, n_out, columns.data(), bits.data());
    std::uint64_t mix = 0;
    for (auto b : got.opened)
        mix = (mix << 1) | b;
    return mix;
}

inline rss::seed_bundle fixed_bundle()
{
    rss::seed_bundle b{};
    auto fill = [](rss::seed_block & s, std::uint8_t tag) {
        auto * p = reinterpret_cast<std::uint8_t *>(&s);
        for (std::size_t i = 0; i < sizeof(s); ++i)
            p[i] = static_cast<std::uint8_t>(tag + i * 17u);
    };
    fill(b.k01, 1);
    fill(b.k12, 2);
    fill(b.k20, 3);
    return b;
}

inline void shuffle_outbound(unsigned me, unsigned pass, std::size_t n,
    std::uint8_t * out, std::size_t out_n)
{
    if (out_n != n * sizeof(std::uint64_t))
        throw std::logic_error("bench shuffle slot");
    const unsigned order[3] = {2u, 0u, 1u};
    const unsigned left = order[pass];
    if (me == left)
    {
        std::memset(out, 0, out_n);
        return;
    }
    const auto bundle = fixed_bundle();
    std::vector<std::uint64_t> column(n);
    std::iota(column.begin(), column.end(), 0);
    shuffle::shuffle_party_view<std::uint64_t> held[3];
    for (unsigned p = 0; p < 3; ++p)
    {
        held[p].own.assign(n, 0);
        held[p].next.assign(n, 0);
    }
    std::uint64_t rng = 0xA5A5A5A5A5A5A5A5ull;
    auto draw = [&] {
        rng = rng * 6364136223846793005ull + 1u;
        return rng;
    };
    for (std::size_t i = 0; i < n; ++i)
    {
        const auto a = draw();
        const auto b = draw();
        const auto c = column[i] - a - b;
        held[0].own[i] = a;
        held[0].next[i] = b;
        held[1].own[i] = b;
        held[1].next[i] = c;
        held[2].own[i] = c;
        held[2].next[i] = a;
    }
    for (unsigned step = 0; step <= pass; ++step)
    {
        const unsigned L = order[step];
        const unsigned u = shuffle::hidden_u_party(L);
        const unsigned side = shuffle::hidden_side_party(L);
        auto u_step = shuffle::shuffle_hidden_pass<std::uint64_t>(
            u, rss::party_seeds::from_bundle(bundle, u), held[u], 0, L, nullptr);
        auto s_step = shuffle::shuffle_hidden_pass<std::uint64_t>(side,
            rss::party_seeds::from_bundle(bundle, side), held[side], 0, L,
            &u_step.out.data);
        std::vector<std::uint64_t> side_msg = s_step.out.data;
        auto l_step = shuffle::shuffle_hidden_pass<std::uint64_t>(L,
            rss::party_seeds::from_bundle(bundle, L), held[L], 0, L, &side_msg);
        if (step == pass)
        {
            const auto & msg = (me == u) ? u_step.out.data : s_step.out.data;
            std::memcpy(out, msg.data(), out_n);
        }
        held[u] = std::move(u_step.view);
        held[side] = std::move(s_step.view);
        held[L] = std::move(l_step.view);
    }
}

inline std::uint64_t heavy(kind k)
{
    switch (k)
    {
    case kind::proj5:
        return run_proj(5);
    case kind::proj17:
        return run_proj(17);
    case kind::proj64:
        return run_proj(64);
    case kind::mul5:
        return run_mul(5);
    case kind::mul7:
        return run_mul(7);
    case kind::mul11:
        return run_mul(11);
    case kind::thresh8:
        return run_thresh(8);
    case kind::thresh16:
        return run_thresh(16);
    case kind::chain7:
        return run_chain();
    case kind::yao_if4:
        return run_if(4, 2);
    case kind::yao_if16:
        return run_if(16, 8);
    case kind::yao_hot4:
        return run_hot(4);
    case kind::yao_hot8:
        return run_hot(8);
    case kind::flute2:
        return run_flute(2, 1);
    case kind::flute4:
        return run_flute(4, 1);
    case kind::flute8:
        return run_flute(8, 1);
    case kind::flute4x8:
        return run_flute(4, 8);
    default:
        return 0;
    }
}

inline std::size_t payload_of(kind k)
{
    switch (k)
    {
    case kind::proj5:
        return proj_bytes(5);
    case kind::proj17:
        return proj_bytes(17);
    case kind::proj64:
        return proj_bytes(64);
    case kind::mul5:
        return mul_bytes(5);
    case kind::mul7:
        return mul_bytes(7);
    case kind::mul11:
        return mul_bytes(11);
    case kind::thresh8:
        return thresh_bytes(8);
    case kind::thresh16:
        return thresh_bytes(16);
    case kind::chain7:
        return chain_bytes();
    case kind::yao_if4:
        return yao_if_bytes(4, 2);
    case kind::yao_if16:
        return yao_if_bytes(16, 8);
    case kind::yao_hot4:
        return yao_hot_bytes(4);
    case kind::yao_hot8:
        return yao_hot_bytes(8);
    case kind::flute2:
    case kind::flute4:
    case kind::flute8:
        return 8;
    case kind::flute4x8:
        return 16;
    case kind::shuf16:
        return 16 * sizeof(std::uint64_t);
    case kind::shuf64:
        return 64 * sizeof(std::uint64_t);
    case kind::shuf256:
        return 256 * sizeof(std::uint64_t);
    }
    throw std::invalid_argument("bench: cell");
}

inline protocol::plan peer_plan(std::size_t, kind k)
{
    protocol::composer c(0);
    auto done = c.bench_peer(static_cast<std::uint32_t>(k), payload_of(k));
    (void)done;
    return c.default_plan();
}

inline protocol::plan ring_plan(std::size_t, kind k)
{
    protocol::composer c(0);
    auto done = c.bench_ring(static_cast<std::uint32_t>(k), payload_of(k));
    (void)done;
    return c.default_plan();
}

}  // namespace detail

struct cell
{
    const char * name;
    int parties;
    kind id;
};

/// @brief The battery. Secret indexes stay on the DPF plans beside these.
inline std::vector<cell> battery()
{
    return {
        {"arith_proj_m5", 2, kind::proj5},
        {"arith_proj_m17", 2, kind::proj17},
        {"arith_proj_m64", 2, kind::proj64},
        {"arith_mul_p5", 2, kind::mul5},
        {"arith_mul_p7", 2, kind::mul7},
        {"arith_mul_p11", 2, kind::mul11},
        {"arith_thresh_b8", 2, kind::thresh8},
        {"arith_thresh_b16", 2, kind::thresh16},
        {"arith_chain_mul4", 2, kind::chain7},
        {"yao_if_4_2", 2, kind::yao_if4},
        {"yao_if_16_8", 2, kind::yao_if16},
        {"yao_onehot_k4", 2, kind::yao_hot4},
        {"yao_onehot_k8", 2, kind::yao_hot8},
        {"flute_d2", 2, kind::flute2},
        {"flute_d4", 2, kind::flute4},
        {"flute_d8", 2, kind::flute8},
        {"flute_d4_o8", 2, kind::flute4x8},
        {"shuffle_n16", 3, kind::shuf16},
        {"shuffle_n64", 3, kind::shuf64},
        {"shuffle_n256", 3, kind::shuf256},
    };
}

inline protocol::plan plan_for(std::size_t party, kind id)
{
    if (is_shuffle(id))
        return detail::ring_plan(party, id);
    return detail::peer_plan(party, id);
}

/// @brief Local half of a cell. Party 0 emits the payload. Party 1 applies it.
///        Every shuffle party emits its own ring slot.
inline void run_cell(std::uint32_t opcode, std::size_t party, std::uint32_t aux,
    std::uint8_t * out, std::size_t out_n)
{
    if (out == nullptr && out_n != 0)
        throw std::invalid_argument("bench cell buffer");
    const auto k = kind_of(aux);
    if (is_shuffle(k))
    {
        if (opcode != protocol::opcodes::bench_emit)
            return;
        detail::shuffle_outbound(static_cast<unsigned>(party), pass_of(aux),
            shuffle_n(k), out, out_n);
        return;
    }
    if (opcode == protocol::opcodes::bench_apply)
    {
        if (party != 1)
        {
            if (out_n != 0)
                out[0] = 0;
            return;
        }
        const auto mix = detail::heavy(k);
        if (out_n >= sizeof(mix))
            std::memcpy(out, &mix, sizeof(mix));
        return;
    }
    if (party != 0)
    {
        if (out_n != 0)
            std::memset(out, 0, out_n);
        return;
    }
    detail::paint(out, out_n, detail::heavy(k));
}

}  // namespace bench
}  // namespace dpf

#endif