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

#include "grotto/offset_poly.hpp"

#include "dpf/verifiable.hpp"

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

namespace
{

uint64_t wrapped_pow(uint8_t center, uint8_t eta, std::size_t degree)
{
    const uint64_t point = static_cast<uint64_t>(static_cast<uint8_t>(center + eta));
    uint64_t acc = 0;
    uint64_t pow = 1;
    std::vector<uint64_t> coeff(degree + 1, 0);
    coeff[0] = 3;
    coeff[degree] = 1;
    for (uint64_t c : coeff)
    {
        acc += c * pow;
        pow *= point;
    }
    return acc;
}

}  // namespace

TEST(OffsetPoly, PublicCoefficientsMatchTheWrappedPolynomial)
{
    const std::size_t degree = 4;
    const auto mat = grotto::make_offset_poly_keys<uint8_t>(2, degree);
    std::vector<uint64_t> coeff(degree + 1, 0);
    coeff[0] = 3;
    coeff[degree] = 1;
    const std::vector<uint8_t> knots{0};
    for (int eta = 0; eta < 256; eta += 17)
    {
        const auto e = static_cast<uint8_t>(eta);
        const uint64_t s0 = grotto::offset_poly_eval<0>(mat, knots, {coeff}, e);
        const uint64_t s1 = grotto::offset_poly_eval<1>(mat, knots, {coeff}, e);
        const uint64_t clear = grotto::offset_poly_clear<uint8_t>(2, knots, {coeff}, e);
        EXPECT_EQ(s0 + s1, clear);
        EXPECT_EQ(clear, wrapped_pow(2, e, degree));
    }
}

TEST(OffsetPoly, CarrySplitStillEvaluatesTheWrappedPoint)
{
    const auto mat = grotto::make_offset_poly_keys<uint8_t>(200, 1);
    const std::vector<uint64_t> coeff{5, 1};
    const std::vector<uint8_t> knots{0};
    const uint8_t eta = 100;
    const uint64_t got = grotto::offset_poly_eval<0>(mat, knots, {coeff}, eta)
        + grotto::offset_poly_eval<1>(mat, knots, {coeff}, eta);
    EXPECT_EQ(got, uint64_t{5} + 44);
}

TEST(OffsetPoly, SharedCoefficientsUseOneBeaverDot)
{
    const std::size_t degree = 2;
    const uint8_t center = 9;
    const uint8_t eta = 4;
    const auto mat = grotto::make_offset_poly_keys<uint8_t>(center, degree);
    const std::vector<uint64_t> coeff{4, 0, 2};
    const std::vector<uint8_t> knots{0};
    const std::vector<uint64_t> share0{1, 7, 9};
    std::vector<uint64_t> share1(degree + 1);
    for (std::size_t i = 0; i < coeff.size(); ++i)
        share1[i] = coeff[i] - share0[i];

    const auto kappas = grotto::offset_poly_kappas<uint8_t>(knots, degree, eta);
    const auto p0 = grotto::offset_poly_power_shares<0>(mat, knots, eta);
    const auto p1 = grotto::offset_poly_power_shares<1>(mat, knots, eta);
    ASSERT_EQ(p0.size(), kappas.size());

    std::vector<uint64_t> q0, q1, y0, y1;
    for (std::size_t piece = 0; piece < kappas.size(); ++piece)
    {
        const auto a0 = grotto::offset_poly_shift_share(share0, kappas[piece]);
        const auto a1 = grotto::offset_poly_shift_share(share1, kappas[piece]);
        for (std::size_t m = 0; m <= degree; ++m)
        {
            q0.push_back(a0[m]);
            q1.push_back(a1[m]);
            y0.push_back(p0[piece][m]);
            y1.push_back(p1[piece][m]);
        }
    }
    auto triple = dpf::beavers::sample_dot<uint64_t>(q0.size());
    std::vector<uint64_t> opened_d(q0.size()), opened_e(y0.size());
    for (std::size_t i = 0; i < q0.size(); ++i)
    {
        // Open only masked differences q-a and power-b (F_Poly).
        opened_d[i] = (q0[i] - triple.x[i].p0) + (q1[i] - triple.x[i].p1);
        opened_e[i] = (y0[i] - triple.y[i].p0) + (y1[i] - triple.y[i].p1);
    }
    const uint64_t v0 = grotto::offset_poly_beaver_share(0, q0, y0, opened_d, opened_e, triple);
    const uint64_t v1 = grotto::offset_poly_beaver_share(1, q1, y1, opened_d, opened_e, triple);
    const uint64_t clear = grotto::offset_poly_clear<uint8_t>(center, knots, {coeff}, eta);
    EXPECT_EQ(v0 + v1, clear);
    const uint64_t point = static_cast<uint8_t>(center + eta);
    EXPECT_EQ(clear, 4 + 2 * point * point);
}

TEST(OffsetPoly, VerifiableTokensRejectATamperedProof)
{
    const std::size_t degree = 2;
    const auto mat = grotto::make_offset_poly_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_poly_eval<0>(mat, knots, {coeff}, eta, a.data());
    const uint64_t s1 = grotto::offset_poly_eval<1>(mat, knots, {coeff}, eta, b.data());
    EXPECT_EQ(s0 + s1, grotto::offset_poly_clear<uint8_t>(2, knots, {coeff}, eta));
    for (std::size_t m = 0; m <= degree; ++m)
        EXPECT_TRUE(dpf::verify(a[m], b[m])) << m;
    a[0][0] = simde_mm_xor_si128(a[0][0], simde_mm_set1_epi8(1));
    EXPECT_FALSE(dpf::verify(a[0], b[0]));
    EXPECT_TRUE(dpf::verify(a[1], b[1]));
}

TEST(OffsetPoly, HornerAndPolyAgreeOnTheSameKnots)
{
    constexpr std::size_t D = 3;
    const uint8_t center = 17;
    const uint8_t eta = 9;
    const std::vector<uint8_t> knots{0, 40, 100};
    std::vector<std::array<uint64_t, D + 1>> horner_coeff{
        {2, 0, 1, 0},
        {0, 3, 0, 1},
        {5, 1, 0, 0},
    };
    std::vector<std::vector<uint64_t>> poly_coeff(horner_coeff.size());
    for (std::size_t i = 0; i < horner_coeff.size(); ++i)
        poly_coeff[i].assign(horner_coeff[i].begin(), horner_coeff[i].end());

    auto hmat = grotto::make_offset_horner_keys<uint8_t, D>(center);
    auto pmat = grotto::make_offset_poly_keys<uint8_t>(center, D);
    const uint64_t h = grotto::offset_horner_eval<0, D>(hmat, knots, horner_coeff, eta)
        + grotto::offset_horner_eval<1, D>(hmat, knots, horner_coeff, eta);
    const uint64_t p = grotto::offset_poly_eval<0>(pmat, knots, poly_coeff, eta)
        + grotto::offset_poly_eval<1>(pmat, knots, poly_coeff, eta);
    EXPECT_EQ(h, p);
    EXPECT_EQ(h, grotto::offset_horner_clear<D>(center, knots, horner_coeff, eta));
    EXPECT_EQ(p, grotto::offset_poly_clear(center, knots, poly_coeff, eta));
}

TEST(OffsetPoly, RejectsADegreePastTheCap)
{
    EXPECT_THROW(grotto::make_offset_poly_keys<uint8_t>(1, grotto::offset_poly_max_degree + 1), std::invalid_argument);
    EXPECT_THROW(grotto::make_offset_poly_keys<uint8_t>(1, grotto::offset_poly_max_degree + 1, dpf::verifiable{}), std::invalid_argument);
}