/// @file grotto/offset_poly.hpp
/// @brief Generic offset evaluation of a polynomial, or of shares of its coefficients.
/// @details This is offset Horner at a runtime degree. The dealer keys one
///          incremental comparison `idcf(gt)` whose payload is the vector of
///          powers `center^m`. The seed spine is stored once. After
///          `eta = x - r` is public and `center = 2r`, the segment walk
///          returns additive shares of `center^m` on the piece that contains
///          the center. A public binomial shift by the piece's carry `kappa`
///          then dots to a share of `f` at the wrapped `x + r`.
///
///          Public coefficients need no further round. Shared coefficients are
///          shifted locally, because the binomial map is linear, and the dot
///          with the power shares is one Beaver inner product.

#ifndef LIBDPF_INCLUDE_GROTTO_OFFSET_POLY_HPP__
#define LIBDPF_INCLUDE_GROTTO_OFFSET_POLY_HPP__

#include "grotto/offset_horner.hpp"

#include "dpf/beaver.hpp"

#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <utility>
#include <vector>

namespace grotto
{

inline constexpr std::size_t offset_poly_max_degree = 16;

template <typename InputT>
struct offset_poly_keys
{
    static_assert(std::is_integral_v<InputT>, "offset poly domain must be an integer group");

    using input_type = InputT;

    std::size_t degree = 0;
    bool verifiable = false;
    /// @brief One `idcf(gt)` of `vec<uint64_t, degree + 1>`. Empty until `make_offset_poly_keys`.
    offset_horner_detail::lane_keys<InputT, false> keys{};
    offset_horner_detail::lane_keys<InputT, true> keys_v{};
    std::vector<std::array<uint64_t, 2>> wrap_share;
};
/// \complexity One `dpf::make_dpf` of `idcf(gt)` on a `degree + 1` lane vector. `degree` is rejected above `offset_poly_max_degree` (16).
/// The seed spine is one key. Value words grow with `degree`.
/// \rounds No party interaction.
/// \communication None inside this function.
/// \preprocessing One comparison key and `degree + 1` additive payload splits (`wrap_share`, two `uint64_t` each).
/// @see grotto::offset_horner_eval
/// @see grotto::offset_jet_eval

template <typename InputT>
offset_poly_keys<InputT> make_offset_poly_keys(InputT center, std::size_t degree);
/// \complexity One `dpf::make_dpf` of `idcf(gt)` on a `degree + 1` lane vector, with proof tokens. `degree` is rejected above `offset_poly_max_degree` (16).
/// \rounds No party interaction.
/// \communication None inside this function.
/// \preprocessing One verifiable comparison key and `degree + 1` additive payload splits (`wrap_share`, two `uint64_t` each).
/// @see grotto::offset_horner_eval
/// @see grotto::offset_jet_eval

template <typename InputT>
offset_poly_keys<InputT> make_offset_poly_keys(InputT center, std::size_t degree,
    dpf::verifiable);
/// \complexity `prepare` shifts and sorts the knots and inserts the same carry cuts. Then one segment walk of the `degree + 1` lane payload.
/// Refined pieces: the knot vector, plus the cuts `prepare` / `prepare_pieces` inserts (domain minimum, and the carry threshold when the input width is at most 62 bits). Call that count `P`.
/// Time is one walk, `degree ≤ 16`. Extra space `Θ(P · degree)` for the shifted rows.
/// \rounds None. `eta` is already public to the caller.
/// \communication None.
/// \preprocessing None created here. Uses the keys from `make_offset_poly_keys`.
/// @see grotto::offset_horner_eval
/// @see grotto::offset_jet_eval

template <std::size_t Party, typename InputT>
uint64_t offset_poly_eval(
    const offset_poly_keys<InputT> & mat,
    const std::vector<InputT> & knots,
    const std::vector<std::vector<uint64_t>> & coeff,
    InputT eta, dpf::proof_token * tokens = nullptr);
/// \complexity Piece preparation plus a Horner of `degree + 1` terms on the hot piece. No keys. `degree ≤ 16`.
/// @see grotto::offset_poly_eval

template <typename InputT>
uint64_t offset_poly_clear(
    InputT center,
    const std::vector<InputT> & knots,
    const std::vector<std::vector<uint64_t>> & coeff,
    InputT eta);

namespace offset_poly_detail
{

inline void check_degree(std::size_t degree)
{
    if (degree > offset_poly_max_degree)
        throw std::invalid_argument("offset poly: degree exceeds 16");
}

inline uint64_t binom_u64(unsigned n, unsigned k)
{
    if (k > n)
        return 0;
    if (k > n - k)
        k = n - k;
    unsigned __int128 c = 1;
    for (unsigned i = 1; i <= k; ++i)
        c = c * (n - k + i) / i;
    return static_cast<uint64_t>(c);
}

inline std::vector<uint64_t> binomial_shift(const std::vector<uint64_t> & coeff, uint64_t limb)
{
    const std::size_t degree = coeff.empty() ? 0 : coeff.size() - 1;
    std::vector<uint64_t> pow(degree + 1, 1);
    for (std::size_t m = 1; m <= degree; ++m)
        pow[m] = pow[m - 1] * limb;
    std::vector<uint64_t> out(degree + 1, 0);
    for (std::size_t m = 0; m < coeff.size(); ++m)
    {
        for (std::size_t k = 0; k <= m; ++k)
            out[k] += coeff[m] * binom_u64(static_cast<unsigned>(m), static_cast<unsigned>(k)) * pow[m - k];
    }
    return out;
}

inline uint64_t horner(const std::vector<uint64_t> & coeff, uint64_t point)
{
    if (coeff.empty())
        return 0;
    uint64_t acc = coeff.back();
    for (std::size_t i = coeff.size() - 1; i-- > 0; )
        acc = acc * point + coeff[i];
    return acc;
}

template <typename InputT>
struct prepared
{
    InputT knot{};
    std::vector<uint64_t> coeff;
    std::int64_t kappa = 0;
};

template <typename InputT>
std::vector<prepared<InputT>> prepare(
    const std::vector<InputT> & knots,
    const std::vector<std::vector<uint64_t>> & coeff,
    InputT eta)
{
    using namespace offset_horner_detail;
    if (knots.empty() || knots.size() != coeff.size())
        throw std::invalid_argument("offset poly: knots and coefficient rows differ");
    const std::size_t width = coeff.front().size();
    for (const auto & row : coeff)
    {
        if (row.size() != width)
            throw std::invalid_argument("offset poly: coefficient rows differ in length");
    }
    for (std::size_t i = 1; i < knots.size(); ++i)
    {
        if (!(knots[i - 1] < knots[i]))
            throw std::invalid_argument("offset poly: knots must be strictly increasing");
    }

    std::vector<prepared<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 prepared<InputT> & a, const prepared<InputT> & b) { return a.knot < b.knot; });

