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

#include "grotto/offset_repr.hpp"

#include "dpf/verifiable.hpp"

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

namespace
{

uint64_t fib_ref(uint64_t n)
{
    if (n == 0)
        return 0;
    uint64_t a = 0, b = 1;
    for (uint64_t i = 1; i < n; ++i)
    {
        const uint64_t c = a + b;
        a = b;
        b = c;
    }
    return b;
}

std::vector<uint64_t> matvec2(
    const std::vector<std::vector<uint64_t>> & M,
    const std::vector<uint64_t> & v)
{
    return {M[0][0] * v[0] + M[0][1] * v[1], M[1][0] * v[0] + M[1][1] * v[1]};
}

}  // namespace

TEST(OffsetRepr, FibonacciStateMatchesReference)
{
    for (uint64_t n = 0; n < 90; ++n)
    {
        const auto S = grotto::offset_repr_fibonacci_state(n);
        ASSERT_EQ(S.size(), 2u);
        EXPECT_EQ(S[1], fib_ref(n)) << "F_" << n;
        EXPECT_EQ(S[0], fib_ref(n + 1)) << "F_" << (n + 1);
    }
}

TEST(OffsetRepr, MatrixPowAdvancesFibonacci)
{
    const auto M = grotto::offset_repr_fibonacci_matrix();
    const auto S0 = grotto::offset_repr_fibonacci_state(0);
    for (std::int64_t k = 0; k < 40; ++k)
    {
        const auto Mk = grotto::offset_repr_matrix_pow(M, k);
        const auto advanced = matvec2(Mk, S0);
        const auto expect = grotto::offset_repr_fibonacci_state(static_cast<uint64_t>(k));
        EXPECT_EQ(advanced, expect) << "k=" << k;
    }
    // Negative: M^{-k} S_k = S_0.
    for (std::int64_t k = 1; k < 20; ++k)
    {
        const auto Sk = grotto::offset_repr_fibonacci_state(static_cast<uint64_t>(k));
        const auto Minv = grotto::offset_repr_matrix_pow(M, -k);
        const auto back = matvec2(Minv, Sk);
        EXPECT_EQ(back, S0) << "back from k=" << k;
    }
}

TEST(OffsetRepr, GeometricIsScalarPow)
{
    const uint64_t lambda = 3;
    const auto M = grotto::offset_repr_geometric_matrix(lambda);
    uint64_t expect = 1;
    for (std::int64_t e = 0; e < 20; ++e)
    {
        const auto P = grotto::offset_repr_matrix_pow(M, e);
        EXPECT_EQ(P[0][0], expect);
        expect *= lambda;
    }
    const auto Minv = grotto::offset_repr_matrix_pow(M, -1);
    EXPECT_EQ(Minv[0][0] * lambda, uint64_t{1});
}

TEST(OffsetRepr, Crc32JumpMatchesNaive)
{
    constexpr std::uint32_t poly = 0xEDB88320u;
    auto step1 = [](std::uint32_t s) {
        return (s >> 1) ^ (poly & (0u - (s & 1u)));
    };
    for (std::uint32_t seed : {0u, 1u, 0xFFFFFFFFu, 0x12345678u})
    {
        for (unsigned steps = 0; steps < 200; ++steps)
        {
            std::uint32_t naive = seed;
            for (unsigned i = 0; i < steps; ++i)
                naive = step1(naive);
            EXPECT_EQ(grotto::offset_repr_crc32_jump(seed, steps), naive)
                << "seed=" << seed << " steps=" << steps;
        }
    }
    // Doubling path for large jumps.
    const std::uint32_t seed = 0xA5A5A5A5u;
    std::uint32_t naive = seed;
    for (unsigned i = 0; i < 1000; ++i)
        naive = step1(naive);
    EXPECT_EQ(grotto::offset_repr_crc32_jump(seed, 1000), naive);
}

TEST(OffsetRepr, KeyedFibonacciMatchesClear)
{
    const uint8_t center = 10;
    const auto state = grotto::offset_repr_fibonacci_state(center);
    const auto M = grotto::offset_repr_fibonacci_matrix();
    const auto mat = grotto::make_offset_repr_keys<uint8_t>(center, state);
    const std::vector<uint8_t> knots{0};
    for (int eta = 0; eta < 256; eta += 17)
    {
        const auto e = static_cast<uint8_t>(eta);
        const auto s0 = grotto::offset_repr_eval<0>(mat, M, knots, e);
        const auto s1 = grotto::offset_repr_eval<1>(mat, M, knots, e);
        const auto clear = grotto::offset_repr_clear(center, state, M, knots, e);
        ASSERT_EQ(s0.size(), 2u);
        ASSERT_EQ(s1.size(), 2u);
        ASSERT_EQ(clear.size(), 2u);
        EXPECT_EQ(s0[0] + s1[0], clear[0]);
        EXPECT_EQ(s0[1] + s1[1], clear[1]);
        const uint8_t wrapped = static_cast<uint8_t>(center + e);
        const auto expect = grotto::offset_repr_fibonacci_state(wrapped);
        EXPECT_EQ(clear, expect) << "eta=" << eta;
    }
}

TEST(OffsetRepr, CarrySplitStillAdvances)
{
    const uint8_t center = 200;
    const uint8_t eta = 100;  // wraps: 44
    const auto state = grotto::offset_repr_fibonacci_state(center);
    const auto M = grotto::offset_repr_fibonacci_matrix();
    const auto mat = grotto::make_offset_repr_keys<uint8_t>(center, state);
    const std::vector<uint8_t> knots{0};
    const auto got0 = grotto::offset_repr_eval<0>(mat, M, knots, eta);
    const auto got1 = grotto::offset_repr_eval<1>(mat, M, knots, eta);
    const auto expect = grotto::offset_repr_fibonacci_state(44);
    EXPECT_EQ(got0[0] + got1[0], expect[0]);
    EXPECT_EQ(got0[1] + got1[1], expect[1]);
}

TEST(OffsetRepr, VerifiableProofs)
{
    const uint8_t center = 7;
    const auto state = grotto::offset_repr_fibonacci_state(center);
    const auto M = grotto::offset_repr_fibonacci_matrix();
    const auto mat = grotto::make_offset_repr_keys<uint8_t>(center, state, dpf::verifiable{});
    const std::vector<uint8_t> knots{0};
    const uint8_t eta = 5;
    std::vector<dpf::proof_token> a(2), b(2);
    const auto s0 = grotto::offset_repr_eval<0>(mat, M, knots, eta, a.data());
    const auto s1 = grotto::offset_repr_eval<1>(mat, M, knots, eta, b.data());
    const auto clear = grotto::offset_repr_clear(center, state, M, knots, eta);
    EXPECT_EQ(s0[0] + s1[0], clear[0]);
    EXPECT_EQ(s0[1] + s1[1], clear[1]);
    EXPECT_TRUE(dpf::verify(a[0], b[0]));
    EXPECT_TRUE(dpf::verify(a[1], b[1]));
}

TEST(OffsetRepr, RejectsBadDim)
{
    EXPECT_THROW(grotto::make_offset_repr_keys<uint8_t>(1, {}), std::invalid_argument);
    EXPECT_THROW(grotto::make_offset_repr_keys<uint8_t>(1,
            std::vector<uint64_t>(grotto::offset_repr_max_dim + 1, 0)), std::invalid_argument);
}