/// @file grotto/offset_horner.hpp
/// @brief Noninteractive cubic evaluation after the public offset is opened.
/// @details The dealer keys one comparison at `center` per power
///          `1, center, center^2, center^3` in Z/2^64. After the parties open
///          `eta`, each party shifts the knots by `eta` and sorts them. The
///          sign-respecting segment walk then returns additive shares of
///          `center^m` on the piece that contains the wrapped sum
///          `center + eta`, and shares of 0 on the other pieces. A public
///          binomial combination of those shares is a share of the coefficients
///          of that piece as a polynomial in `eta`. Horner at the public `eta`
///          needs no further round.
///
///          The opened value is that polynomial at `lift(center) + lift(eta)`
///          in Z/2^64. `lift` sign-extends a signed domain element and
///          zero-extends an unsigned one. This equals the polynomial at the
///          wrapped group element only when the domain addition does not
///          overflow. Piece selection still follows the wrapped element.
///
///          `offset_horner_at_x_plus_r` is the wiring from the reconstruction
///          the parties already do: `eta = x - r` and `center = 2r`.

#ifndef LIBDPF_INCLUDE_GROTTO_OFFSET_HORNER_HPP__
#define LIBDPF_INCLUDE_GROTTO_OFFSET_HORNER_HPP__

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

#include "dpf.hpp"
#include "grotto/prefix_parity.hpp"

