/// @file party/flows_bulk.cpp
/// @brief Parametric beaver / DCF / Horner flows for coverage and benchmarks.

#include "cases.hpp"
#include "flow_util.hpp"
#include "key_io.hpp"
#include "registry.hpp"

#include <cstdint>
#include <string>
#include <vector>

#include "simde/simde/x86/avx2.h"

#include "dpf/beaver.hpp"
#include "dpf/dcf.hpp"
#include "dpf/dpf_key.hpp"
#include "dpf/eval_point.hpp"
#include "dpf/eval_unified.hpp"
#include "dpf/incremental.hpp"
#include "dpf/prg_aes.hpp"

#include <tuple>

namespace dpf
{
namespace party
{
namespace bulk
{

using util::Counter;
using util::evaluate_online;
using util::install_and_bind_u64;
using util::open_additive;
using util::require;
using util::role;
using util::share_bits;
using util::trio;
using util::u64;

// ---- parametric beaver product ----

template <u64 X, u64 Y>
int beaver_product_xy(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(x * y);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x, y}, {X, Y});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == X * Y, "product_xy");
    return 0;
}

// ---- parametric dot ----

template <std::size_t N>
int beaver_dot_n(role self, trio & net)
{
    beavers::session<u64> s;
    std::vector<beavers::session<u64>::wire> xs, ys;
    std::vector<u64> secrets;
    u64 expect = 0;
    for (std::size_t i = 0; i < N; ++i)
    {
        xs.push_back(s.input());
        ys.push_back(s.input());
        u64 a = static_cast<u64>(i + 1);
        u64 b = static_cast<u64>(N - i);
        secrets.push_back(a);
        expect += a * b;
    }
    for (std::size_t i = 0; i < N; ++i)
        secrets.push_back(static_cast<u64>(N - i));
    auto z = s.dot(xs, ys);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        require(s.round_of(z) == 1, "dot round");
        net::deal_session(net, s);
        return 0;
    }
    std::vector<beavers::session<u64>::wire> wires = xs;
    wires.insert(wires.end(), ys.begin(), ys.end());
    install_and_bind_u64(s, net, self, wires, secrets);
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == expect, "dot_n");
    return 0;
}

// ---- parametric scale ----

template <std::size_t N>
int beaver_scale_n(role self, trio & net)
{
    beavers::session<u64> s;
    auto sc = s.input();
    std::vector<beavers::session<u64>::wire> lanes;
    std::vector<u64> secrets{3u};
    for (std::size_t i = 0; i < N; ++i)
    {
        lanes.push_back(s.input());
        secrets.push_back(static_cast<u64>(i + 1));
    }
    auto out = s.scale(sc, lanes);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    std::vector<beavers::session<u64>::wire> wires{sc};
    wires.insert(wires.end(), lanes.begin(), lanes.end());
    install_and_bind_u64(s, net, self, wires, secrets);
    evaluate_online(s, net, self);
    std::vector<u64> mine;
    mine.reserve(N);
    for (std::size_t i = 0; i < N; ++i)
        mine.push_back(s.value_party(out[i]));
    auto opened = open_additive(net, self, mine);
    if (self == role::p0)
    {
        for (std::size_t i = 0; i < N; ++i)
            require(opened[i] == 3u * static_cast<u64>(i + 1), "scale_n");
    }
    return 0;
}

// ---- parametric stream ----

