#include <gtest/gtest.h>

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

#include <cstdint>
#include <utility>

namespace
{

template <typename Beta, typename A, typename B, typename X>
Beta recon_cmp(const A & a, const B & b, X x)
{
    return dpf::reconstruct(
        dpf::eval_point<Beta>(dpf::cmp, a, x),
        dpf::eval_point<Beta>(dpf::cmp, b, x));
}

template <typename Beta, typename A, typename B, typename X>
Beta recon_prefix4(const A & a, const B & b, X x)
{
    return dpf::reconstruct(
        dpf::eval_point<4, Beta>(dpf::cmp_prefix<4>, a, x),
        dpf::eval_point<4, Beta>(dpf::cmp_prefix<4>, b, x));
}

}  // namespace

TEST(WidePayload, VectorPointAddsComponentwise)
{
    using out_t = dpf::vec<std::uint32_t, 4>;
    const std::uint16_t alpha = 0x1234;
    out_t y;
    y[0] = 1;
    y[1] = 0xffffffffu;
    y[2] = 7;
    y[3] = 100;
    auto [k0, k1] = dpf::make_dpf(alpha, y);
    auto at = [&](std::uint16_t x) {
        return dpf::reconstruct(*dpf::eval_point(k0, x), *dpf::eval_point(k1, x));
    };
    EXPECT_EQ(at(alpha), y);
    EXPECT_EQ(at(0), out_t{});
    EXPECT_EQ(at(static_cast<std::uint16_t>(alpha + 1)), out_t{});
}

TEST(WidePayload, Uint128Comparison)
{
    using beta = simde_uint128;
    const std::uint8_t alpha = 0x20;
    const beta hi = (beta{1} << 80) + 9;
    const beta lo = beta{3};
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(hi, lo));
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<std::uint8_t>(x);
        const beta got = recon_cmp<beta>(k0, k1, q);
        const beta want = q > alpha ? hi : lo;
        EXPECT_EQ(got, want) << int(q);
    }
}

TEST(WidePayload, FixedpointComparison)
{
    using beta = grotto::fixedpoint<8, std::uint32_t>;
    const std::uint8_t alpha = 10;
    const beta hi = beta::from_raw(0x01020304u);
    const beta lo = beta::from_raw(0x00000007u);
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::lt(hi, lo));
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{9}), hi);
    EXPECT_EQ(recon_cmp<beta>(k0, k1, alpha), lo);
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{11}), lo);
}

TEST(WidePayload, WideFixedpointComparison)
{
    using beta = grotto::fixedpoint<4, simde_uint128>;
    const std::uint8_t alpha = 4;
    const beta hi = beta::from_raw((simde_uint128{1} << 70) + 11);
    const beta lo = beta::from_raw(simde_uint128{2});
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::geq(hi, lo));
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{3}), lo);
    EXPECT_EQ(recon_cmp<beta>(k0, k1, alpha), hi);
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{5}), hi);
}

TEST(WidePayload, VectorComparison)
{
    using beta = dpf::vec<std::uint16_t, 3>;
    const std::uint8_t alpha = 7;
    beta hi;
    hi[0] = 9;
    hi[1] = 1000;
    hi[2] = 4;
    beta lo;
    lo[1] = 1;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(hi, lo));
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{6}), lo);
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{8}), hi);
}

TEST(WidePayload, IntervalContainmentUint128)
{
    using beta = simde_uint128;
    const std::uint8_t r = 20;
    const beta hi = (beta{1} << 96) + 5;
    auto [k0, k1] = dpf::make_dpf(r, dpf::ic(std::uint8_t{3}, std::uint8_t{5}, hi));
    auto at = [&](std::uint8_t opened) {
        return dpf::reconstruct(
            dpf::eval_point(dpf::ic, k0, opened),
            dpf::eval_point(dpf::ic, k1, opened));
    };
    EXPECT_EQ(at(23), hi);
    EXPECT_EQ(at(25), hi);
    EXPECT_EQ(at(22), beta{});
    EXPECT_EQ(at(26), beta{});
}

TEST(WidePayload, IdcfUint128MatchesFullComparison)
{
    using beta = simde_uint128;
    const std::uint8_t alpha = 0x6e;
    const beta y = (beta{1} << 100) + 13;
    auto [i0, i1] = dpf::make_dpf(alpha, dpf::idcf(dpf::gt(y)));
    auto [f0, f1] = dpf::make_dpf(alpha, dpf::gt(y));
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<std::uint8_t>(x);
        EXPECT_EQ(recon_cmp<beta>(i0, i1, q), recon_cmp<beta>(f0, f1, q));
        const bool above = (q >> 4) > (alpha >> 4);
        const beta prefix = recon_prefix4<beta>(i0, i1, q);
        const beta want = above ? y : beta{};
        EXPECT_EQ(prefix, want) << int(q);
    }
}

