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

#include "grotto/nmod.hpp"

#include <cstdint>
#include <random>

using u128 = unsigned __int128;

namespace
{

grotto::nmod_result oracle(std::int64_t x_raw, unsigned x_bits, std::uint64_t recip,
    unsigned recip_bits, unsigned residue_bits)
{
    const unsigned scale = x_bits + recip_bits;
    const __int128 prod = static_cast<__int128>(x_raw) * static_cast<__int128>(recip);
    const bool neg = prod < 0;
    const auto mag = static_cast<u128>(neg ? -prod : prod);
    const u128 mask = scale >= 128 ? ~u128{0} : (u128{1} << scale) - 1;
    const u128 quot_mag = scale >= 128 ? 0 : mag >> scale;
    const u128 rem = scale >= 128 ? mag : mag & mask;
    grotto::nmod_result out;
    if (!neg)
        out.quotient = static_cast<std::int64_t>(quot_mag);
    else if (rem == 0)
        out.quotient = -static_cast<std::int64_t>(quot_mag);
    else
        out.quotient = -static_cast<std::int64_t>(quot_mag) - 1;
    if (residue_bits == 0 || rem == 0)
        return out;
    u128 field = rem;
    if (neg)
        field = (u128{1} << scale) - rem;
    if (scale >= residue_bits)
        field >>= scale - residue_bits;
    else
        field <<= residue_bits - scale;
    const u128 unit = u128{1} << residue_bits;
    out.residue = static_cast<std::int64_t>(field & (unit - 1));
    return out;
}

}  // namespace

TEST(Nmod, SplitsAnIntegerModulusOnTheFractionalBoundary)
{
    // 3.25 = 13/4, modulo 1.
    const auto split = grotto::nmod(13, 2, 1, 0, 2);
    EXPECT_EQ(split.quotient, 3);
    EXPECT_EQ(split.residue, 1);

    // -1.25 = -5/4. floor is -2 and the residue is 0.75.
    const auto neg = grotto::nmod(-5, 2, 1, 0, 2);
    EXPECT_EQ(neg.quotient, -2);
    EXPECT_EQ(neg.residue, 3);

    // Coarser residue truncates toward -infinity: floor(0.75 * 2) = 1.
    EXPECT_EQ(grotto::nmod(-5, 2, 1, 0, 1).residue, 1);
}

TEST(Nmod, ExactNegativeIntegersHaveAZeroResidue)
{
    const auto split = grotto::nmod(-8, 2, 1, 0, 2);
    EXPECT_EQ(split.quotient, -2);
    EXPECT_EQ(split.residue, 0);
    EXPECT_EQ(grotto::nmod(0, 8, 1, 0, 8).quotient, 0);
    EXPECT_EQ(grotto::nmod(0, 8, 1, 0, 8).residue, 0);
}

TEST(Nmod, FloorDoesNotRoundUpToTheNextQuotient)
{
    // 1 - 2^{-16}. A round-to-nearest reciprocal product would become 1.
    const auto split = grotto::nmod(1, 0, (std::uint64_t{1} << 16) - 1, 16, 8);
    EXPECT_EQ(split.quotient, 0);
    EXPECT_EQ(split.residue, 255);
}

TEST(Nmod, PowerOfTwoModulusIsAnExactShift)
{
    // 1.5 / 2^{-1} = 3 exactly.
    const auto half = grotto::nmod_pow2(6, 2, 1, 2);
    EXPECT_EQ(half.quotient, 3);
    EXPECT_EQ(half.residue, 0);

    // 1.5 / 2 = 0.75.
    const auto two = grotto::nmod_pow2(6, 2, -1, 2);
    EXPECT_EQ(two.quotient, 0);
    EXPECT_EQ(two.residue, 3);

    // 2^x splits at the integer.
    const auto unit = grotto::nmod_pow2(6, 2, 0, 2);
    EXPECT_EQ(unit.quotient, 1);
    EXPECT_EQ(unit.residue, 2);

    const auto via_recip = grotto::nmod(6, 2, 2, 0, 2);
    EXPECT_EQ(half.quotient, via_recip.quotient);
    EXPECT_EQ(half.residue, via_recip.residue);
}

