#include <gtest/gtest.h>

#include "dpf.hpp"
#include "grotto/fixedpoint.hpp"
#include "grotto/fixedpoint_mul.hpp"
#include "grotto/principal_lut.hpp"

#include <cstdint>
#include <cstring>
#include <limits>
#include <stdexcept>
#include <tuple>
#include <vector>

namespace
{

template <typename T>
void expect_same(const T & got, const T & point)
{
    EXPECT_EQ(std::memcmp(&got, &point, sizeof(T)), 0);
}

template <typename In, typename Out>
void expect_wrapping_interval(In from, In to, In alpha, Out y, std::size_t leaves)
{
    auto [k0, k1] = dpf::make_dpf(alpha, y);
    using key_t = std::decay_t<decltype(k0)>;
    EXPECT_EQ((dpf::utils::get_nodes_in_interval<key_t>(from, to)), leaves);

    auto [buf0, it0] = dpf::eval_interval(k0, from, to);
    auto [buf1, it1] = dpf::eval_interval(k1, from, to);
    auto a = std::begin(it0);
    auto b = std::begin(it1);
    const auto a_end = std::end(it0);
    const auto b_end = std::end(it1);

    std::size_t n = 0;
    In cur = from;
    for (;;)
    {
        ASSERT_NE(a, a_end);
        ASSERT_NE(b, b_end);
        auto p0 = *dpf::eval_point(k0, cur);
        auto p1 = *dpf::eval_point(k1, cur);
        expect_same(*a, p0);
        expect_same(*b, p1);
        ++a;
        ++b;
        ++n;
        if (cur == to)
            break;
        cur = static_cast<In>(static_cast<std::uint64_t>(cur) + 1u);
        ASSERT_LT(n, std::size_t{1} << 20);
    }
    EXPECT_EQ(a, a_end);
    EXPECT_EQ(b, b_end);
    const std::uint64_t width = std::uint64_t{1} << dpf::utils::bitlength_of_v<In>;
    const std::uint64_t masked = (static_cast<std::uint64_t>(to)
        - static_cast<std::uint64_t>(from)) & (width - 1);
    EXPECT_EQ(n, masked + 1);
}

}  // namespace

TEST(CornerGaps, SameLeafWrapUint8Uint32)
{
    expect_wrapping_interval<uint8_t, uint32_t>(10, 9, 40, 0x11111111u, 65);
}

TEST(CornerGaps, AdjacentLeafWrapUint8Uint32)
{
    expect_wrapping_interval<uint8_t, uint32_t>(8, 7, 40, 0x22222222u, 64);
}

TEST(CornerGaps, LongWrapStillMatchesPoint)
{
    expect_wrapping_interval<uint8_t, uint32_t>(200, 10, 3, 0x33333333u, 17);

    auto [k0, k1] = dpf::make_dpf(uint8_t{3}, uint32_t{0x33333333u});
    using key_t = std::decay_t<decltype(k0)>;
    auto memo0 = dpf::make_full_tree_interval_memoizer<key_t>(uint8_t{200}, uint8_t{10});
    auto memo1 = dpf::make_full_tree_interval_memoizer<key_t>(uint8_t{200}, uint8_t{10});
    auto [buf0, it0] = dpf::eval_interval(k0, uint8_t{200}, uint8_t{10}, std::move(memo0));
    auto [buf1, it1] = dpf::eval_interval(k1, uint8_t{200}, uint8_t{10}, std::move(memo1));
    auto p0 = *dpf::eval_point(k0, uint8_t{200});
    auto p1 = *dpf::eval_point(k1, uint8_t{200});
    expect_same(*std::begin(it0), p0);
    expect_same(*std::begin(it1), p1);
    EXPECT_EQ(dpf::reconstruct(*std::begin(it0), *std::begin(it1)), 0u);
    auto back0 = std::begin(it0);
    auto back1 = std::begin(it1);
    for (int i = 0; i < 66; ++i, ++back0, ++back1) {}
    auto q0 = *dpf::eval_point(k0, uint8_t{10});
    auto q1 = *dpf::eval_point(k1, uint8_t{10});
    expect_same(*back0, q0);
    expect_same(*back1, q1);
}

TEST(CornerGaps, OneOutputPerLeafWrap)
{
    expect_wrapping_interval<uint16_t, simde_uint128>(5, 4, 9, simde_uint128{7}, 65536);
}