namespace grotto
{

inline constexpr std::size_t offset_horner_max_degree = 3;

template <typename T>
T offset_horner_group_add(T a, T b) noexcept
{
    using u = std::make_unsigned_t<T>;
    return static_cast<T>(static_cast<u>(static_cast<u>(a) + static_cast<u>(b)));
}

template <typename T>
T offset_horner_group_sub(T a, T b) noexcept
{
    using u = std::make_unsigned_t<T>;
    return static_cast<T>(static_cast<u>(static_cast<u>(a) - static_cast<u>(b)));
}

/// `eta = x - r` and `center = 2r`, both in the input group.
template <typename T>
struct offset_horner_x_plus_r
{
    T eta{};
    T center{};
};

template <typename T>
offset_horner_x_plus_r<T> offset_horner_at_x_plus_r(T x, T r) noexcept
{
    return offset_horner_x_plus_r<T>{
        offset_horner_group_sub(x, r),
        offset_horner_group_add(r, r)};
}

template <typename InputT, std::size_t Degree>
struct offset_horner_keys;

namespace offset_horner_detail
{

inline constexpr uint64_t binom[4][4] = {
    {1, 0, 0, 0},
    {1, 1, 0, 0},
    {1, 2, 1, 0},
    {1, 3, 3, 1},
};

template <typename T>
uint64_t lift(T v) noexcept
{
    if constexpr (std::is_signed_v<T>)
        return static_cast<uint64_t>(static_cast<std::int64_t>(v));
    else
        return static_cast<uint64_t>(v);
}

template <std::size_t Degree>
uint64_t horner_at(const std::array<uint64_t, Degree + 1> & coeff, uint64_t point) noexcept
{
    uint64_t acc = coeff[Degree];
    for (std::size_t k = Degree; k-- > 0; )
        acc = acc * point + coeff[k];
    return acc;
}

template <std::size_t Degree>
void fill_payloads(uint64_t base, uint64_t (&payload)[Degree + 1]) noexcept
{
    uint64_t pow = 1;
    for (std::size_t m = 0; m <= Degree; ++m)
    {
        payload[m] = pow;
        pow *= base;
    }
}

template <typename InputT, std::size_t Degree, std::size_t... M>
auto make_key_array(InputT center, const uint64_t (&payload)[Degree + 1],
    std::index_sequence<M...>)
{
    using pair = typename offset_horner_keys<InputT, Degree>::key_pair;
    return std::array<pair, Degree + 1>{
        dpf::make_dpf(center, dpf::gt(payload[M]))...};
}

template <typename InputT>
void check_knots(const std::vector<InputT> & knots, std::size_t coeff_rows)
{
    if (knots.empty() || knots.size() != coeff_rows)
        throw std::invalid_argument("offset horner: knots and coefficient rows differ");
    for (std::size_t i = 1; i < knots.size(); ++i)
    {
        if (!(knots[i - 1] < knots[i]))
            throw std::invalid_argument("offset horner: knots must be strictly increasing");
    }
}

template <std::size_t Degree, typename InputT>
struct shifted_piece
{
    InputT knot{};
    std::array<uint64_t, Degree + 1> coeff{};
};

template <std::size_t Degree, typename InputT>
std::vector<shifted_piece<Degree, InputT>> shift_and_sort(
    const std::vector<InputT> & knots,
    const std::vector<std::array<uint64_t, Degree + 1>> & coeff,
    InputT eta)
{
    std::vector<shifted_piece<Degree, InputT>> rows(knots.size());
    for (std::size_t i = 0; i < knots.size(); ++i)
    {
        rows[i].knot = offset_horner_group_sub(knots[i], eta);
        rows[i].coeff = coeff[i];
    }
    std::sort(rows.begin(), rows.end(),
        [](const shifted_piece<Degree, InputT> & a, const shifted_piece<Degree, InputT> & b) {
            return a.knot < b.knot;
        });
    return rows;
}

template <typename Key, typename InputT>
std::vector<uint64_t> segments_of(const Key & key, const std::vector<InputT> & knots,
    uint64_t wrap_share)
{
    using namespace dpf::detail::dcf_impl;
    const std::size_t n = knots.size();
    const uint64_t mask = key.cmp().mask;
    if (n == 1)
        return std::vector<uint64_t>{wrap_share & mask};

    std::vector<uint64_t> prefix(n);
    signed_prefix_parities_into(key, knots.data(), n, prefix.data());
    std::vector<uint64_t> seg(n);
    for (std::size_t i = 0; i < n; ++i)
    {
        const uint64_t nxt = prefix[(i + 1) % n];
        seg[i] = (nxt + neg_m(prefix[i], mask)) & mask;
    }
    seg[n - 1] = (seg[n - 1] + wrap_share) & mask;
    return seg;
}

template <std::size_t Degree>
void accumulate(std::array<uint64_t, Degree + 1> & out,
    const std::array<std::vector<uint64_t>, Degree + 1> & seg,
    const std::vector<std::array<uint64_t, Degree + 1>> & coeff)
{
    const std::size_t n = coeff.size();
    for (std::size_t i = 0; i < n; ++i)
    {
        for (std::size_t m = 0; m <= Degree; ++m)
        {
            for (std::size_t k = 0; k <= m; ++k)
            {
                // seg[m-k] opens to center^{m-k} on this piece.
                const uint64_t weight = seg[m - k][i];
                out[k] += weight * coeff[i][m] * binom[m][k];
            }
        }
    }
}

template <std::size_t Degree, typename InputT>
int piece_index(InputT point, const std::vector<InputT> & sorted_knots)
{
    const std::size_t n = sorted_knots.size();
    if (n <= 1)
        return 0;
    for (std::size_t i = 0; i + 1 < n; ++i)
    {
        if (point >= sorted_knots[i] && point < sorted_knots[i + 1])
            return static_cast<int>(i);
    }
    return static_cast<int>(n - 1);
}

template <std::size_t Degree>
std::array<uint64_t, Degree + 1> binomial_coefficients(
    const std::array<uint64_t, Degree + 1> & a, uint64_t center_limb)
{
    std::array<uint64_t, Degree + 1> c{};
    uint64_t center_pow[Degree + 1];
    center_pow[0] = 1;
    for (std::size_t m = 1; m <= Degree; ++m)
        center_pow[m] = center_pow[m - 1] * center_limb;
    for (std::size_t m = 0; m <= Degree; ++m)
    {
        for (std::size_t k = 0; k <= m; ++k)
            c[k] += a[m] * binom[m][k] * center_pow[m - k];
    }
    return c;
}

}  // namespace offset_horner_detail

/// Both parties' comparison keys and wrap-piece shares for one center.
template <typename InputT, std::size_t Degree = offset_horner_max_degree>
struct offset_horner_keys
{
    static_assert(Degree <= offset_horner_max_degree, "offset horner degree is at most 3");
    static_assert(std::is_integral_v<InputT>, "offset horner domain must be an integer group");

    static constexpr std::size_t degree = Degree;
    using input_type = InputT;
    using key_pair = decltype(dpf::make_dpf(std::declval<InputT>(), dpf::gt(uint64_t{0})));

