#include <gtest/gtest.h>

#include "grotto/ring_switch.hpp"

#include "dpf/field128.hpp"
#include "dpf/p256_scalar.hpp"
#include "dpf/verifiable.hpp"

#include <cstdint>

TEST(Residue, Zn64Arithmetic)
{
    using Z = grotto::zn64<7>;
    EXPECT_EQ((Z{3} + Z{5}).raw(), 1u);
    EXPECT_EQ((-Z{3}).raw(), 4u);
    EXPECT_EQ((Z{3} - Z{5}).raw(), 5u);
}

TEST(Residue, Zn128Arithmetic)
{
    using Z = grotto::zn128<1009, 0>;
    EXPECT_EQ((Z{1000} + Z{20}).lo(), 11u);
    EXPECT_EQ((-Z{1}).lo(), 1008u);
}

TEST(Residue, CrtFactor)
{
    using Z = grotto::zn64<15>;
    const Z share{11};
    EXPECT_EQ(grotto::ring_switch_factor<3>(share).raw(), 2u);
    EXPECT_EQ(grotto::ring_switch_factor<5>(share).raw(), 1u);
}

TEST(RingSwitch, Zn64WrapAndNoWrap)
{
    using Z = grotto::zn64<1009>;
    const std::uint8_t r = 200;
    const std::uint8_t eta_nw = 10;   // 200+10 < 256
    const std::uint8_t eta_w = 100;   // 200+100 >= 256
    const std::uint8_t x_nw = static_cast<std::uint8_t>(r + eta_nw);
    const std::uint8_t x_w = static_cast<std::uint8_t>(r + eta_w);

    auto mat = grotto::make_ring_switch_keys<Z>(r);
    const Z s0 = grotto::ring_switch_eval<0>(mat, eta_nw);
    const Z s1 = grotto::ring_switch_eval<1>(mat, eta_nw);
    EXPECT_EQ(s0 + s1, grotto::ring_switch_clear<Z>(x_nw, r, eta_nw));
    EXPECT_EQ((s0 + s1).raw(), static_cast<std::uint64_t>(x_nw) % 1009);

    const Z t0 = grotto::ring_switch_eval<0>(mat, eta_w);
    const Z t1 = grotto::ring_switch_eval<1>(mat, eta_w);
    EXPECT_EQ(t0 + t1, grotto::ring_switch_clear<Z>(x_w, r, eta_w));
    EXPECT_EQ((t0 + t1).raw(), static_cast<std::uint64_t>(x_w) % 1009);
}

TEST(RingSwitch, Zn64EtaZero)
{
    using Z = grotto::zn64<97>;
    const std::uint8_t r = 55;
    const std::uint8_t eta = 0;
    auto mat = grotto::make_ring_switch_keys<Z>(r);
    const Z got = grotto::ring_switch_eval<0>(mat, eta)
        + grotto::ring_switch_eval<1>(mat, eta);
    EXPECT_EQ(got.raw(), 55u);
}

TEST(RingSwitch, FullWidthLimb)
{
    using Z = grotto::zn64<10007>;
    const std::uint64_t r = 0xffff'ffff'ffff'ff00ull;
    const std::uint64_t eta = 0x200ull;  // wraps
    const std::uint64_t x = r + eta;    // wraps in uint64
    auto mat = grotto::make_ring_switch_keys<Z>(r);
    const Z got = grotto::ring_switch_eval<0>(mat, eta)
        + grotto::ring_switch_eval<1>(mat, eta);
    EXPECT_EQ(got, grotto::ring_switch_clear<Z>(x, r, eta));
    EXPECT_EQ(got.raw(), x % 10007);
}

TEST(RingSwitch, Field128)
{
    const std::uint16_t r = 40000;
    const std::uint16_t eta = 30000;  // wraps
    const std::uint16_t x = static_cast<std::uint16_t>(r + eta);
    auto mat = grotto::make_ring_switch_keys<dpf::field128>(r);
    const auto got = grotto::ring_switch_eval<0>(mat, eta)
        + grotto::ring_switch_eval<1>(mat, eta);
    EXPECT_EQ(got, grotto::ring_switch_clear<dpf::field128>(x, r, eta));
    EXPECT_EQ(got, dpf::field128{x});
}

TEST(RingSwitch, P256Scalar)
{
    const std::uint8_t r = 200;
    const std::uint8_t eta = 100;
    const std::uint8_t x = static_cast<std::uint8_t>(r + eta);
    auto mat = grotto::make_ring_switch_keys<dpf::p256_scalar>(r);
    const auto got = grotto::ring_switch_eval<0>(mat, eta)
        + grotto::ring_switch_eval<1>(mat, eta);
    EXPECT_EQ(got, grotto::ring_switch_clear<dpf::p256_scalar>(x, r, eta));
    EXPECT_EQ(got, dpf::p256_scalar{x});
}

TEST(RingSwitch, Zn128)
{
    using Z = grotto::zn128<0x9a57'0000'0000'0001ull, 0>;
    const std::uint8_t r = 10;
    const std::uint8_t eta = 20;
    const std::uint8_t x = 30;
    auto mat = grotto::make_ring_switch_keys<Z>(r);
    const Z got = grotto::ring_switch_eval<0>(mat, eta)
        + grotto::ring_switch_eval<1>(mat, eta);
    EXPECT_EQ(got, grotto::ring_switch_clear<Z>(x, r, eta));
    EXPECT_EQ(got.lo(), 30u);
}

TEST(RingSwitch, Verifiable)
{
    using Z = grotto::zn64<1009>;
    const std::uint8_t r = 17;
    const std::uint8_t eta = 200;
    const std::uint8_t x = static_cast<std::uint8_t>(r + eta);
    auto mat = grotto::make_ring_switch_keys<Z>(r, dpf::verifiable{});
    dpf::proof_token a{}, b{};
    const Z s0 = grotto::ring_switch_eval<0>(mat, eta, &a);
    const Z s1 = grotto::ring_switch_eval<1>(mat, eta, &b);
    EXPECT_EQ(s0 + s1, grotto::ring_switch_clear<Z>(x, r, eta));
    EXPECT_TRUE(dpf::verify(a, b));
}