TEST(CornerGaps, SaturatedUint64LeafCount)
{
    using in_t = uint64_t;
    using out_t = simde_uint128;
    using key_t = dpf::utils::dpf_type_t<dpf::prg::aes128, dpf::prg::aes128, in_t, out_t>;
    EXPECT_EQ((dpf::utils::get_nodes_in_interval<key_t>(in_t{1}, ~in_t{0})),
        std::numeric_limits<std::size_t>::max());
    EXPECT_THROW((dpf::utils::get_nodes_in_interval<key_t>(in_t{0}, ~in_t{0})),
        std::length_error);
    EXPECT_THROW((dpf::utils::get_nodes_in_interval<key_t>(in_t{5}, in_t{4})),
        std::length_error);
}

TEST(CornerGaps, MemoizerRejectsALargerInterval)
{
    auto [k0, k1] = dpf::make_dpf(uint8_t{4}, uint32_t{1});
    using key_t = std::decay_t<decltype(k0)>;
    auto memo = dpf::make_basic_interval_memoizer<key_t>(uint8_t{0}, uint8_t{10});
    auto buf = dpf::make_output_buffer_for_interval<key_t>(uint8_t{0}, uint8_t{100});
    EXPECT_THROW(dpf::eval_interval(k0, uint8_t{0}, uint8_t{100}, buf, memo),
        std::length_error);
    (void)k1;
}

TEST(CornerGaps, OddStartInteriorTail)
{
    auto [k0, k1] = dpf::make_dpf(uint16_t{3}, simde_uint128{11});
    for (uint16_t to : {uint16_t{9}, uint16_t{10}})
    {
        auto [b0, it0] = dpf::eval_interval(k0, uint16_t{1}, to);
        auto [b1, it1] = dpf::eval_interval(k1, uint16_t{1}, to);
        auto a = std::begin(it0);
        auto b = std::begin(it1);
        for (uint16_t q = 1; q <= to; ++q, ++a, ++b)
        {
            auto p0 = *dpf::eval_point(k0, q);
            auto p1 = *dpf::eval_point(k1, q);
            expect_same(*a, p0);
            expect_same(*b, p1);
        }
        EXPECT_EQ(a, std::end(it0));
        EXPECT_EQ(b, std::end(it1));
    }
}

TEST(CornerGaps, PathMemoizerExtremes)
{
    auto [k0, k1] = dpf::make_dpf(uint16_t{0x0102}, uint16_t{9});
    using key_t = std::decay_t<decltype(k0)>;
    dpf::basic_path_memoizer<key_t> m0;
    dpf::basic_path_memoizer<key_t> m1;
    const uint16_t queries[] = {0, 0x8000, 1, 0};
    for (uint16_t q : queries)
    {
        auto a = *dpf::eval_point(k0, q, m0);
        auto b = *dpf::eval_point(k1, q, m1);
        auto fa = *dpf::eval_point(k0, q);
        auto fb = *dpf::eval_point(k1, q);
        EXPECT_EQ(a, fa) << q;
        EXPECT_EQ(b, fb) << q;
    }
}

TEST(CornerGaps, DepthOneBitInterval)
{
    for (uint8_t alpha : {uint8_t{0}, uint8_t{127}, uint8_t{128}, uint8_t{255}})
    {
        auto [k0, k1] = dpf::make_dpf(alpha, dpf::bit::one);
        using key_t = std::decay_t<decltype(k0)>;
        EXPECT_EQ(key_t::depth, 1u);

        auto check = [&](uint8_t from, uint8_t to) {
            auto [b0, it0] = dpf::eval_interval(k0, from, to);
            auto [b1, it1] = dpf::eval_interval(k1, from, to);
            auto a = std::begin(it0);
            auto b = std::begin(it1);
            for (uint8_t q = from; ; )
            {
                const bool on = q == alpha;
                const bool bit = static_cast<bool>(*a) != static_cast<bool>(*b);
                EXPECT_EQ(bit, on) << int(q);
                ++a;
                ++b;
                if (q == to)
                    break;
                ++q;
            }
            EXPECT_EQ(a, std::end(it0));
            EXPECT_EQ(std::end(it1), b);
        };
        check(alpha, alpha);
        check(127, 128);
    }
}