    const auto insert = [&](InputT point) {
        for (const auto & row : rows)
        {
            if (row.knot == point)
                return;
        }
        std::vector<InputT> cuts;
        for (const auto & row : rows)
            cuts.push_back(row.knot);
        std::size_t hot = rows.size() - 1;
        for (std::size_t i = 0; i + 1 < cuts.size(); ++i)
        {
            if (point >= cuts[i] && point < cuts[i + 1])
            {
                hot = i;
                break;
            }
        }
        prepared<InputT> extra;
        extra.knot = point;
        extra.coeff = rows[hot].coeff;
        rows.push_back(std::move(extra));
        std::sort(rows.begin(), rows.end(),
            [](const prepared<InputT> & a, const prepared<InputT> & b) { return a.knot < b.knot; });
    };

    constexpr unsigned bits = dpf::utils::bitlength_of_v<InputT>;
    if (bits <= 62)
    {
        insert(domain_min<InputT>());
        if (const auto cut = carry_threshold(eta))
            insert(*cut);
    }
    for (auto & row : rows)
        row.kappa = kappa_for(row.knot, eta);
    return rows;
}

template <typename Key, typename InputT>
std::vector<uint64_t> segments(const Key & key, const std::vector<InputT> & knots, uint64_t wrap_share,
    dpf::proof_token * pi = nullptr)
{
    const std::size_t n = knots.size();
    if (n == 1)
    {
        if (pi != nullptr)
        {
            dpf::detail::vdpf::init_proof(*pi, key);
            dpf::detail::vdpf::fold_output_binding(*pi, key);
        }
        return std::vector<uint64_t>{wrap_share & key.cmp().mask};
    }
    std::vector<uint64_t> prefix(n);
    signed_prefix_parities_into(key, knots.data(), n, prefix.data(), pi);
    return offset_horner_detail::segments_from_prefixes(prefix, wrap_share, key.cmp().mask);
}

}  // namespace offset_poly_detail
/// \complexity One `dpf::make_dpf` of `idcf(gt)` on a `degree + 1` lane vector. `degree` is rejected above `offset_poly_max_degree` (16).
/// The seed spine is one key. Value words grow with `degree`.
/// \rounds No party interaction.
/// \communication None inside this function.
/// \preprocessing One comparison key and `degree + 1` additive payload splits (`wrap_share`, two `uint64_t` each).
/// @see grotto::offset_horner_eval
/// @see grotto::offset_jet_eval