template <std::size_t N>
int beaver_stream_n(role self, trio & net)
{
    using block = prg::aes128::block_type;
    block seed = simde_mm_set_epi64x(0x9, static_cast<long long>(N));
    if (self == role::p2)
    {
        beavers::oracle<u64> src(seed, 8);
        std::vector<beavers::beaver2<u64>> bulk(N);
        beavers::fill_beaver2(src, 0, bulk.data(), N);
        require(bulk[0].ab.open() == bulk[0].a.open() * bulk[0].b.open(), "s0");
        require(bulk[N - 1].ab.open() == bulk[N - 1].a.open() * bulk[N - 1].b.open(),
            "sN");
        std::vector<u64> p0, p1;
        p0.reserve(N * 4);
        p1.reserve(N * 4);
        for (std::size_t i = 0; i < N; ++i)
        {
            p0.insert(p0.end(),
                {bulk[i].a.p0, bulk[i].b.p0, bulk[i].ab.p0, bulk[i].out.p0});
            p1.insert(p1.end(),
                {bulk[i].a.p1, bulk[i].b.p1, bulk[i].ab.p1, bulk[i].out.p1});
        }
        net.to(role::p0).send_vec(p0);
        net.to(role::p1).send_vec(p1);
        return 0;
    }
    auto shares = net.to(role::p2).recv_vec<u64>();
    require(shares.size() == N * 4, "stream size");
    // Online open copies 0 and N/2.
    auto open_copy = [&](std::size_t i, u64 sx, u64 sy) {
        u64 a = shares[i * 4 + 0];
        u64 b = shares[i * 4 + 1];
        u64 ab = shares[i * 4 + 2];
        auto [x0, x1] = util::split_u64(sx, static_cast<unsigned>(i * 2 + 1));
        auto [y0, y1] = util::split_u64(sy, static_cast<unsigned>(i * 2 + 2));
        u64 xs = self == role::p0 ? x0 : x1;
        u64 ys = self == role::p0 ? y0 : y1;
        role peer = self == role::p0 ? role::p1 : role::p0;
        u64 dx = (xs + a) + net.exchange_with(peer, xs + a);
        u64 dy = (ys + b) + net.exchange_with(peer, ys + b);
        u64 z = ab;
        z -= dx * b;
        z -= dy * a;
        if (self == role::p0)
            z += dx * dy;
        auto open = open_additive(net, self, z);
        if (self == role::p0)
            require(open == sx * sy, "stream online");
    };
    open_copy(0, 6u, 7u);
    open_copy(N / 2, 3u, 11u);
    return 0;
}

// ---- parametric public horner (one round) ----

template <std::size_t Degree>
int beaver_horner_d(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    beavers::session<u64>::wire z;
    if constexpr (Degree == 1)
        z = s.horner(x, {1ull, 2ull});
    else if constexpr (Degree == 2)
        z = s.horner(x, {1ull, 2ull, 3ull});
    else if constexpr (Degree == 3)
        z = s.horner(x, {1ull, 2ull, 3ull, 4ull});
    else if constexpr (Degree == 4)
        z = s.horner(x, {1ull, 2ull, 3ull, 4ull, 5ull});
    else if constexpr (Degree == 5)
        z = s.horner(x, {1ull, 2ull, 3ull, 4ull, 5ull, 6ull});
    else if constexpr (Degree == 6)
        z = s.horner(x, {1ull, 2ull, 3ull, 4ull, 5ull, 6ull, 7ull});
    else if constexpr (Degree == 7)
        z = s.horner(x, {1ull, 2ull, 3ull, 4ull, 5ull, 6ull, 7ull, 8ull});
    else
        z = s.horner(x, {1ull, 2ull, 3ull, 4ull, 5ull, 6ull, 7ull, 8ull, 9ull});

    const u64 xv = 5u;
    u64 expect = 0;
    u64 pow = 1;
    for (std::size_t i = 0; i <= Degree; ++i)
    {
        expect += static_cast<u64>(i + 1) * pow;
        pow *= xv;
    }
    if (self == role::p2)
    {
        require(s.round_of(z) == 1, "horner one round");
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x}, {xv});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == expect, "horner_d");
    return 0;
}

// ---- sign * horner one round ----