TEST(CornerGaps, EmptyAndDuplicateSequence)
{
    auto [k0, k1] = dpf::make_dpf(uint8_t{40}, uint8_t{7});
    std::vector<uint8_t> empty;
    auto [eb0, eit0] = dpf::eval_sequence(k0, empty.begin(), empty.end());
    auto [eb1, eit1] = dpf::eval_sequence(k1, empty.begin(), empty.end());
    EXPECT_EQ(std::begin(eit0), std::end(eit0));
    EXPECT_EQ(std::begin(eit1), std::end(eit1));

    const std::vector<uint8_t> seq{40, 40, 41};
    auto [b0, it0] = dpf::eval_sequence(k0, seq.begin(), seq.end());
    auto [b1, it1] = dpf::eval_sequence(k1, seq.begin(), seq.end());
    auto a = std::begin(it0);
    auto b = std::begin(it1);
    const uint8_t want[] = {7, 7, 0};
    for (int i = 0; i < 3; ++i, ++a, ++b)
        EXPECT_EQ(dpf::reconstruct(*a, *b), want[i]) << i;
    EXPECT_EQ(a, std::end(it0));
}

TEST(CornerGaps, IncrementalAdjacentLaneWrap)
{
    const uint16_t alpha = 0x00ab;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<8>(uint32_t{0xabcdu}));
    auto [b0, it0] = dpf::eval_interval(dpf::out<0, 8>, k0, uint8_t{8}, uint8_t{7});
    auto [b1, it1] = dpf::eval_interval(dpf::out<0, 8>, k1, uint8_t{8}, uint8_t{7});
    auto a = std::begin(it0);
    auto b = std::begin(it1);
    std::size_t n = 0;
    for (int lane = 8; ; )
    {
        const uint16_t query = static_cast<uint16_t>(static_cast<uint16_t>(lane) << 8);
        auto p0 = *dpf::eval_point(dpf::out<0, 8>, k0, query);
        auto p1 = *dpf::eval_point(dpf::out<0, 8>, k1, query);
        EXPECT_EQ(*a, p0) << lane;
        EXPECT_EQ(*b, p1) << lane;
        ++a;
        ++b;
        ++n;
        if (lane == 7)
            break;
        lane = (lane + 1) & 255;
    }
    EXPECT_EQ(n, 256u);
    EXPECT_EQ(a, std::end(it0));
}

TEST(CornerGaps, ModintWideLiteralAndLimbShift)
{
    using namespace dpf::literals;
    EXPECT_EQ(1_u129, dpf::modint<129>{1});
    const uint256_t bit128{1, 0};
    const auto wide = 340282366920938463463374607431768211456_u129;
    EXPECT_EQ(wide, (dpf::modint<129>{bit128}));

    const bool shift10 = (dpf::modint<10>{1} << 10) == dpf::modint<10>{0};
    const bool shift16 = (dpf::modint<10>{1} << 16) == dpf::modint<10>{0};
    const bool shift64 = (dpf::modint<64>{1} << 64) == dpf::modint<64>{0};
    const bool shift128 = (dpf::modint<65>{1} << 128) == dpf::modint<65>{0};
    const bool rshift10 = (dpf::modint<10>{5} >> 10) == dpf::modint<10>{0};
    const bool rshift64 = (dpf::modint<64>{1} >> 64) == dpf::modint<64>{0};
    EXPECT_TRUE(shift10 && shift16 && shift64 && shift128 && rshift10 && rshift64);

    dpf::modint<10> assigned{1};
    assigned <<= 16;
    EXPECT_TRUE(assigned == dpf::modint<10>{0});
    assigned = dpf::modint<10>{7};
    assigned >>= 10;
    EXPECT_TRUE(assigned == dpf::modint<10>{0});
}

TEST(CornerGaps, SetbitEmptyAndSingleAndNarrowLeaf)
{
    dpf::dynamic_bit_array<> zeros(128);
    using iter_t = decltype(zeros.begin());
    dpf::subinterval_iterable<iter_t> all(zeros.begin(), zeros.size(),
        0, zeros.size() - 1, 0, 0);
    auto none = dpf::indices_set_in(all);
    EXPECT_EQ(none.begin(), none.end());

    zeros[0] = true;
    dpf::subinterval_iterable<iter_t> one(zeros.begin(), zeros.size(),
        0, zeros.size() - 1, 0, 0);
    auto set = dpf::indices_set_in(one);
    auto it = set.begin();
    ASSERT_NE(it, set.end());
    EXPECT_EQ(*it, 0u);
    ++it;
    EXPECT_EQ(it, set.end());

    dpf::dynamic_bit_array<> narrow(128);
    narrow[0] = true;
    dpf::subinterval_iterable<iter_t> clipped(narrow.begin(), narrow.size(),
        0, 1, 0, 2);
    auto clipped_set = dpf::indices_set_in(clipped);
    auto cit = clipped_set.begin();
    ASSERT_NE(cit, clipped_set.end());
    EXPECT_EQ(*cit, 0u);
    ++cit;
    EXPECT_EQ(cit, clipped_set.end());
}

