/// @file grotto/offset_repr.hpp
/// @brief Representation shift: advance a linear-recurrence state by a public offset.
/// @details Offset Horner is the unipotent (Pascal) case of a shift-invariant
///          module. Here the dealer keys one incremental comparison whose
///          payload is the state vector
///          \f$S_c\in(\mathbb{Z}/2^{64})^d\f$ at the hidden center. The seed
///          spine is stored once. After `eta`
///          opens, each refined piece has a public carry `kappa`, and the
///          parties apply the public matrix power
///          \f$S_{c+\kappa}=M^{\kappa}S_c.\f$
///          Negative `kappa` uses \f$M^{-1}\f$ (the determinant must be a unit).
///          The wrap branch is multiplication by the public constant
///          \f$M^{\mp 2^n}\f$; when \f$M^{2^n}=I\f$ it is the identity.
///
///          Fibonacci is the companion matrix of \f$T^2-T-1\f$, via the Lucas
///          addition formula \f$F_{c+\kappa}=F_{\kappa}F_{c+1}+F_{\kappa-1}F_c\f$.
///          A \f$1\times 1\f$ matrix \f$[\lambda]\f$ is a geometric jump
///          \f$\lambda^{c+\kappa}=\lambda^{\kappa}\lambda^{c}\f$. CRC / LFSR
///          resume over \f$\mathrm{GF}(2)\f$ is the same checkpoint with XOR
///          shares; the cleartext jump helper documents that twin.
/// @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_REPR_HPP__
#define LIBDPF_INCLUDE_GROTTO_OFFSET_REPR_HPP__

#include "grotto/offset_poly.hpp"

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

namespace grotto
{

inline constexpr std::size_t offset_repr_max_dim = 8;

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

    using input_type = InputT;