int beaver_sign_horner(role self, trio & net)
{
    beavers::session<u64> s;
    auto sgn = s.input();
    auto x = s.input();
    auto z = s.horner(sgn, x, {u64{4}, u64{5}, u64{6}});  // sgn*(4+5x+6x^2)
    if (self == role::p2)
    {
        require(s.round_of(z) == 1, "sign horner round");
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {sgn, x}, {2u, 3u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 2u * (4u + 5u * 3u + 6u * 9u), "sign_horner");
    return 0;
}

// ---- DCF kinds ----

template <int Kind>  // 0=lt 1=leq 2=gt 3=geq
int dcf_kind_point(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x40;
    const Input x0 = 0x15;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const std::uint64_t beta = 9;

    auto check = [&](const auto & key) {
        auto below = eval_point(cmp, key, static_cast<Input>(alpha - 1));
        auto on = eval_point(cmp, key, alpha);
        auto above = eval_point(cmp, key, static_cast<Input>(alpha + 1));
        std::vector<u64> mine{
            share_bits(below) & key.cmp().mask,
            share_bits(on) & key.cmp().mask,
            share_bits(above) & key.cmp().mask};
        auto opened = open_additive(net, self, mine);
        const auto ob = opened[0] & key.cmp().mask;
        const auto oo = opened[1] & key.cmp().mask;
        const auto oa = opened[2] & key.cmp().mask;
        if (self == role::p0)
        {
            if constexpr (Kind == 0)
                require(ob == beta && oo == 0 && oa == 0, "lt");
            else if constexpr (Kind == 1)
                require(ob == beta && oo == beta && oa == 0, "leq");
            else if constexpr (Kind == 2)
                require(ob == 0 && oo == 0 && oa == beta, "gt");
            else
                require(ob == 0 && oo == beta && oa == beta, "geq");
        }
    };

    if constexpr (Kind == 0)
    {
        require_opened(dist_with_cmp_key<prg::aes128, prg::aes128, true>(net, self, x0, x1, lt(beta), check, check),
            self, encoded_xor_point(x0, x1));
    }
    else if constexpr (Kind == 1)
    {
        require_opened(dist_with_cmp_key<prg::aes128, prg::aes128, true>(net, self, x0, x1, leq(beta), check, check),
            self, encoded_xor_point(x0, x1));
    }
    else if constexpr (Kind == 2)
    {
        require_opened(dist_with_cmp_key<prg::aes128, prg::aes128, true>(net, self, x0, x1, gt(beta), check, check),
            self, encoded_xor_point(x0, x1));
    }
    else
    {
        require_opened(dist_with_cmp_key<prg::aes128, prg::aes128, true>(net, self, x0, x1, geq(beta), check, check),
            self, encoded_xor_point(x0, x1));
    }
    return 0;
}

// ---- classic point eval bench ----

template <std::uint8_t Alpha, std::uint64_t Beta>
int dpf_point_ab(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha{Alpha};
    const Input x0 = static_cast<Input>(alpha ^ Input{0x5a});
    const Input x1 = static_cast<Input>(alpha ^ x0);
    auto on = [&](const auto & key) {
        auto y = *eval_point(key, alpha);
        auto open = util::open_subtractive(net, self, share_bits(y));
        if (self == role::p0)
            require(open == Beta, "point");
        auto yz = *eval_point(key, static_cast<Input>(Alpha ^ 1u));
        auto oz = util::open_subtractive(net, self, share_bits(yz));
        if (self == role::p0)
            require(oz == 0u, "point off");
    };
    require_tree_prefix(
        dist_with_point_key<prg::aes128, prg::aes128, true>(net, self, x0, x1, std::uint64_t{Beta}, on, on),
        self, verifiable_tree_prefix<prg::aes128, prg::aes128, std::uint64_t>(
            x0, x1));
    return 0;
}

// ---- multivariate poly one-round ----

int beaver_poly_cubic_cluster(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s.input();
    auto out = s(x * y * z + x * y + z);
    if (self == role::p2)
    {
        // Degree-3 monomials need two interactive rounds (see beaver_test xyz).
        require(s.round_of(out) == 2, "cluster round");
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x, y, z}, {2u, 3u, 5u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(out));
    if (self == role::p0)
        require(o == 2u * 3u * 5u + 2u * 3u + 5u, "cluster");
    return 0;
}

int beaver_like_terms(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(3 * x * y + 2 * x * y);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x, y}, {4u, 5u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 5u * 4u * 5u, "like");
    return 0;
}