TEST(CornerGaps, EmptyRotationIsEmptyAndZeroIsIdentity)
{
    std::vector<int> empty;
    dpf::rotation_iterable<std::vector<int>::iterator> none(
        empty.begin(), empty.end(), 1);
    EXPECT_EQ(none.begin(), none.end());

    std::vector<int> values{1, 2, 3};
    dpf::rotation_iterable<std::vector<int>::iterator> id(
        values.begin(), values.end(), 0);
    std::vector<int> got;
    for (auto it = id.begin(); it != id.end(); ++it)
        got.push_back(*it);
    EXPECT_EQ(got, values);

    dpf::rotation_iterable<std::vector<int>::iterator> rot(
        values.begin(), values.end(), 1);
    got.clear();
    for (auto it = rot.begin(); it != rot.end(); ++it)
        got.push_back(*it);
    EXPECT_EQ(got, (std::vector<int>{2, 3, 1}));
}

TEST(CornerGaps, Party1NegatesSignedMinimum)
{
    using sub = dpf::subtractive_share<int32_t, 1>;
    using add0 = dpf::additive_share<int32_t, 0>;
    const auto raw = std::numeric_limits<int32_t>::min();
    const auto party1 = sub::from_raw(raw).as_additive();
    EXPECT_EQ(party1.raw(), raw);
    EXPECT_EQ(dpf::reconstruct(add0::from_raw(0), party1), raw);
    EXPECT_EQ((-sub::from_raw(raw)).raw(), raw);
}

TEST(CornerGaps, FixedMulFloorsAndPrecisionCastDiffersFromLogicalShift)
{
    using q4 = grotto::fixedpoint<4, std::int32_t>;
    const auto prod = grotto::fixed_mul<8, 4>(q4::from_raw(-3), q4::from_raw(1));
    EXPECT_EQ(prod.integral_representation(), -1);

    const auto neg_pair = grotto::fixed_mul<8, 4>(q4::from_raw(-3), q4::from_raw(-2));
    EXPECT_EQ(neg_pair.integral_representation(), 0);

    auto quarter = q4::from_raw(-12);
    const auto casted = grotto::precision_cast<0>(quarter);
    EXPECT_EQ(casted.integral_representation(), -1);
    quarter >>= 4;
    EXPECT_EQ(quarter.integral_representation(), 268435455);
}

TEST(CornerGaps, Int64MinFactorIsDefined)
{
    using grotto::principal_detail::w_from_i128;
    using grotto::principal_detail::w_mul_i64;
    using grotto::principal_detail::w_mul_u64;
    using grotto::principal_detail::w_neg;
    const auto value = w_from_i128(3);
    const auto got = w_mul_i64(value, std::numeric_limits<std::int64_t>::min());
    const auto want = w_neg(w_mul_u64(value, std::uint64_t{1} << 63));
    EXPECT_EQ(got.lo, want.lo);
    EXPECT_EQ(got.hi, want.hi);
}

TEST(CornerGaps, PrgRejectsUint32Seam)
{
    alignas(64) simde__m128i seed = simde_mm_set_epi64x(1, 2);
    alignas(64) simde__m128i out[4];
    const auto pos = static_cast<psnip_uint32_t>(UINT32_MAX - 1u);
    EXPECT_THROW(dpf::prg::aes128::eval(seed, out, 4, pos), std::invalid_argument);
    EXPECT_THROW(dpf::prg::lowmc128::eval(seed, out, 4, pos), std::invalid_argument);

    const auto ok = static_cast<psnip_uint32_t>(UINT32_MAX - 3u);
    dpf::prg::aes128::eval(seed, out, 4, ok);
    for (psnip_uint32_t i = 0; i < 4; ++i)
    {
        const auto one = dpf::prg::aes128::eval(seed, ok + i);
        EXPECT_EQ(std::memcmp(&out[i], &one, sizeof(one)), 0) << i;
    }
    dpf::prg::lowmc128::eval(seed, out, 4, ok);
    for (psnip_uint32_t i = 0; i < 4; ++i)
    {
        const auto one = dpf::prg::lowmc128::eval(seed, ok + i);
        EXPECT_EQ(std::memcmp(&out[i], &one, sizeof(one)), 0) << i;
    }

    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    EXPECT_THROW((dpf::randomness::detail::lane_codec<dpf::prg::aes128, simde__m128i>::fill(
        seed, static_cast<std::uint64_t>(UINT32_MAX) - 1u, out, 4)),
        std::invalid_argument);
    HEDLEY_PRAGMA(GCC diagnostic pop)
}