    std::size_t dim = 0;
    bool verifiable = false;
    /// @brief One `idcf(gt)` of `vec<uint64_t, dim>`.
    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_repr_detail
{

inline void check_dim(std::size_t dim)
{
    if (dim == 0 || dim > offset_repr_max_dim)
        throw std::invalid_argument("offset repr: dimension must be in 1..8");
}

inline void check_square(const std::vector<std::vector<uint64_t>> & M, std::size_t dim)
{
    if (M.size() != dim)
        throw std::invalid_argument("offset repr: matrix row count differs from dim");
    for (const auto & row : M)
    {
        if (row.size() != dim)
            throw std::invalid_argument("offset repr: matrix is not square");
    }
}

inline std::vector<std::vector<uint64_t>> identity(std::size_t dim)
{
    std::vector<std::vector<uint64_t>> I(dim, std::vector<uint64_t>(dim, 0));
    for (std::size_t i = 0; i < dim; ++i)
        I[i][i] = 1;
    return I;
}

inline std::vector<std::vector<uint64_t>> matmul(
    const std::vector<std::vector<uint64_t>> & A,
    const std::vector<std::vector<uint64_t>> & B)
{
    const std::size_t n = A.size();
    std::vector<std::vector<uint64_t>> C(n, std::vector<uint64_t>(n, 0));
    for (std::size_t i = 0; i < n; ++i)
        for (std::size_t k = 0; k < n; ++k)
            for (std::size_t j = 0; j < n; ++j)
                C[i][j] += A[i][k] * B[k][j];
    return C;
}

inline std::vector<uint64_t> matvec(
    const std::vector<std::vector<uint64_t>> & M,
    const std::vector<uint64_t> & v)
{
    const std::size_t n = M.size();
    std::vector<uint64_t> out(n, 0);
    for (std::size_t i = 0; i < n; ++i)
        for (std::size_t j = 0; j < n; ++j)
            out[i] += M[i][j] * v[j];
    return out;
}

inline uint64_t inv_u64(uint64_t a)
{
    if ((a & 1u) == 0)
        throw std::invalid_argument("offset repr: matrix determinant is not a unit");
    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 Matrix inverse over \f$\mathbb{Z}/2^{64}\f$ by Gauss–Jordan.
inline std::vector<std::vector<uint64_t>> inverse(
    std::vector<std::vector<uint64_t>> A)
{
    const std::size_t n = A.size();
    auto I = identity(n);
    for (std::size_t k = 0; k < n; ++k)
    {
        std::size_t piv = k;
        while (piv < n && (A[piv][k] & 1u) == 0)
            ++piv;
        if (piv == n)
            throw std::invalid_argument("offset repr: matrix is not invertible over Z/2^64");
        std::swap(A[piv], A[k]);
        std::swap(I[piv], I[k]);
        const uint64_t inv = inv_u64(A[k][k]);
        for (std::size_t j = 0; j < n; ++j)
        {
            A[k][j] *= inv;
            I[k][j] *= inv;
        }
        for (std::size_t i = 0; i < n; ++i)
        {
            if (i == k)
                continue;
            const uint64_t f = A[i][k];
            for (std::size_t j = 0; j < n; ++j)
            {
                A[i][j] -= f * A[k][j];
                I[i][j] -= f * I[k][j];
            }
        }
    }
    return I;
}

}  // namespace offset_repr_detail

/// @brief \f$M^{e}\f$ for a signed exponent, over \f$\mathbb{Z}/2^{64}\f$.
/// \complexity If `exp < 0`, one Gauss–Jordan inverse (`inverse`, three nested loops over `n = M.size()`).
/// Then one squaring per bit of `|exp|`: the `while (e != 0)` loop, each step two `matmul`s (`Θ(n³)` each). `exp` is an `int64_t`, so at most 63 squarings.
/// Extra space `Θ(n²)` for `acc` and `base`. This function does not apply `offset_repr_max_dim`; callers that go through `check_dim` use `n ≤ 8`.
/// @see grotto::offset_repr_eval
inline std::vector<std::vector<uint64_t>> offset_repr_matrix_pow(
    std::vector<std::vector<uint64_t>> M, std::int64_t exp)
{
    const std::size_t n = M.size();
    offset_repr_detail::check_square(M, n);
    if (exp < 0)
    {
        M = offset_repr_detail::inverse(std::move(M));
        exp = -exp;
    }
    auto acc = offset_repr_detail::identity(n);
    auto base = std::move(M);
    auto e = static_cast<std::uint64_t>(exp);
    while (e != 0)
    {
        if (e & 1u)
            acc = offset_repr_detail::matmul(acc, base);
        base = offset_repr_detail::matmul(base, base);
        e >>= 1;
    }
    return acc;
}

/// @brief Companion matrix of Fibonacci: \f$S_n=(F_{n+1},F_n)\f$, \f$S_{n+1}=MS_n\f$.
/// \complexity Returns a 2×2 constant. `Θ(1)`.
/// @see grotto::offset_repr_fibonacci_state
inline std::vector<std::vector<uint64_t>> offset_repr_fibonacci_matrix()
{
    return {{1, 1}, {1, 0}};
}

/// @brief Fibonacci state at index `n`: \f$(F_{n+1}, F_n)\f$ in \f$\mathbb{Z}/2^{64}\f$.
/// \complexity One `offset_repr_matrix_pow` of the 2×2 companion at exponent `n`, then a matrix-vector product. `Θ(log n)` 2×2 multiplies.
/// @see grotto::offset_repr_matrix_pow
inline std::vector<uint64_t> offset_repr_fibonacci_state(std::uint64_t n)
{
    if (n == 0)
        return {1, 0};  // F_1=1, F_0=0
    auto M = offset_repr_matrix_pow(offset_repr_fibonacci_matrix(),
        static_cast<std::int64_t>(n));
    // S_0 = (1, 0); S_n = M^n S_0.
    return offset_repr_detail::matvec(M, {1, 0});
}

/// @brief \f$1\times 1\f$ geometric matrix \f$[\lambda]\f$.
/// \complexity Returns a 1×1 matrix. `Θ(1)`.
/// @see grotto::offset_repr_matrix_pow
inline std::vector<std::vector<uint64_t>> offset_repr_geometric_matrix(uint64_t lambda)
{
    return {{lambda}};
}

/// @brief Cleartext CRC-32 jump by `steps` bits (ISO / Ethernet polynomial).
/// @details Public twin of a GF(2) representation shift: one register word,
///          multiply by \f$x^{\mathrm{steps}}\f$ in \f$\mathrm{GF}(2)[x]/(p)\f$.
///          Keyed CRC wants XOR shares of the register; this helper is the
///          public \f$A^{\kappa}\f$ factor those shares would be multiplied by.
/// \complexity Builds a `64 × 32` jump table once (function-local static), then
///          applies one table lookup per set bit of `steps` (at most 64).
/// Extra space is that `64 × 32` table of `uint32_t`. Cleartext; no keys.
/// @see grotto::offset_repr_eval
/// @note This is the public twin. It does not take XOR shares.
inline std::uint32_t offset_repr_crc32_jump(std::uint32_t state, std::uint64_t steps)
{
    constexpr std::uint32_t poly = 0xEDB88320u;
    // jump[k][i] = image of the singleton bit `1<<i` after `2^k` steps.
    // Depends only on `poly`, so build once.
    static const auto jump = [] {
        std::array<std::array<std::uint32_t, 32>, 64> table{};
        auto step1 = [](std::uint32_t s) {
            return (s >> 1) ^ (poly & (0u - (s & 1u)));
        };
        for (unsigned i = 0; i < 32; ++i)
            table[0][i] = step1(std::uint32_t{1u} << i);
        for (unsigned k = 1; k < 64; ++k)
            for (unsigned i = 0; i < 32; ++i)
            {
                std::uint32_t s = table[k - 1][i];
                std::uint32_t out = 0;
                for (unsigned b = 0; b < 32; ++b)
                    if ((s >> b) & 1u)
                        out ^= table[k - 1][b];
                table[k][i] = out;
            }
        return table;
    }();
    for (unsigned k = 0; k < 64; ++k)
    {
        if (((steps >> k) & 1u) == 0)
            continue;
        std::uint32_t nxt = 0;
        for (unsigned b = 0; b < 32; ++b)
            if ((state >> b) & 1u)
                nxt ^= jump[k][b];
        state = nxt;
    }
    return state;
}
/// \complexity One `dpf::make_dpf` of `idcf(gt)` on a `dim`-lane vector. `check_dim` requires the length in `1 .. offset_repr_max_dim` (8).
/// The seed spine is one key. Value words grow with `dim`.
/// \rounds No party interaction.
/// \communication None inside this function.
/// \preprocessing One comparison key and `dim` additive splits of the state lanes.
/// @see grotto::offset_horner_eval
/// @see [Representation shift](@ref offset_repr)

template <typename InputT>
offset_repr_keys<InputT> make_offset_repr_keys(
    InputT center, const std::vector<uint64_t> & state)
{
    using namespace offset_horner_detail;
    offset_repr_detail::check_dim(state.size());
    offset_repr_keys<InputT> mat;
    mat.dim = state.size();
    mat.verifiable = false;
    mat.wrap_share.reserve(mat.dim);
    for (std::size_t i = 0; i < mat.dim; ++i)
    {
        const uint64_t blind = dpf::uniform_sample<uint64_t>();
        mat.wrap_share.push_back({blind, state[i] - blind});
    }
    mat.keys = offset_horner_detail::make_lane_keys<InputT, false>(
        center, state.data(), state.size());
    return mat;
}
/// \complexity One `dpf::make_dpf` of `idcf(gt)` on a `dim`-lane vector. `check_dim` requires the length in `1 .. offset_repr_max_dim` (8).
/// The seed spine is one key. Value words grow with `dim`.
/// \rounds No party interaction.
/// \communication None inside this function.
/// \preprocessing One comparison key and `dim` additive splits of the state lanes.
/// @see grotto::offset_horner_eval
/// @see [Representation shift](@ref offset_repr)

template <typename InputT>
offset_repr_keys<InputT> make_offset_repr_keys(
    InputT center, const std::vector<uint64_t> & state, dpf::verifiable)
{
    using namespace offset_horner_detail;
    offset_repr_detail::check_dim(state.size());
    offset_repr_keys<InputT> mat;
    mat.dim = state.size();
    mat.verifiable = true;
    mat.wrap_share.reserve(mat.dim);
    for (std::size_t i = 0; i < mat.dim; ++i)
    {
        const uint64_t blind = dpf::uniform_sample<uint64_t>();
        mat.wrap_share.push_back({blind, state[i] - blind});
    }
    mat.keys_v = offset_horner_detail::make_lane_keys<InputT, true>(
        center, state.data(), state.size());
    return mat;
}

/// @brief Shares of \f$S_c\f$ on every refined piece (pre-advance), for one party.
/// \complexity One segment walk of the `dim`-lane payload (`dim ≤ 8`) over the refined pieces. Extra space `Θ(P · dim)`.
/// \rounds None.
/// \communication None.
/// \preprocessing None created here.
/// @see grotto::offset_repr_eval
template <std::size_t Party, typename InputT>
std::vector<std::vector<uint64_t>> offset_repr_state_shares(
    const offset_repr_keys<InputT> & mat,
    const std::vector<InputT> & knots,
    InputT eta,
    dpf::proof_token * tokens = nullptr)
{
    static_assert(Party < 2, "offset repr 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.dim)
        throw std::invalid_argument("offset repr: comparison payload width differs from the state");
    std::vector<std::vector<uint64_t>> dummy(knots.size(),
        std::vector<uint64_t>(mat.dim, 0));
    const auto pieces = prepare(knots, dummy, eta);
    std::vector<InputT> shifted;
    for (const auto & piece : pieces)
        shifted.push_back(piece.knot);
    dpf::proof_token * pi = (mat.verifiable && tokens != nullptr) ? &tokens[0] : nullptr;
    auto out = 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.dim);
    return out;
}

/// @brief One party's share of \f$M^{\kappa}S_c\f$ at the wrapped point.
/// \complexity One segment walk of the `dim`-lane payload (`dim ≤ 8`), then `offset_repr_matrix_pow(M, kappa)` and a matrix-vector product on every refined piece.
/// `offset_repr_matrix_pow` squares a `dim × dim` matrix once per bit of `|kappa|` (`matmul` is three nested `dim` loops).
/// 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 plus `Θ(P · bitlength(kappa) · dim³)` field operations. Extra space `Θ(P · dim)` for the state table and `Θ(dim²)` for the matrix powers.
/// \rounds None. `eta` is an argument.
/// \communication None.
/// \preprocessing None created here. Uses the keys from `make_offset_repr_keys`.
/// @note A negative `kappa` inverts `M`. `inverse` rejects an even determinant (not a unit in `Z/2^64`).
/// @see grotto::offset_horner_eval
/// @see [Representation shift](@ref offset_repr)
template <std::size_t Party, typename InputT>
std::vector<uint64_t> offset_repr_eval(
    const offset_repr_keys<InputT> & mat,
    const std::vector<std::vector<uint64_t>> & M,
    const std::vector<InputT> & knots,
    InputT eta,
    dpf::proof_token * tokens = nullptr)
{
    offset_repr_detail::check_square(M, mat.dim);
    using namespace offset_poly_detail;
    std::vector<std::vector<uint64_t>> dummy(knots.size(),
        std::vector<uint64_t>(mat.dim, 0));
    const auto pieces = prepare(knots, dummy, eta);
    const auto states = offset_repr_state_shares<Party>(mat, knots, eta, tokens);
    std::vector<uint64_t> out(mat.dim, 0);
    for (std::size_t p = 0; p < pieces.size(); ++p)
    {
        const auto Mk = offset_repr_matrix_pow(M, pieces[p].kappa);
        const auto advanced = offset_repr_detail::matvec(Mk, states[p]);
        for (std::size_t i = 0; i < mat.dim; ++i)
            out[i] += advanced[i];
    }
    return out;
}

/// @brief Cleartext \f$M^{\kappa}S_c\f$ on the hot piece.
/// \complexity Same matrix powers as `offset_repr_eval`, but on the single hot piece and with the clear state instead of shares. `Θ(bitlength(kappa) · dim³)`.
/// @see grotto::offset_repr_eval
template <typename InputT>
std::vector<uint64_t> offset_repr_clear(
    InputT center,
    const std::vector<uint64_t> & state,
    const std::vector<std::vector<uint64_t>> & M,
    const std::vector<InputT> & knots,
    InputT eta)
{
    offset_repr_detail::check_dim(state.size());
    offset_repr_detail::check_square(M, state.size());
    using namespace offset_poly_detail;
    std::vector<std::vector<uint64_t>> dummy(knots.size(),
        std::vector<uint64_t>(state.size(), 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 auto Mk = offset_repr_matrix_pow(M, pieces[hot].kappa);
    return offset_repr_detail::matvec(Mk, state);
}

}  // namespace grotto

#endif  // LIBDPF_INCLUDE_GROTTO_OFFSET_REPR_HPP__