TEST(WidePayload, DealerMatchesDoernerShelat)
{
    using beta = simde_uint128;
    const std::uint16_t alpha = 0x0102;
    const std::uint16_t x0 = 0x00f0;
    const std::uint16_t x1 = static_cast<std::uint16_t>(alpha ^ x0);
    const beta y = (beta{1} << 77) + 4;
    auto dealer = dpf::make_dpf(alpha, dpf::gt(y));
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    auto ds = dpf::make_dpf_doerner_shelat(x0, x1,
        dpf::ds_randomness<simde__m128i (*)(), dpf::detail::urandom_pad_rng>{
            dpf::uniform_sample<simde__m128i>, {}},
        dpf::gt(y));
    HEDLEY_PRAGMA(GCC diagnostic pop)
    for (std::uint16_t q : {std::uint16_t{0}, std::uint16_t{0x0101}, alpha,
            std::uint16_t{0x0103}, std::uint16_t{0xffff}})
    {
        const beta from_dealer = recon_cmp<beta>(dealer.first, dealer.second, q);
        const beta from_ds = recon_cmp<beta>(ds.first, ds.second, q);
        EXPECT_EQ(from_dealer, from_ds) << q;
    }
}

TEST(WidePayload, XorWrapperComparison)
{
    using beta = dpf::xor_wrapper<std::uint32_t>;
    const std::uint8_t alpha = 9;
    const beta hi{0x01020304u};
    const beta lo{0x0000000fu};
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(hi, lo));
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{8}), lo);
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{10}), hi);
}

TEST(WidePayload, WildcardAssignUint128)
{
    using beta = simde_uint128;
    const std::uint8_t alpha = 15;
    const beta hi = (beta{1} << 90) + 8;
    const beta lo = beta{2};
    const auto wild = dpf::wildcard<beta>;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::gt(wild));
    dpf::assign_cmp(k0, k1, hi, lo);
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{14}), lo);
    EXPECT_EQ(recon_cmp<beta>(k0, k1, std::uint8_t{16}), hi);
}

TEST(WidePayload, IntervalContainmentUint128FalseAndTopBound)
{
    using beta = simde_uint128;
    const std::uint8_t r = 10;
    const beta hi = (beta{1} << 90) + 8;
    const beta lo = beta{3};
    auto [k0, k1] = dpf::make_dpf(r,
        dpf::ic(std::uint8_t{5}, std::uint8_t{255}, hi, lo));
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<std::uint8_t>(x);
        const auto w = static_cast<std::uint8_t>(q - r);
        const beta want = (w >= 5) ? hi : lo;
        EXPECT_EQ(dpf::reconstruct(
                dpf::eval_point(dpf::ic, k0, q),
                dpf::eval_point(dpf::ic, k1, q)),
            want) << int(q);
    }
}

TEST(WidePayload, IntervalContainmentXorGroup)
{
    using beta = dpf::xor_wrapper<std::uint32_t>;
    const std::uint8_t r = 40;
    const beta hi{0x11111111u};
    const beta lo{0x01010101u};
    auto [k0, k1] = dpf::make_dpf(r,
        dpf::ic(std::uint8_t{1}, std::uint8_t{255}, hi, lo));
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<std::uint8_t>(x);
        const auto w = static_cast<std::uint8_t>(q - r);
        const beta want = (w >= 1) ? hi : lo;
        EXPECT_EQ(dpf::reconstruct(
                dpf::eval_point(dpf::ic, k0, q),
                dpf::eval_point(dpf::ic, k1, q)),
            want) << int(q);
    }
}

TEST(WidePayload, IntervalContainmentVector)
{
    using beta = dpf::vec<std::uint32_t, 2>;
    beta hi;
    hi[0] = 9;
    hi[1] = 0xffffffffu;
    beta lo;
    lo[0] = 1;
    lo[1] = 2;
    const std::uint8_t r = 8;
    auto [k0, k1] = dpf::make_dpf(r,
        dpf::ic(std::uint8_t{2}, std::uint8_t{4}, hi, lo));
    auto at = [&](std::uint8_t opened) {
        return dpf::reconstruct(
            dpf::eval_point(dpf::ic, k0, opened),
            dpf::eval_point(dpf::ic, k1, opened));
    };
    EXPECT_EQ(at(10), hi);
    EXPECT_EQ(at(12), hi);
    EXPECT_EQ(at(9), lo);
    EXPECT_EQ(at(13), lo);
}