template <typename InputT>
offset_poly_keys<InputT> make_offset_poly_keys(InputT center, std::size_t degree)
{
    using namespace offset_horner_detail;
    offset_poly_detail::check_degree(degree);
    offset_poly_keys<InputT> mat;
    mat.degree = degree;
    mat.verifiable = false;
    std::vector<uint64_t> payload(degree + 1);
    mat.wrap_share.reserve(degree + 1);
    uint64_t pow = 1;
    const uint64_t base = lift(center);
    for (std::size_t m = 0; m <= degree; ++m)
    {
        payload[m] = pow;
        const uint64_t blind = dpf::uniform_sample<uint64_t>();
        mat.wrap_share.push_back({blind, pow - blind});
        pow *= base;
    }
    mat.keys = make_lane_keys<InputT, false>(center, payload.data(), payload.size());
    return mat;
}
/// \complexity One `dpf::make_dpf` of `idcf(gt)` on a `degree + 1` lane vector, with proof tokens. `degree` is rejected above `offset_poly_max_degree` (16).
/// \rounds No party interaction.
/// \communication None inside this function.
/// \preprocessing One verifiable comparison key and `degree + 1` additive payload splits (`wrap_share`, two `uint64_t` each).
/// @see grotto::offset_horner_eval
/// @see grotto::offset_jet_eval

template <typename InputT>
offset_poly_keys<InputT> make_offset_poly_keys(InputT center, std::size_t degree,
    dpf::verifiable)
{
    using namespace offset_horner_detail;
    offset_poly_detail::check_degree(degree);
    offset_poly_keys<InputT> mat;
    mat.degree = degree;
    mat.verifiable = true;
    std::vector<uint64_t> payload(degree + 1);
    mat.wrap_share.reserve(degree + 1);
    uint64_t pow = 1;
    const uint64_t base = lift(center);
    for (std::size_t m = 0; m <= degree; ++m)
    {
        payload[m] = pow;
        const uint64_t blind = dpf::uniform_sample<uint64_t>();
        mat.wrap_share.push_back({blind, pow - blind});
        pow *= base;
    }
    mat.keys_v = make_lane_keys<InputT, true>(center, payload.data(), payload.size());
    return mat;
}
/// \complexity Piece preparation plus a Horner of `degree + 1` terms on the hot piece. No keys. `degree ≤ 16`.
/// @see grotto::offset_poly_eval

template <typename InputT>
uint64_t offset_poly_clear(
    InputT center,
    const std::vector<InputT> & knots,
    const std::vector<std::vector<uint64_t>> & coeff,
    InputT eta)
{
    using namespace offset_poly_detail;
    const auto pieces = prepare(knots, coeff, eta);
    std::vector<InputT> cuts;
    for (const auto & piece : pieces)
        cuts.push_back(piece.knot);
    std::size_t hot = pieces.size() - 1;
    for (std::size_t i = 0; i + 1 < cuts.size(); ++i)
    {
        if (center >= cuts[i] && center < cuts[i + 1])
        {
            hot = i;
            break;
        }
    }
    const auto q = binomial_shift(pieces[hot].coeff, static_cast<uint64_t>(pieces[hot].kappa));
    return horner(q, offset_horner_detail::lift(center));
}
/// \complexity `prepare` shifts and sorts the knots and inserts the same carry cuts. Then one segment walk of the `degree + 1` lane payload.
/// Refined pieces: the knot vector, plus the cuts `prepare` / `prepare_pieces` inserts (domain minimum, and the carry threshold when the input width is at most 62 bits). Call that count `P`.
/// Time is one walk, `degree ≤ 16`. Extra space `Θ(P · degree)` for the shifted rows.
/// \rounds None. `eta` is already public to the caller.
/// \communication None.
/// \preprocessing None created here. Uses the key from `make_offset_poly_keys`.
/// @see grotto::offset_horner_eval
/// @see grotto::offset_jet_eval

