#include <gtest/gtest.h>
#include <tuple>

#include "grotto/offset_jet.hpp"

#include "dpf/verifiable.hpp"

#include <cstdint>
#include <stdexcept>
#include <vector>

namespace
{

uint64_t binom_ref(std::int64_t n, unsigned k)
{
    if (k == 0)
        return 1;
    if (n >= 0)
    {
        unsigned __int128 acc = 1;
        for (unsigned i = 1; i <= k; ++i)
            acc = acc * static_cast<unsigned __int128>(n - static_cast<std::int64_t>(k - i)) / i;
        return static_cast<uint64_t>(acc);
    }
    const uint64_t mag = binom_ref(-n + static_cast<std::int64_t>(k) - 1, k);
    return (k & 1u) ? static_cast<uint64_t>(0) - mag : mag;
}

}  // namespace

TEST(OffsetJet, BinomMatchesReference)
{
    for (uint64_t n = 0; n < 40; ++n)
        for (unsigned k = 0; k <= 16; ++k)
            EXPECT_EQ(grotto::offset_jet_binom(n, k), binom_ref(static_cast<std::int64_t>(n), k))
                << "n=" << n << " k=" << k;

    for (std::int64_t n = -20; n < 0; ++n)
        for (unsigned k = 0; k <= 10; ++k)
            EXPECT_EQ(grotto::offset_jet_binom(n, k), binom_ref(n, k))
                << "n=" << n << " k=" << k;

    // Chu–Vandermonde sanity for a large center and a public kappa.
    const uint64_t c = (uint64_t{1} << 40) + 17;
    const std::int64_t kappa = -3;
    for (unsigned k = 0; k <= 8; ++k)
    {
        uint64_t lhs = grotto::offset_jet_binom(
            static_cast<std::int64_t>(c) + kappa, k);
        uint64_t rhs = 0;
        for (unsigned j = 0; j <= k; ++j)
            rhs += grotto::offset_jet_binom(c, j)
                * grotto::offset_jet_binom(kappa, k - j);
        EXPECT_EQ(lhs, rhs) << "k=" << k;
    }
}

TEST(OffsetJet, PublicCoefficientsMatchTheWrappedPoint)
{
    const std::size_t degree = 3;
    const uint8_t center = 9;
    const auto mat = grotto::make_offset_jet_keys<uint8_t>(center, degree);
    // f(x) = 2 + 3 C(x,1) + C(x,3)
    const std::vector<uint64_t> coeff{2, 3, 0, 1};
    const std::vector<uint8_t> knots{0};
    for (int eta = 0; eta < 256; eta += 19)
    {
        const auto e = static_cast<uint8_t>(eta);
        const uint64_t s0 = grotto::offset_jet_eval<0>(mat, knots, coeff, e);
        const uint64_t s1 = grotto::offset_jet_eval<1>(mat, knots, coeff, e);
        const uint64_t clear = grotto::offset_jet_clear<uint8_t>(center, knots, coeff, e);
        EXPECT_EQ(s0 + s1, clear);
        const uint8_t wrapped = static_cast<uint8_t>(center + e);
        uint64_t expect = 0;
        for (std::size_t k = 0; k <= degree; ++k)
            expect += coeff[k] * grotto::offset_jet_binom(static_cast<uint64_t>(wrapped),
                static_cast<unsigned>(k));
        EXPECT_EQ(clear, expect);
    }
}

TEST(OffsetJet, CarrySplitStillEvaluates)
{
    const uint8_t center = 200;
    const uint8_t eta = 100;
    const auto mat = grotto::make_offset_jet_keys<uint8_t>(center, 1);
    const std::vector<uint64_t> coeff{5, 1};
    const std::vector<uint8_t> knots{0};
    const uint64_t got = grotto::offset_jet_eval<0>(mat, knots, coeff, eta)
        + grotto::offset_jet_eval<1>(mat, knots, coeff, eta);
    // wrapped = 200+100 = 44; f = 5 + C(44,1) = 5+44
    EXPECT_EQ(got, uint64_t{5} + 44);
}

