/// @file grotto/offset_jet.hpp
/// @brief Binomial-basis jet after the public offset is opened.
/// @details The dealer keys one incremental comparison whose payload is the
///          vector of \f$\binom{\mathrm{center}}{k}\f$ in \f$\mathbb{Z}/2^{64}\f$,
///          for \f$k = 0,\ldots,d\f$. The seed spine is stored once. After `eta` opens, the same knot shift and
///          carry cut as offset poly refine the pieces. A public
///          Chu–Vandermonde convolution by the piece's carry `kappa`
///          \f$\binom{c+\kappa}{k} =\sum_j\binom{c}{j}\binom{\kappa}{k-j}\f$
///          yields additive shares of the jet at the wrapped input. Public
///          dots against that jet are free:
///          - value of \f$\sum a_k\binom{x}{k}\f$;
///          - forward difference, by Pascal's
///            \f$\binom{x+1}{k}-\binom{x}{k}=\binom{x}{k-1}\f$;
///          - hockey-stick prefix \f$\sum_{i<x}f(i)\f$, which needs one extra
///            degree so \f$\binom{x}{k+1}\f$ is in the key.
///
///          Binomials modulo \f$2^{64}\f$ use a falling factorial modulo
///          \f$2^{64+v_2(k!)}\f$ (\f$v_2(16!)=15\f$), then multiply by the
///          inverse of the odd part of \f$k!\f$. Dividing by \f$k!\f$ inside
///          \f$\mathbb{Z}/2^{64}\f$ alone is not exact.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license.

#ifndef LIBDPF_INCLUDE_GROTTO_OFFSET_JET_HPP__
#define LIBDPF_INCLUDE_GROTTO_OFFSET_JET_HPP__

#include "grotto/offset_poly.hpp"

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