TEST(Nmod, IntegerReciprocalAgreesWithFloorDivision)
{
    // round(2^100 / 3) == floor(2^100 / 3) for this width.
    const u128 recip = (u128{1} << 100) / 3;
    const auto split = grotto::nmod(10, 0, recip, 100, 16);
    EXPECT_EQ(split.quotient, 3);
    EXPECT_EQ(split.residue, 21845);
}

TEST(Nmod, WideQuarterTurnReciprocal)
{
    // RN(4/π · 2^80). x = 2.5 at 8 fractional bits, residue at 16 bits.
    const u128 recip = (u128{83443} << 64) | 494442167743545356ULL;
    const auto split = grotto::nmod(640, 8, recip, 80, 16);
    EXPECT_EQ(split.quotient, 3);
    EXPECT_EQ(split.residue, 11999);
}

TEST(Nmod, MatchesFloorOnRandomReciprocals)
{
    std::mt19937 rng(0x6d6f64u);
    std::uniform_int_distribution<int> values(-4000, 4000);
    std::uniform_int_distribution<int> bits(0, 20);
    std::uniform_int_distribution<unsigned> recip_dist(1, 100000);
    for (int i = 0; i < 4000; ++i)
    {
        const unsigned x_bits = static_cast<unsigned>(bits(rng));
        const unsigned recip_bits = static_cast<unsigned>(bits(rng));
        const unsigned residue_bits = static_cast<unsigned>(bits(rng) % 16);
        const std::int64_t x_raw = values(rng);
        const std::uint64_t recip = recip_dist(rng);
        const auto got = grotto::nmod(x_raw, x_bits, recip, recip_bits, residue_bits);
        const auto want = oracle(x_raw, x_bits, recip, recip_bits, residue_bits);
        EXPECT_EQ(got.quotient, want.quotient) << i;
        EXPECT_EQ(got.residue, want.residue) << i;
        if (got.quotient != want.quotient || got.residue != want.residue)
            break;
    }
}

TEST(Nmod, PowerOfTwoAgreesWithTheGeneralSplit)
{
    std::mt19937 rng(13);
    std::uniform_int_distribution<int> values(-2000, 2000);
    std::uniform_int_distribution<int> exps(-12, 12);
    for (int i = 0; i < 500; ++i)
    {
        const int exp = exps(rng);
        const unsigned residue_bits = static_cast<unsigned>(i % 10);
        const std::int64_t x_raw = values(rng);
        const auto got = grotto::nmod_pow2(x_raw, 8, exp, residue_bits);
        grotto::nmod_result want;
        if (exp >= 0)
            want = grotto::nmod(x_raw, 8, std::uint64_t{1} << exp, 0, residue_bits);
        else
            want = grotto::nmod(x_raw, 8, 1, static_cast<unsigned>(-exp), residue_bits);
        EXPECT_EQ(got.quotient, want.quotient);
        EXPECT_EQ(got.residue, want.residue);
    }
}

TEST(Nmod, RejectsAZeroReciprocalAndAHugeQuotient)
{
    EXPECT_THROW(grotto::nmod(1, 0, 0, 0, 4), std::invalid_argument);
    EXPECT_THROW(grotto::nmod(1, 0, 1, 0, 64), std::invalid_argument);
    EXPECT_THROW(grotto::nmod_pow2(1, 0, 128, 4), std::overflow_error);
    EXPECT_THROW(grotto::nmod(INT64_MAX, 0, u128{1} << 80, 0, 4), std::overflow_error);
}

TEST(Nmod, ScalePastTheProductWindow)
{
    const auto pos = grotto::nmod(1, 200, 1, 60, 4);
    EXPECT_EQ(pos.quotient, 0);
    EXPECT_EQ(pos.residue, 0);

    const auto neg = grotto::nmod(-1, 200, 1, 60, 4);
    EXPECT_EQ(neg.quotient, -1);
    EXPECT_EQ(neg.residue, 15);

    EXPECT_THROW(grotto::nmod(1, 100001u, 1, 0, 4), std::invalid_argument);
}

TEST(Nmod, MostNegativeInputModuloOne)
{
    const auto split = grotto::nmod(INT64_MIN, 0, 1, 0, 4);
    EXPECT_EQ(split.quotient, INT64_MIN);
    EXPECT_EQ(split.residue, 0);
}