template <std::size_t Party, typename InputT>
uint64_t offset_poly_eval(
    const offset_poly_keys<InputT> & mat,
    const std::vector<InputT> & knots,
    const std::vector<std::vector<uint64_t>> & coeff,
    InputT eta, dpf::proof_token * tokens)
{
    static_assert(Party < 2, "offset poly party is 0 or 1");
    using namespace offset_horner_detail;
    using namespace offset_poly_detail;
    const std::size_t lanes = mat.verifiable ? mat.keys_v.index() : mat.keys.index();
    if (lanes != mat.degree + 1)
        throw std::invalid_argument("offset poly: comparison payload width differs from the degree");
    const auto pieces = prepare(knots, coeff, eta);
    std::vector<InputT> shifted;
    shifted.reserve(pieces.size());
    for (const auto & piece : pieces)
        shifted.push_back(piece.knot);

    uint64_t value = 0;
    std::vector<std::vector<uint64_t>> shifted_coeff(pieces.size());
    for (std::size_t i = 0; i < pieces.size(); ++i)
        shifted_coeff[i] = binomial_shift(pieces[i].coeff,
            static_cast<uint64_t>(pieces[i].kappa));
    dpf::proof_token * pi = (mat.verifiable && tokens != nullptr) ? &tokens[0] : nullptr;
    const auto table = mat.verifiable
        ? lane_segments<Party>(mat.keys_v, shifted, mat.wrap_share, pi)
        : lane_segments<Party>(mat.keys, shifted, mat.wrap_share, nullptr);
    if (mat.verifiable)
        replicate_proof(tokens, mat.degree + 1);
    for (std::size_t m = 0; m <= mat.degree; ++m)
        for (std::size_t i = 0; i < pieces.size(); ++i)
            value += table[i][m] * shifted_coeff[i][m];
    return value;
}

/// @brief Public carry `kappa` of each refined piece, in knot order.
/// \complexity One `prepare` (sort plus carry cuts) and a copy of each piece's `kappa`. Time dominated by the sort, `Θ(P log P)`. Extra space `Θ(P)`.
/// @see grotto::offset_poly_eval
template <typename InputT>
std::vector<std::int64_t> offset_poly_kappas(
    const std::vector<InputT> & knots,
    std::size_t degree,
    InputT eta)
{
    std::vector<std::vector<uint64_t>> dummy(knots.size(), std::vector<uint64_t>(degree + 1, 0));
    const auto pieces = offset_poly_detail::prepare(knots, dummy, eta);
    std::vector<std::int64_t> out;
    out.reserve(pieces.size());
    for (const auto & piece : pieces)
        out.push_back(piece.kappa);
    return out;
}

/// @brief One party's binomial shift of an additive coefficient share.
/// @details `eta`'s carry `kappa` is public, so each party runs the same linear map.
/// \complexity `binomial_shift` over the coefficient vector. The implementation is a `Θ(degree²)` Pascal update (`degree` is `coeff_share.size() - 1`, at most 16).
/// @see grotto::offset_poly_power_shares
/// @note `kappa` is public, so each party runs the same linear map.
inline std::vector<uint64_t> offset_poly_shift_share(
    const std::vector<uint64_t> & coeff_share, std::int64_t kappa)
{
    return offset_poly_detail::binomial_shift(coeff_share, static_cast<uint64_t>(kappa));
}

/// @brief Shares of `center^m` on every refined piece, for one party.
/// \complexity One segment walk of the `degree + 1` lane payload over the refined pieces. Same `P` and `degree ≤ 16` as `offset_poly_eval`. Extra space `Θ(P · degree)`.
/// \rounds None.
/// \communication None.
/// \preprocessing None created here.
/// @see grotto::offset_poly_eval
template <std::size_t Party, typename InputT>
std::vector<std::vector<uint64_t>> offset_poly_power_shares(
    const offset_poly_keys<InputT> & mat,
    const std::vector<InputT> & knots,
    InputT eta)
{
    static_assert(Party < 2, "offset poly party is 0 or 1");
    using namespace offset_poly_detail;
    std::vector<std::vector<uint64_t>> dummy(knots.size(), std::vector<uint64_t>(mat.degree + 1, 0));
    const auto pieces = prepare(knots, dummy, eta);
    std::vector<InputT> shifted;
    for (const auto & piece : pieces)
        shifted.push_back(piece.knot);
    const auto table = mat.verifiable
        ? offset_horner_detail::lane_segments<Party>(mat.keys_v, shifted, mat.wrap_share, nullptr)
        : offset_horner_detail::lane_segments<Party>(mat.keys, shifted, mat.wrap_share, nullptr);
    return table;
}