#define REG(name, tags, fn, benchable) \
    register_flow(flow{#name, tags, fn, benchable})

}  // namespace bulk

void register_bulk_flows()
{
    using namespace bulk;
    REG(beaver_product_3_5, "beaver product bench", (beaver_product_xy<3, 5>), true);
    REG(beaver_product_11_13, "beaver product bench", (beaver_product_xy<11, 13>), true);
    REG(beaver_product_100_200, "beaver product bench", (beaver_product_xy<100, 200>), true);
    REG(beaver_product_max8, "beaver product bench",
        (beaver_product_xy<0xffffffffffffffffull, 2>), true);

    REG(beaver_dot_n1, "beaver dot bench", beaver_dot_n<1>, true);
    REG(beaver_dot_n2, "beaver dot bench", beaver_dot_n<2>, true);
    REG(beaver_dot_n4, "beaver dot bench", beaver_dot_n<4>, true);
    REG(beaver_dot_n8, "beaver dot bench", beaver_dot_n<8>, true);
    REG(beaver_dot_n16, "beaver dot bench", beaver_dot_n<16>, true);
    REG(beaver_dot_n32, "beaver dot bench", beaver_dot_n<32>, true);

    REG(beaver_scale_n1, "beaver scale bench", beaver_scale_n<1>, true);
    REG(beaver_scale_n4, "beaver scale bench", beaver_scale_n<4>, true);
    REG(beaver_scale_n16, "beaver scale bench", beaver_scale_n<16>, true);
    REG(beaver_scale_n64, "beaver scale bench", beaver_scale_n<64>, true);

    REG(beaver_stream_n8, "beaver stream bench", beaver_stream_n<8>, true);
    REG(beaver_stream_n32, "beaver stream bench", beaver_stream_n<32>, true);
    REG(beaver_stream_n128, "beaver stream bench", beaver_stream_n<128>, true);
    REG(beaver_stream_n512, "beaver stream bench", beaver_stream_n<512>, true);
    REG(beaver_stream_n2048, "beaver stream bench", beaver_stream_n<2048>, true);

    REG(beaver_horner_d1, "beaver horner grotto bench", beaver_horner_d<1>, true);
    REG(beaver_horner_d2, "beaver horner grotto bench", beaver_horner_d<2>, true);
    REG(beaver_horner_d3, "beaver horner grotto bench", beaver_horner_d<3>, true);
    REG(beaver_horner_d4, "beaver horner grotto bench", beaver_horner_d<4>, true);
    REG(beaver_horner_d5, "beaver horner grotto bench", beaver_horner_d<5>, true);
    REG(beaver_horner_d6, "beaver horner grotto bench", beaver_horner_d<6>, true);
    REG(beaver_horner_d7, "beaver horner grotto bench", beaver_horner_d<7>, true);
    REG(beaver_horner_d8, "beaver horner grotto bench", beaver_horner_d<8>, true);
    REG(beaver_sign_horner, "beaver horner grotto bench", beaver_sign_horner, true);

    REG(beaver_poly_cubic_cluster, "beaver poly bench", beaver_poly_cubic_cluster, true);
    REG(beaver_like_terms, "beaver poly bench", beaver_like_terms, true);

    REG(dcf_lt, "dcf grotto bench", dcf_kind_point<0>, true);
    REG(dcf_leq, "dcf grotto bench", dcf_kind_point<1>, true);
    REG(dcf_gt, "dcf grotto bench", dcf_kind_point<2>, true);
    REG(dcf_geq, "dcf grotto bench", dcf_kind_point<3>, true);

    REG(dpf_point_2a_7, "dpf geneval bench", (dpf_point_ab<0x2a, 7>), true);
    REG(dpf_point_00_1, "dpf geneval bench", (dpf_point_ab<0x00, 1>), true);
    REG(dpf_point_ff_99, "dpf geneval bench", (dpf_point_ab<0xff, 99>), true);
    REG(dpf_point_80_42, "dpf geneval bench", (dpf_point_ab<0x80, 42>), true);
}

#undef REG

}  // namespace party
}  // namespace dpf