    InputT center{};
    /// `keys[m]` is `gt(center^m)` keyed at `center`. `.first` is party 0.
    std::array<key_pair, Degree + 1> keys;
    /// Random additive split of `center^m`, indexed `[power][party]`.
    std::array<std::array<uint64_t, 2>, Degree + 1> wrap_share{};
};

template <typename InputT, std::size_t Degree = offset_horner_max_degree>
offset_horner_keys<InputT, Degree> make_offset_horner_keys(InputT center)
{
    using namespace offset_horner_detail;
    uint64_t payload[Degree + 1];
    fill_payloads<Degree>(lift(center), payload);

    offset_horner_keys<InputT, Degree> mat{
        center,
        make_key_array<InputT, Degree>(center, payload, std::make_index_sequence<Degree + 1>{}),
        {}};
    for (std::size_t m = 0; m <= Degree; ++m)
    {
        const uint64_t blind = dpf::uniform_sample<uint64_t>();
        mat.wrap_share[m][0] = blind;
        mat.wrap_share[m][1] = payload[m] - blind;
    }
    return mat;
}

/// Cleartext coefficients of the selected piece, shifted to `center`, in Z/2^64.
template <std::size_t Degree, typename InputT>
std::array<uint64_t, Degree + 1> offset_horner_clear_coefficients(
    InputT center,
    const std::vector<InputT> & knots,
    const std::vector<std::array<uint64_t, Degree + 1>> & coeff,
    InputT eta)
{
    using namespace offset_horner_detail;
    check_knots(knots, coeff.size());
    const auto rows = shift_and_sort<Degree>(knots, coeff, eta);
    std::vector<InputT> shifted(rows.size());
    for (std::size_t i = 0; i < rows.size(); ++i)
        shifted[i] = rows[i].knot;
    const int hot = piece_index<Degree>(center, shifted);
    return binomial_coefficients<Degree>(rows[static_cast<std::size_t>(hot)].coeff, lift(center));
}

/// Cleartext value: selected piece at `lift(center) + lift(eta)` in Z/2^64.
template <std::size_t Degree, typename InputT>
uint64_t offset_horner_clear(
    InputT center,
    const std::vector<InputT> & knots,
    const std::vector<std::array<uint64_t, Degree + 1>> & coeff,
    InputT eta)
{
    const auto c = offset_horner_clear_coefficients<Degree>(center, knots, coeff, eta);
    return offset_horner_detail::horner_at<Degree>(c, offset_horner_detail::lift(eta));
}

/// One party's coefficient shares. `Party` is 0 or 1.
template <std::size_t Party, std::size_t Degree, typename InputT>
std::array<uint64_t, Degree + 1> offset_horner_coefficient_share(
    const offset_horner_keys<InputT, Degree> & mat,
    const std::vector<InputT> & knots,
    const std::vector<std::array<uint64_t, Degree + 1>> & coeff,
    InputT eta)
{
    static_assert(Party < 2, "offset horner party is 0 or 1");
    using namespace offset_horner_detail;
    check_knots(knots, coeff.size());
    const auto rows = shift_and_sort<Degree>(knots, coeff, eta);
    std::vector<InputT> shifted(rows.size());
    std::vector<std::array<uint64_t, Degree + 1>> ordered(rows.size());
    for (std::size_t i = 0; i < rows.size(); ++i)
    {
        shifted[i] = rows[i].knot;
        ordered[i] = rows[i].coeff;
    }

    std::array<std::vector<uint64_t>, Degree + 1> seg;
    for (std::size_t m = 0; m <= Degree; ++m)
    {
        seg[m] = segments_of(std::get<Party>(mat.keys[m]), shifted,
            mat.wrap_share[m][Party]);
    }
    std::array<uint64_t, Degree + 1> out{};
    accumulate<Degree>(out, seg, ordered);
    return out;
}

/// One party's share of the cubic at `lift(center) + lift(eta)`.
template <std::size_t Party, std::size_t Degree, typename InputT>
uint64_t offset_horner_eval(
    const offset_horner_keys<InputT, Degree> & mat,
    const std::vector<InputT> & knots,
    const std::vector<std::array<uint64_t, Degree + 1>> & coeff,
    InputT eta)
{
    const auto shares = offset_horner_coefficient_share<Party, Degree>(mat, knots, coeff, eta);
    return offset_horner_detail::horner_at<Degree>(shares, offset_horner_detail::lift(eta));
}

}  // namespace grotto

#endif  // LIBDPF_INCLUDE_GROTTO_OFFSET_HORNER_HPP__