#include <gtest/gtest.h>

#include <cstdint>
#include <cstring>
#include <limits>
#include <sstream>

#include "simde/simde/x86/avx2.h"

#include "dpf/fp61.hpp"
#include "dpf/leaf_arithmetic.hpp"

namespace
{

using dpf::fp61;
using dpf::fp61_mod;

}  // namespace

TEST(Fp61, ReduceCanonicalizesTheModulus)
{
    EXPECT_EQ(fp61::reduce(0), 0u);
    EXPECT_EQ(fp61::reduce(1), 1u);
    EXPECT_EQ(fp61::reduce(fp61_mod), 0u);
    EXPECT_EQ(fp61::reduce(fp61_mod + 1), 1u);
    EXPECT_EQ(fp61::reduce(std::uint64_t{1} << 61), 1u);
    EXPECT_EQ(fp61::reduce(std::numeric_limits<std::uint64_t>::max()), 7u);
    EXPECT_EQ(fp61::reduce(fp61::reduce(fp61_mod + 9)), fp61::reduce(9));
}

TEST(Fp61, ArithmeticWrapsInTheField)
{
    const fp61 a{fp61_mod - 1};
    const fp61 b{3};
    EXPECT_EQ((a + b).raw(), 2u);
    EXPECT_EQ((b - a).raw(), 4u);
    EXPECT_EQ((-a).raw(), 1u);
    EXPECT_EQ((-fp61{0}).raw(), 0u);
    EXPECT_EQ((a - a).raw(), 0u);
    EXPECT_EQ((-(-a)).raw(), a.raw());
    EXPECT_EQ((a * fp61{1}).raw(), a.raw());
    EXPECT_EQ((a * fp61{0}).raw(), 0u);
    EXPECT_EQ((fp61{2} * fp61{3}).raw(), 6u);
}

TEST(Fp61, MultiplicationDistributes)
{
    const fp61 a{1000};
    const fp61 b{fp61_mod - 5};
    const fp61 c{17};
    EXPECT_EQ(((a + b) * c).raw(), (a * c + b * c).raw());
    EXPECT_EQ((a * b).raw(), (b * a).raw());
}

TEST(Fp61, StreamPrintsTheReducedValue)
{
    std::ostringstream os;
    os << fp61{fp61_mod + 4};
    EXPECT_EQ(os.str(), "4");
}

TEST(Fp61, FromSeedMersenneFoldsWideWords)
{
    std::uint64_t words[2] = {fp61_mod - 1, 1};
    const auto a = fp61::from_seed(words, sizeof(words));
    // (p-1) + 8·1  (since 2^64 ≡ 8), then reduce — not a 61-bit mask of the lo word.
    EXPECT_NE(a.raw(), 0u);
    EXPECT_EQ(a.raw(), fp61::reduce((fp61_mod - 1) + 8));
}

TEST(Fp61, LeafAddReducesNearModulus)
{
    alignas(16) std::uint64_t left[2] = {fp61_mod - 1, fp61_mod - 2};
    alignas(16) std::uint64_t right[2] = {3, 5};
    alignas(16) std::uint64_t out[2]{};
    simde__m128i a{}, b{};
    std::memcpy(&a, left, sizeof(left));
    std::memcpy(&b, right, sizeof(right));
    const auto sum = dpf::leaf_arithmetic::add_t<fp61, simde__m128i>{}(a, b);
    std::memcpy(out, &sum, sizeof(out));
    EXPECT_EQ(out[0], 2u);
    EXPECT_EQ(out[1], 3u);
}