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

#include "grotto/offset_twist.hpp"

#include "dpf/verifiable.hpp"

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

namespace
{

uint64_t pow_u64(uint64_t base, uint64_t exp)
{
    uint64_t acc = 1;
    while (exp != 0)
    {
        if (exp & 1u)
            acc *= base;
        base *= base;
        exp >>= 1;
    }
    return acc;
}

uint64_t twisted_poly(uint64_t x, uint64_t lambda, const std::vector<uint64_t> & coeff)
{
    uint64_t acc = 0;
    uint64_t pow = 1;
    for (uint64_t c : coeff)
    {
        acc += c * pow;
        pow *= x;
    }
    return acc * pow_u64(lambda, x);
}

uint64_t twisted_half(uint64_t x, const std::vector<uint64_t> & coeff)
{
    uint64_t acc = 0;
    uint64_t pow = 1;
    for (uint64_t c : coeff)
    {
        acc += c * pow;
        pow *= x;
    }
    if (x >= 64)
        return 0;
    return acc >> static_cast<unsigned>(x);
}

}  // namespace

TEST(OffsetTwist, OddLambdaMatchesClear)
{
    const uint8_t center = 9;
    const uint64_t lambda = 3;
    const std::size_t degree = 2;
    const auto mat = grotto::make_offset_twist_keys<uint8_t>(center, degree, lambda);
    // f(x) = (2 + 5x + x^2) * 3^x
    const std::vector<uint64_t> coeff{2, 5, 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_twist_eval<0>(mat, knots, coeff, e);
        const uint64_t s1 = grotto::offset_twist_eval<1>(mat, knots, coeff, e);
        const uint64_t clear = grotto::offset_twist_clear(
            center, lambda, knots, coeff, e);
        EXPECT_EQ(s0 + s1, clear) << "eta=" << eta;
        const uint8_t wrapped = static_cast<uint8_t>(center + e);
        EXPECT_EQ(clear, twisted_poly(wrapped, lambda, coeff)) << "eta=" << eta;
    }
}

TEST(OffsetTwist, CarrySplitStillTwists)
{
    const uint8_t center = 200;
    const uint8_t eta = 100;  // wraps to 44
    const uint64_t lambda = 5;
    const auto mat = grotto::make_offset_twist_keys<uint8_t>(center, 1, lambda);
    const std::vector<uint64_t> coeff{1, 2};  // (1 + 2x) 5^x
    const std::vector<uint8_t> knots{0};
    const uint64_t got = grotto::offset_twist_eval<0>(mat, knots, coeff, eta)
        + grotto::offset_twist_eval<1>(mat, knots, coeff, eta);
    EXPECT_EQ(got, twisted_poly(44, lambda, coeff));
}

TEST(OffsetTwist, HalfShiftsOnShares)
{
    const uint8_t center = 4;
    const std::size_t degree = 2;
    const auto mat = grotto::make_offset_twist_keys<uint8_t>(center, degree, grotto::twist_half);
    // Dyadic path returns shares of the shifted value; untwisted sum stays shared.
    const std::vector<uint64_t> coeff{7, 3, 1};
    const std::vector<uint8_t> knots{0};
    for (int eta = 0; eta < 40; eta += 3)
    {
        const auto e = static_cast<uint8_t>(eta);
        const uint64_t s0 = grotto::offset_twist_eval<0>(mat, knots, coeff, e);
        const uint64_t s1 = grotto::offset_twist_eval<1>(mat, knots, coeff, e);
        const uint8_t wrapped = static_cast<uint8_t>(center + e);
        const uint64_t clear = grotto::offset_twist_clear(
            center, grotto::twist_half, knots, coeff, e);
        EXPECT_EQ(s0 + s1, clear) << "eta=" << eta;
        EXPECT_EQ(clear, twisted_half(wrapped, coeff)) << "eta=" << eta;
    }
}

TEST(OffsetTwist, ArithmeticoGeometricClosedForm)
{
    const uint64_t lambda = 3;
    for (uint64_t n = 1; n < 30; ++n)
    {
        uint64_t naive = 0;
        for (uint64_t k = 1; k <= n; ++k)
            naive += k * pow_u64(lambda, k);
        EXPECT_EQ(grotto::offset_twist_arithmetico_geometric(n, lambda), naive)
            << "n=" << n;
    }
}

TEST(OffsetTwist, VerifiableProofs)
{
    const uint8_t center = 6;
    const uint64_t lambda = 7;
    const std::size_t degree = 2;
    const auto mat = grotto::make_offset_twist_keys<uint8_t>(
        center, degree, lambda, dpf::verifiable{});
    const std::vector<uint64_t> coeff{1, 0, 4};
    const std::vector<uint8_t> knots{0};
    const uint8_t eta = 3;
    std::vector<dpf::proof_token> a(degree + 1), b(degree + 1);
    const uint64_t s0 = grotto::offset_twist_eval<0>(mat, knots, coeff, eta, a.data());
    const uint64_t s1 = grotto::offset_twist_eval<1>(mat, knots, coeff, eta, b.data());
    EXPECT_EQ(s0 + s1,
        grotto::offset_twist_clear(center, lambda, knots, coeff, eta));
    for (std::size_t m = 0; m <= degree; ++m)
        EXPECT_TRUE(dpf::verify(a[m], b[m]));
}

TEST(OffsetTwist, RejectsEvenLambda)
{
    EXPECT_THROW(grotto::make_offset_twist_keys<uint8_t>(1, 1, uint64_t{2}), std::invalid_argument);
}

TEST(OffsetTwist, RejectsOversizeDegree)
{
    EXPECT_THROW(grotto::make_offset_twist_keys<uint8_t>(
            1, grotto::offset_twist_max_degree + 1, uint64_t{3}), std::invalid_argument);
}