TEST(WidePayload, IntervalWildcardAssignUint128)
{
    using beta = simde_uint128;
    const std::uint8_t r = 15;
    const beta hi = (beta{1} << 90) + 8;
    const beta lo = beta{2};
    auto [k0, k1] = dpf::make_dpf(r,
        dpf::ic(std::uint8_t{2}, std::uint8_t{8}, dpf::wildcard<beta>));
    EXPECT_THROW(dpf::eval_point(dpf::ic, k0, std::uint8_t{0}), std::invalid_argument);
    dpf::assign_cmp(k0, k1, hi, lo);
    auto at = [&](std::uint8_t opened) {
        return dpf::reconstruct(
            dpf::eval_point(dpf::ic, k0, opened),
            dpf::eval_point(dpf::ic, k1, opened));
    };
    EXPECT_EQ(at(17), hi);
    EXPECT_EQ(at(23), hi);
    EXPECT_EQ(at(16), lo);
    EXPECT_EQ(at(24), lo);
}

TEST(WidePayload, IntervalDoernerShelatUint128)
{
    using beta = simde_uint128;
    const std::uint8_t r0 = 9;
    const std::uint8_t r1 = 100;
    const std::uint8_t r = static_cast<std::uint8_t>(r0 ^ r1);
    const beta hi = (beta{1} << 70) + 4;
    const beta lo = beta{1};
    auto dealer = dpf::make_dpf(r, dpf::ic(std::uint8_t{3}, std::uint8_t{50}, hi, lo));
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    auto ds = dpf::make_dpf_doerner_shelat(r0, r1,
        dpf::ds_randomness<simde__m128i (*)(), dpf::detail::urandom_pad_rng>{
            dpf::uniform_sample<simde__m128i>, {}},
        dpf::ic(std::uint8_t{3}, std::uint8_t{50}, hi, lo));
    HEDLEY_PRAGMA(GCC diagnostic pop)
    for (int x = 0; x < 256; x += 17)
    {
        const auto q = static_cast<std::uint8_t>(x);
        EXPECT_EQ(dpf::reconstruct(
                dpf::eval_point(dpf::ic, dealer.first, q),
                dpf::eval_point(dpf::ic, dealer.second, q)),
            dpf::reconstruct(
                dpf::eval_point(dpf::ic, ds.first, q),
                dpf::eval_point(dpf::ic, ds.second, q)))
            << int(q);
    }
}

TEST(WidePayload, Modint96ComparisonAndInterval)
{
    using beta = dpf::modint<96>;
    const beta hi{simde_uint128{1} << 80};
    const beta lo{simde_uint128{6}};
    const std::uint8_t alpha = 12;
    auto cmp = dpf::make_dpf(alpha, dpf::lt(hi, lo));
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<std::uint8_t>(x);
        const beta want = q < alpha ? hi : lo;
        EXPECT_EQ(recon_cmp<beta>(cmp.first, cmp.second, q), want) << int(q);
    }

    const std::uint8_t r = 4;
    auto ic = dpf::make_dpf(r, dpf::ic(std::uint8_t{1}, std::uint8_t{3}, hi, lo));
    auto at = [&](std::uint8_t opened) {
        return dpf::reconstruct(
            dpf::eval_point(dpf::ic, ic.first, opened),
            dpf::eval_point(dpf::ic, ic.second, opened));
    };
    EXPECT_EQ(at(5), hi);
    EXPECT_EQ(at(7), hi);
    EXPECT_EQ(at(4), lo);
    EXPECT_EQ(at(8), lo);
}

TEST(WidePayload, VectorLaneArithmeticDoesNotCarry)
{
    dpf::vec<std::uint32_t, 2> a;
    dpf::vec<std::uint32_t, 2> b;
    a[0] = 0xffffffffu;
    a[1] = 1u;
    b[0] = 2u;
    b[1] = 3u;
    const auto sum = a + b;
    EXPECT_EQ(sum[0], 1u);
    EXPECT_EQ(sum[1], 4u);
    const auto prod = a * b;
    EXPECT_EQ(prod[0], 0xfffffffeu);
    EXPECT_EQ(prod[1], 3u);
    const auto neg = -a;
    EXPECT_EQ(neg[0], 1u);
    EXPECT_EQ(neg[1], 0xffffffffu);
    EXPECT_EQ(a, a);
    EXPECT_NE(a, b);
}