/// @brief Beaver dot of shifted coefficient shares with power shares.
/// @details `q` and `power` are flattened in the same order. `opened_d` and
///          `opened_e` are the reconstructed masked differences `q - a` and
///          `power - b` (each party opens its own `q_share - a_share` and
///          `power_share - b_share`; their sums are the only openings).
///          Party 0 adds the public `d·e` term. F_Poly does not open `q` or
///          `center^m`.
/// \complexity One pass over the flattened vectors: `Θ(N)` multiplies, `N = q_share.size()`. Extra space `Θ(1)`.
/// \rounds None in this function. `opened_d` and `opened_e` are arguments; the opening of those masked differences is not performed here.
/// \communication None.
/// @note Party 0 adds the public `d·e` term. This matches the comment above the declaration.
/// @see grotto::offset_poly_power_shares
inline uint64_t offset_poly_beaver_share(
    std::size_t party,
    const std::vector<uint64_t> & q_share,
    const std::vector<uint64_t> & power_share,
    const std::vector<uint64_t> & opened_d,
    const std::vector<uint64_t> & opened_e,
    const dpf::beavers::dot_beaver<uint64_t> & triple)
{
    if (party > 1)
        throw std::invalid_argument("offset poly party is 0 or 1");
    if (q_share.size() != power_share.size() || q_share.size() != opened_d.size()
        || opened_d.size() != opened_e.size() || q_share.size() != triple.x.size())
        throw std::invalid_argument("offset poly: beaver dot length differs");
    uint64_t acc = party == 0 ? triple.cross.p0 : triple.cross.p1;
    for (std::size_t i = 0; i < q_share.size(); ++i)
    {
        const uint64_t a = party == 0 ? triple.x[i].p0 : triple.x[i].p1;
        const uint64_t b = party == 0 ? triple.y[i].p0 : triple.y[i].p1;
        const uint64_t d = opened_d[i];
        const uint64_t e = opened_e[i];
        acc += e * a + d * b;
        if (party == 0)
            acc += d * e;
    }
    (void)q_share;
    (void)power_share;
    return acc;
}

/// @brief Authenticated product share: plain Beaver dot plus a dealer tag share.
/// @details Opt-in. Callers that do not need a MAC keep the overload above.
///          `tag_share` is this party's share of `(v·Δ)` for the reconstructed
///          product `v`; sample it with `mac_share_value(v0+v1, key)`.
/// @param party 0 or 1
/// @param q_share shifted coefficient share, flattened
/// @param power_share power share, same order as `q_share`
/// @param opened_d already-opened `q - a`
/// @param opened_e already-opened `power - b`
/// @param triple dealer Beaver dot triple
/// @param tag_share this party's MAC tag share
/// @return the product share and the tag
/// @see grotto::offset_poly_beaver_share
/// \complexity One call of the plain Beaver dot (`Θ(N)`) plus storing `tag_share`. Extra space `Θ(1)`.
/// \rounds None. `opened_d`, `opened_e`, and `tag_share` are arguments.
/// \communication None.
inline dpf::mac_share<uint64_t> offset_poly_beaver_share(
    std::size_t party,
    const std::vector<uint64_t> & q_share,
    const std::vector<uint64_t> & power_share,
    const std::vector<uint64_t> & opened_d,
    const std::vector<uint64_t> & opened_e,
    const dpf::beavers::dot_beaver<uint64_t> & triple,
    uint64_t tag_share)
{
    return dpf::mac_share<uint64_t>{
        offset_poly_beaver_share(party, q_share, power_share, opened_d,
            opened_e, triple),
        tag_share};
}

}  // namespace grotto

#endif  // LIBDPF_INCLUDE_GROTTO_OFFSET_POLY_HPP__