namespace grotto
{

inline constexpr std::size_t offset_jet_max_degree = offset_poly_max_degree;

template <typename InputT>
struct offset_jet_keys
{
    static_assert(std::is_integral_v<InputT>, "offset jet 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>`.
    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;
};

namespace offset_jet_detail
{

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

/// @brief \f$v_2(k!)\f$ for \f$k \le 16\f$.
inline unsigned val2_factorial(unsigned k) noexcept
{
    unsigned v = 0;
    for (unsigned p = 2; p <= k; p <<= 1)
        v += k / p;
    return v;
}

/// @brief Odd part of \f$k!\f$ and \f$v_2(k!)\f$.
inline uint64_t odd_factorial(unsigned k, unsigned & val2) noexcept
{
    val2 = 0;
    uint64_t odd = 1;
    for (unsigned i = 1; i <= k; ++i)
    {
        unsigned x = i;
        while ((x & 1u) == 0u)
        {
            x >>= 1u;
            ++val2;
        }
        odd *= x;
    }
    return odd;
}

/// @brief Modular inverse of an odd unit modulo \f$2^{64}\f$.
inline uint64_t inv_odd_u64(uint64_t a) noexcept
{
    uint64_t x = a;
    x *= 2 - a * x;
    x *= 2 - a * x;
    x *= 2 - a * x;
    x *= 2 - a * x;
    x *= 2 - a * x;
    return x;
}

/// @brief \f$\binom{n}{k} \bmod 2^{64}\f$ for a nonnegative upper index.
HEDLEY_CONST
HEDLEY_NO_THROW
inline uint64_t binom_u64(uint64_t n, unsigned k) noexcept
{
    if (k == 0)
        return 1;
    if (k > offset_jet_max_degree)
        return 0;
    unsigned val2 = 0;
    const uint64_t odd = odd_factorial(k, val2);
    const unsigned bits = 64u + val2;
    const unsigned __int128 mask =
        (bits >= 128) ? ~static_cast<unsigned __int128>(0)
                      : (static_cast<unsigned __int128>(1) << bits) - 1u;
    unsigned __int128 falling = 1;
    for (unsigned i = 0; i < k; ++i)
    {
        const unsigned __int128 term = static_cast<unsigned __int128>(n - i) & mask;
        falling = (falling * term) & mask;
    }
    falling = (falling * inv_odd_u64(odd)) & mask;
    falling >>= val2;
    return static_cast<uint64_t>(falling);
}

/// @brief \f$\binom{n}{k} \bmod 2^{64}\f$ for a signed upper index (carry `kappa`).
HEDLEY_CONST
HEDLEY_NO_THROW
inline uint64_t binom_i64(std::int64_t n, unsigned k) noexcept
{
    if (k == 0)
        return 1;
    if (k > offset_jet_max_degree)
        return 0;
    if (n >= 0)
        return binom_u64(static_cast<uint64_t>(n), k);
    // \f$\binom{-m}{k}=(-1)^k\binom{m+k-1}{k}\f$ with \f$m=-n>0\f$.
    const uint64_t m = static_cast<uint64_t>(-n);
    const uint64_t mag = binom_u64(m + k - 1u, k);
    return (k & 1u) ? static_cast<uint64_t>(0) - mag : mag;
}

inline std::vector<uint64_t> chu_vandermonde(
    const std::vector<uint64_t> & center_jet, std::int64_t kappa)
{
    const std::size_t degree = center_jet.empty() ? 0 : center_jet.size() - 1;
    std::vector<uint64_t> out(degree + 1, 0);
    for (std::size_t k = 0; k <= degree; ++k)
    {
        uint64_t acc = 0;
        for (std::size_t j = 0; j <= k; ++j)
            acc += center_jet[j] * binom_i64(kappa, static_cast<unsigned>(k - j));
        out[k] = acc;
    }
    return out;
}

}  // namespace offset_jet_detail

/// @brief \f$\binom{n}{k} \bmod 2^{64}\f$ for nonnegative \f$n\f$.
/// \complexity The product for `k ≤ 16` runs `k` steps, widening by `v_2(k!)` (15 when `k` is 16) and multiplying by the odd inverse. `Θ(k)`, extra space `Θ(1)`.
/// @warning Not division by `k!` inside `Z/2^64`. The falling factorial uses the extra `v_2(k!)` bits, then the inverse of the odd part of `k!`.
/// @see grotto::offset_jet_shares
/// @param n upper index
/// @param k lower index, at most 16 in the keyed jet
/// @return `binom(n, k)` modulo `2^64`
HEDLEY_CONST
HEDLEY_NO_THROW
inline uint64_t offset_jet_binom(uint64_t n, unsigned k) noexcept
{
    return offset_jet_detail::binom_u64(n, k);
}

/// @brief \f$\binom{n}{k} \bmod 2^{64}\f$ for a signed upper index.
/// \complexity The product for `k ≤ 16` runs `k` steps, widening by `v_2(k!)` (15 when `k` is 16) and multiplying by the odd inverse. `Θ(k)`, extra space `Θ(1)`.
/// @warning Not division by `k!` inside `Z/2^64`. The falling factorial uses the extra `v_2(k!)` bits, then the inverse of the odd part of `k!`.
/// @see grotto::offset_jet_shares
/// @param n upper index
/// @param k lower index, at most 16 in the keyed jet
/// @return `binom(n, k)` modulo `2^64`
HEDLEY_CONST
HEDLEY_NO_THROW
inline uint64_t offset_jet_binom(std::int64_t n, unsigned k) noexcept
{
    return offset_jet_detail::binom_i64(n, k);
}
/// \complexity One `dpf::make_dpf` of `idcf(gt)` on a `degree + 1` lane vector of binomials. `check_degree` rejects `degree > offset_jet_max_degree` (16, the offset-poly cap).
/// Each binomial is one call to `binom_i64`. 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 splits of those binomials.
/// @warning `binom_i64` is a falling factorial modulo `2^{64+v_2(k!)}`, then a multiply by the inverse of the odd part of `k!`. Dividing by `k!` in `Z/2^64` alone is not exact. `v_2(16!)` is 15.
/// @see grotto::offset_poly_eval
/// @see grotto::offset_horner_eval
/// @see [Binomial jet](@ref offset_jet)

template <typename InputT>
offset_jet_keys<InputT> make_offset_jet_keys(InputT center, std::size_t degree)
{
    using namespace offset_horner_detail;
    offset_jet_detail::check_degree(degree);
    offset_jet_keys<InputT> mat;
    mat.degree = degree;
    mat.verifiable = false;
    std::vector<uint64_t> payload(degree + 1);
    mat.wrap_share.reserve(degree + 1);
    // Signed representative: lift(center) is wrong for degree >= 2 when center < 0.
    const std::int64_t base = math_lift(center);
    for (std::size_t k = 0; k <= degree; ++k)
    {
        payload[k] = offset_jet_detail::binom_i64(base, static_cast<unsigned>(k));
        const uint64_t blind = dpf::uniform_sample<uint64_t>();
        mat.wrap_share.push_back({blind, payload[k] - blind});
    }
    mat.keys = offset_horner_detail::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 of binomials. `check_degree` rejects `degree > offset_jet_max_degree` (16, the offset-poly cap).
/// Each binomial is one call to `binom_i64`. 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 splits of those binomials.
/// @warning `binom_i64` is a falling factorial modulo `2^{64+v_2(k!)}`, then a multiply by the inverse of the odd part of `k!`. Dividing by `k!` in `Z/2^64` alone is not exact. `v_2(16!)` is 15.
/// @see grotto::offset_poly_eval
/// @see grotto::offset_horner_eval
/// @see [Binomial jet](@ref offset_jet)

template <typename InputT>
offset_jet_keys<InputT> make_offset_jet_keys(InputT center, std::size_t degree,
    dpf::verifiable)
{
    using namespace offset_horner_detail;
    offset_jet_detail::check_degree(degree);
    offset_jet_keys<InputT> mat;
    mat.degree = degree;
    mat.verifiable = true;
    std::vector<uint64_t> payload(degree + 1);
    mat.wrap_share.reserve(degree + 1);
    const std::int64_t base = math_lift(center);
    for (std::size_t k = 0; k <= degree; ++k)
    {
        payload[k] = offset_jet_detail::binom_i64(base, static_cast<unsigned>(k));
        const uint64_t blind = dpf::uniform_sample<uint64_t>();
        mat.wrap_share.push_back({blind, payload[k] - blind});
    }
    mat.keys_v = offset_horner_detail::make_lane_keys<InputT, true>(
        center, payload.data(), payload.size());
    return mat;
}

/// @brief Coefficient vector of the forward difference in the binomial basis.
/// @details \f$\Delta f(x)=\sum a_k\binom{x}{k-1}\f$ for \f$k\ge 1\f$, so the
///          difference coefficients are \f$(a_1,a_2,\ldots,a_d,0)\f$.
/// \complexity One copy of `coeff`, shifted down by one index. `Θ(degree)`.
/// @see grotto::offset_jet_dot
/// @see grotto::offset_jet_prefix_coeff
inline std::vector<uint64_t> offset_jet_difference_coeff(
    const std::vector<uint64_t> & coeff)
{
    if (coeff.empty())
        return {};
    std::vector<uint64_t> out(coeff.size(), 0);
    for (std::size_t k = 1; k < coeff.size(); ++k)
        out[k - 1] = coeff[k];
    return out;
}

/// @brief Coefficient vector of the hockey-stick prefix sum.
/// @details \f$\sum_{i<z}a_k\binom{i}{k}=a_k\binom{z}{k+1}\f$, so the prefix
///          coefficients are \f$(0,a_0,\ldots,a_{d-1})\f$. The key degree must
///          be one larger than the polynomial degree.
/// \complexity One copy of `coeff`, shifted up by one index. The result is one longer than `coeff`, so the key degree must be one larger. `Θ(degree)`.
/// @see grotto::offset_jet_dot
inline std::vector<uint64_t> offset_jet_prefix_coeff(
    const std::vector<uint64_t> & coeff)
{
    std::vector<uint64_t> out(coeff.size() + 1, 0);
    for (std::size_t k = 0; k < coeff.size(); ++k)
        out[k + 1] = coeff[k];
    return out;
}

/// @brief Public carry `kappa` of each refined piece, in knot order.
/// \complexity Forwards to `offset_poly_kappas`: one piece preparation, `Θ(P log P)`.
/// @see grotto::offset_poly_kappas
template <typename InputT>
std::vector<std::int64_t> offset_jet_kappas(
    const std::vector<InputT> & knots,
    std::size_t degree,
    InputT eta)
{
    return offset_poly_kappas(knots, degree, eta);
}

/// @brief Dot of a public coefficient vector with a jet share.
/// \complexity One multiply-add per coefficient. `Θ(degree)`, extra space `Θ(1)`.
/// @see grotto::offset_jet_shares
inline uint64_t offset_jet_dot(
    const std::vector<uint64_t> & coeff,
    const std::vector<uint64_t> & jet)
{
    if (coeff.size() != jet.size())
        throw std::invalid_argument("offset jet: coefficient and jet lengths differ");
    uint64_t acc = 0;
    for (std::size_t k = 0; k < coeff.size(); ++k)
        acc += coeff[k] * jet[k];
    return acc;
}

/// @brief Shares of the shifted binomial jet on the hot piece, for one party.
/// \complexity One segment walk of the `degree + 1` lane payload, then a Chu–Vandermonde update.
/// The update loops `m`, then pieces `i`, then `k ≥ m`, so the arithmetic after the walks is `Θ(P · degree²)`.
/// 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`. `degree ≤ 16`.
/// Extra space is the jet (`degree + 1` words) and the per-piece binomial table `Θ(P · degree)`.
/// \rounds None. `eta` is an argument.
/// \communication None.
/// \preprocessing None created here. Uses the keys from `make_offset_jet_keys`.
/// @warning The binomials in that update are `offset_jet_binom`: not division by `k!` inside `Z/2^64`.
/// @see grotto::offset_poly_eval
/// @see grotto::offset_horner_eval
/// @see [Binomial jet](@ref offset_jet)
template <std::size_t Party, typename InputT>
std::vector<uint64_t> offset_jet_shares(
    const offset_jet_keys<InputT> & mat,
    const std::vector<InputT> & knots,
    InputT eta,
    dpf::proof_token * tokens = nullptr)
{
    static_assert(Party < 2, "offset jet 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 jet: comparison payload width differs from the degree");
    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;
    shifted.reserve(pieces.size());
    for (const auto & piece : pieces)
        shifted.push_back(piece.knot);

    std::vector<uint64_t> jet(mat.degree + 1, 0);
    // Public kappa table and Chu factors are independent of the power index.
    std::vector<std::vector<uint64_t>> kappa_binom(pieces.size());
    for (std::size_t i = 0; i < pieces.size(); ++i)
    {
        kappa_binom[i].resize(mat.degree + 1);
        for (std::size_t t = 0; t <= mat.degree; ++t)
            kappa_binom[i][t] = offset_jet_detail::binom_i64(
                pieces[i].kappa, static_cast<unsigned>(t));
    }
    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)
        {
            // Chu–Vandermonde: only column m of the center jet is nonzero.
            for (std::size_t k = m; k <= mat.degree; ++k)
                jet[k] += table[i][m] * kappa_binom[i][k - m];
        }
    }
    return jet;
}

/// @brief One party's share of \f$\sum a_k\binom{x}{k}\f$ at the wrapped point.
/// \complexity One segment walk of the `degree + 1` lane payload, then a Chu–Vandermonde update.
/// The update loops `m`, then pieces `i`, then `k ≥ m`, so the arithmetic after the walks is `Θ(P · degree²)`.
/// 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`. `degree ≤ 16`.
/// Extra space is the jet (`degree + 1` words) and the per-piece binomial table `Θ(P · degree)`.
/// \rounds None. `eta` is an argument.
/// \communication None.
/// \preprocessing None created here. Uses the keys from `make_offset_jet_keys`.
/// @warning The binomials in that update are `offset_jet_binom`: not division by `k!` inside `Z/2^64`.
/// @see grotto::offset_poly_eval
/// @see grotto::offset_horner_eval
/// @see [Binomial jet](@ref offset_jet)
template <std::size_t Party, typename InputT>
uint64_t offset_jet_eval(
    const offset_jet_keys<InputT> & mat,
    const std::vector<InputT> & knots,
    const std::vector<uint64_t> & coeff,
    InputT eta,
    dpf::proof_token * tokens = nullptr)
{
    if (coeff.size() != mat.degree + 1)
        throw std::invalid_argument("offset jet: coefficient length must be degree+1");
    const auto jet = offset_jet_shares<Party>(mat, knots, eta, tokens);
    return offset_jet_dot(coeff, jet);
}

/// @brief Cleartext value of \f$\sum a_k\binom{\mathrm{center}+\kappa}{k}\f$.
/// \complexity Builds the center jet with `degree + 1` binomials and one Chu–Vandermonde (`Θ(degree²)`), after the piece sort. No keys.
/// @see grotto::offset_jet_eval
template <typename InputT>
uint64_t offset_jet_clear(
    InputT center,
    const std::vector<InputT> & knots,
    const std::vector<uint64_t> & coeff,
    InputT eta)
{
    using namespace offset_poly_detail;
    if (coeff.empty())
        return 0;
    const std::size_t degree = coeff.size() - 1;
    std::vector<std::vector<uint64_t>> dummy(knots.size(),
        std::vector<uint64_t>(degree + 1, 0));
    const auto pieces = prepare(knots, dummy, 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 std::int64_t base = offset_horner_detail::math_lift(center);
    std::vector<uint64_t> center_jet(degree + 1);
    for (std::size_t k = 0; k <= degree; ++k)
        center_jet[k] = offset_jet_detail::binom_i64(base, static_cast<unsigned>(k));
    const auto jet =
        offset_jet_detail::chu_vandermonde(center_jet, pieces[hot].kappa);
    return offset_jet_dot(coeff, jet);
}

}  // namespace grotto

#endif  // LIBDPF_INCLUDE_GROTTO_OFFSET_JET_HPP__