TEST(OffsetJet, DifferenceAndPrefixFromOneJet)
{
    const std::size_t degree = 3;
    const uint8_t center = 12;
    const uint8_t eta = 3;
    // Need degree+1 for prefix of a degree-2 polynomial; use degree 3 key.
    const auto mat = grotto::make_offset_jet_keys<uint8_t>(center, degree);
    const std::vector<uint64_t> poly{4, 2, 1};  // 4 + 2 C(x,1) + C(x,2)
    std::vector<uint64_t> coeff = poly;
    coeff.push_back(0);  // pad to degree 3
    const std::vector<uint8_t> knots{0};

    const auto j0 = grotto::offset_jet_shares<0>(mat, knots, eta);
    const auto j1 = grotto::offset_jet_shares<1>(mat, knots, eta);
    std::vector<uint64_t> jet(degree + 1);
    for (std::size_t k = 0; k <= degree; ++k)
        jet[k] = j0[k] + j1[k];

    const uint8_t x = static_cast<uint8_t>(center + eta);
    const uint64_t value = grotto::offset_jet_dot(coeff, jet);
    uint64_t expect_v = 0;
    for (std::size_t k = 0; k < poly.size(); ++k)
        expect_v += poly[k] * grotto::offset_jet_binom(static_cast<uint64_t>(x),
            static_cast<unsigned>(k));
    EXPECT_EQ(value, expect_v);

    const auto dcoeff = grotto::offset_jet_difference_coeff(coeff);
    const uint64_t diff = grotto::offset_jet_dot(dcoeff, jet);
    const uint64_t fx1 = expect_v
        - poly[0] * 0  // recompute f(x+1) - f(x)
        + 0;
    uint64_t expect_fx1 = 0;
    for (std::size_t k = 0; k < poly.size(); ++k)
        expect_fx1 += poly[k] * grotto::offset_jet_binom(
            static_cast<uint64_t>(static_cast<uint8_t>(x + 1)),
            static_cast<unsigned>(k));
    EXPECT_EQ(diff, static_cast<uint64_t>(expect_fx1 - expect_v));
    (void)fx1;

    // Prefix needs degree+1: key degree 3, poly degree <= 2.
    const auto pcoeff = grotto::offset_jet_prefix_coeff(poly);
    ASSERT_EQ(pcoeff.size(), degree + 1u);
    const uint64_t prefix = grotto::offset_jet_dot(pcoeff, jet);
    uint64_t expect_p = 0;
    for (uint8_t i = 0; i < x; ++i)
    {
        for (std::size_t k = 0; k < poly.size(); ++k)
            expect_p += poly[k] * grotto::offset_jet_binom(static_cast<uint64_t>(i),
                static_cast<unsigned>(k));
    }
    EXPECT_EQ(prefix, expect_p);
}

TEST(OffsetJet, VerifiableProofs)
{
    const std::size_t degree = 2;
    const auto mat = grotto::make_offset_jet_keys<uint8_t>(2, degree, dpf::verifiable{});
    const std::vector<uint64_t> coeff{1, 0, 3};
    const std::vector<uint8_t> knots{0};
    const uint8_t eta = 5;
    std::vector<dpf::proof_token> a(degree + 1), b(degree + 1);
    const uint64_t s0 = grotto::offset_jet_eval<0>(mat, knots, coeff, eta, a.data());
    const uint64_t s1 = grotto::offset_jet_eval<1>(mat, knots, coeff, eta, b.data());
    EXPECT_EQ(s0 + s1, grotto::offset_jet_clear<uint8_t>(2, knots, coeff, eta));
    for (std::size_t m = 0; m <= degree; ++m)
        EXPECT_TRUE(dpf::verify(a[m], b[m]));
}

TEST(OffsetJet, RejectsOversizeDegree)
{
    EXPECT_THROW(grotto::make_offset_jet_keys<uint8_t>(1, grotto::offset_jet_max_degree + 1), std::invalid_argument);
}

TEST(OffsetJet, NegativeCenterDegreeTwoPlus)
{
    const std::size_t degree = 3;
    const int8_t center = -7;
    const auto mat = grotto::make_offset_jet_keys<int8_t>(center, degree);
    // f(x) = 1 + C(x,1) + 3 C(x,2) + C(x,3)
    const std::vector<uint64_t> coeff{1, 1, 3, 1};
    const std::vector<int8_t> knots{static_cast<int8_t>(-128)};
    for (int eta = -20; eta <= 20; eta += 5)
    {
        const auto e = static_cast<int8_t>(eta);
        const uint64_t s0 = grotto::offset_jet_eval<0>(mat, knots, coeff, e);
        const uint64_t s1 = grotto::offset_jet_eval<1>(mat, knots, coeff, e);
        const uint64_t clear = grotto::offset_jet_clear<int8_t>(center, knots, coeff, e);
        EXPECT_EQ(s0 + s1, clear) << "eta=" << eta;

        const int8_t wrapped = grotto::offset_horner_group_add(center, e);
        uint64_t expect = 0;
        for (std::size_t k = 0; k <= degree; ++k)
            expect += coeff[k] * grotto::offset_jet_binom(
                static_cast<std::int64_t>(wrapped), static_cast<unsigned>(k));
        EXPECT_EQ(clear, expect) << "eta=" << eta;
    }
}