#include <gtest/gtest.h>

#include "dpf.hpp"

#include <cstdint>
#include <utility>

namespace
{

template <typename A, typename B, typename Target>
uint64_t recon_cmp(const A & a, const B & b, Target target, uint8_t x)
{
    return dpf::reconstruct(dpf::eval_point(target, a, x),
        dpf::eval_point(target, b, x));
}

uint64_t lcp_len(uint8_t x, uint8_t alpha, std::size_t n = 8, std::size_t width = 8)
{
    for (std::size_t i = 0; i < n; ++i)
    {
        const uint8_t shift = static_cast<uint8_t>(width - 1 - i);
        if (((x >> shift) & 1) != ((alpha >> shift) & 1))
            return i;
    }
    return n;
}

uint64_t high_prefix(uint8_t alpha, uint64_t d, std::size_t n = 8)
{
    if (d == 0)
        return 0;
    if (d >= n)
        return alpha & ((1u << n) - 1u);
    const unsigned drop = static_cast<unsigned>(n - d);
    return (static_cast<unsigned>(alpha) >> drop) << drop;
}

uint64_t low_prefix(uint8_t alpha, uint64_t d, std::size_t n = 8)
{
    if (d == 0)
        return 0;
    if (d >= n)
        return alpha;
    return static_cast<unsigned>(alpha) >> (n - d);
}

template <typename Make, typename Unit>
void expect_domain(Make make, Unit unit_of)
{
    constexpr uint8_t alpha = 0xB4;
    auto [k0, k1] = make(alpha);
    for (int x = 0; x < 256; ++x)
    {
        const auto got = recon_cmp(k0, k1, dpf::cmp, static_cast<uint8_t>(x));
        EXPECT_EQ(got, unit_of(static_cast<uint8_t>(x), alpha))
            << "x=" << x;
    }
}

}  // namespace

TEST(PathRecipe, LengthMaskPrefixBreakAndPacked)
{
    constexpr uint8_t alpha = 0xB4;
    expect_domain(
        [](uint8_t a) { return dpf::make_dpf(a, dpf::lcp(uint64_t{1})); },
        [](uint8_t x, uint8_t a) { return lcp_len(x, a); });
    expect_domain(
        [](uint8_t a) { return dpf::make_dpf(a, dpf::common_prefix(uint64_t{1})); },
        [](uint8_t x, uint8_t a) { return high_prefix(a, lcp_len(x, a)); });
    expect_domain(
        [](uint8_t a) { return dpf::make_dpf(a, dpf::prefix_mask(uint64_t{1})); },
        [](uint8_t x, uint8_t a) { return high_prefix(0xFF, lcp_len(x, a)); });
    expect_domain(
        [](uint8_t a) { return dpf::make_dpf(a, dpf::diverge_one_hot(uint64_t{1})); },
        [](uint8_t x, uint8_t a) { return 1ULL << lcp_len(x, a); });
    expect_domain(
        [](uint8_t a) { return dpf::make_dpf(a, dpf::break_bit(uint64_t{3})); },
        [](uint8_t x, uint8_t a) {
            const auto d = lcp_len(x, a);
            if (d >= 8)
                return 0ULL;
            return 3ULL * ((a >> (7 - d)) & 1);
        });
    expect_domain(
        [](uint8_t a) {
            return dpf::make_dpf(a, dpf::prefix_with_length<4>(uint64_t{1}));
        },
        [](uint8_t x, uint8_t a) {
            const auto d = lcp_len(x, a);
            return (low_prefix(a, d) << 4) | d;
        });
    expect_domain(
        [](uint8_t a) {
            return dpf::make_dpf(a, dpf::lcp(uint64_t{5}, uint64_t{2}));
        },
        [](uint8_t x, uint8_t a) { return 2ULL + 3ULL * lcp_len(x, a); });
    expect_domain(
        [](uint8_t a) {
            return dpf::make_dpf(a, dpf::path_paint(
                [](std::size_t matched, uint64_t, bool) {
                    return static_cast<uint64_t>(matched);
                }));
        },
        [](uint8_t x, uint8_t a) { return lcp_len(x, a); });

    auto [p0, p1] = dpf::make_dpf(alpha, dpf::lcp_at<4>(uint64_t{1}));
    for (int x = 0; x < 256; ++x)
    {
        const auto got = recon_cmp(p0, p1, dpf::cmp, static_cast<uint8_t>(x));
        EXPECT_EQ(got, lcp_len(static_cast<uint8_t>(x), alpha, 4, 8)) << x;
    }
}

TEST(PathRecipe, WildcardScaleAssignsLength)
{
    constexpr uint8_t alpha = 0x3C;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::lcp(dpf::wildcard<uint64_t>));
    dpf::assign_cmp(k0, k1, uint64_t{4});
    for (int x = 0; x < 256; ++x)
    {
        EXPECT_EQ(recon_cmp(k0, k1, dpf::cmp, static_cast<uint8_t>(x)),
            4ULL * lcp_len(static_cast<uint8_t>(x), alpha))
            << x;
    }
}

TEST(PathRecipe, IdpfSlotsArePrefixPointFunctions)
{
    constexpr uint8_t alpha = 0xA6;
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::idpf(uint64_t{11}, uint64_t{22}, uint64_t{33}));
    auto slot = [&](auto target, uint8_t x) {
        return dpf::reconstruct(*dpf::eval_point(target, k0, x),
            *dpf::eval_point(target, k1, x));
    };
    for (int x = 0; x < 256; ++x)
    {
        const auto d = lcp_len(static_cast<uint8_t>(x), alpha);
        EXPECT_EQ(slot(dpf::out<0>, static_cast<uint8_t>(x)), d >= 1 ? 11u : 0u);
        EXPECT_EQ(slot(dpf::out<1>, static_cast<uint8_t>(x)), d >= 2 ? 22u : 0u);
        EXPECT_EQ(slot(dpf::out<2>, static_cast<uint8_t>(x)), d >= 3 ? 33u : 0u);
    }

    auto [s0, s1] = dpf::make_dpf(alpha, dpf::idpf_at<4, 7>(uint8_t{9}, uint8_t{8}));
    for (int x = 0; x < 256; ++x)
    {
        const auto d = lcp_len(static_cast<uint8_t>(x), alpha);
        auto at = [&](auto target) {
            return dpf::reconstruct(*dpf::eval_point(target, s0, static_cast<uint8_t>(x)),
                *dpf::eval_point(target, s1, static_cast<uint8_t>(x)));
        };
        EXPECT_EQ(at(dpf::out<0>), d >= 4 ? 9u : 0u);
        EXPECT_EQ(at(dpf::out<1>), d >= 7 ? 8u : 0u);
    }
}

TEST(PathRecipe, IdcfMatchesComparisonAtEveryPrefix)
{
    constexpr uint8_t alpha = 0x6E;
    auto check = [&](auto idcf_spec, auto at_spec, auto full_spec, std::size_t L,
                      auto target) {
        auto [i0, i1] = dpf::make_dpf(alpha, idcf_spec);
        auto [n0, n1] = dpf::make_dpf(alpha, at_spec);
        auto [f0, f1] = dpf::make_dpf(alpha, full_spec);
        for (int x = 0; x < 256; ++x)
        {
            const auto q = static_cast<uint8_t>(x);
            EXPECT_EQ(recon_cmp(i0, i1, target, q), recon_cmp(n0, n1, dpf::cmp, q))
                << "L=" << L << " x=" << x;
            EXPECT_EQ(recon_cmp(i0, i1, dpf::cmp, q), recon_cmp(f0, f1, dpf::cmp, q))
                << "full x=" << x;
        }
        EXPECT_EQ(i0.prefix_cw(8), i0.cw_last());
    };
    check(dpf::idcf(dpf::lt(uint64_t{1})), dpf::lt_at<4>(uint64_t{1}),
        dpf::lt(uint64_t{1}), 4, dpf::cmp_prefix<4>);
    check(dpf::idcf(dpf::leq(uint64_t{1})), dpf::leq_at<3>(uint64_t{1}),
        dpf::leq(uint64_t{1}), 3, dpf::cmp_prefix<3>);
    check(dpf::idcf(dpf::gt(uint64_t{1})), dpf::gt_at<5>(uint64_t{1}),
        dpf::gt(uint64_t{1}), 5, dpf::cmp_prefix<5>);
    check(dpf::idcf(dpf::geq(uint64_t{1})), dpf::geq_at<1>(uint64_t{1}),
        dpf::geq(uint64_t{1}), 1, dpf::cmp_prefix<1>);

    auto [z0, z1] = dpf::make_dpf(alpha, dpf::idcf(dpf::lt(uint64_t{1})));
    auto [e0, e1] = dpf::make_dpf(alpha, dpf::idcf(dpf::leq(uint64_t{1})));
    auto [g0, g1] = dpf::make_dpf(alpha, dpf::idcf(dpf::gt(uint64_t{1})));
    auto [q0, q1] = dpf::make_dpf(alpha, dpf::idcf(dpf::geq(uint64_t{1})));
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<uint8_t>(x);
        EXPECT_EQ(recon_cmp(z0, z1, dpf::cmp_prefix<0>, q), 0u);
        EXPECT_EQ(recon_cmp(e0, e1, dpf::cmp_prefix<0>, q), 1u);
        EXPECT_EQ(recon_cmp(g0, g1, dpf::cmp_prefix<0>, q), 0u);
        EXPECT_EQ(recon_cmp(q0, q1, dpf::cmp_prefix<0>, q), 1u);
    }
}

TEST(PathRecipe, DoernerShelatAndGenevalMatchLength)
{
    constexpr uint8_t alpha = 0x91;
    const uint8_t x0 = 0x10;
    const uint8_t x1 = static_cast<uint8_t>(alpha ^ x0);
    auto [k0, k1] = dpf::make_dpf_doerner_shelat(x0, x1, dpf::lcp(uint64_t{1}));
    for (int x = 0; x < 256; ++x)
    {
        EXPECT_EQ(recon_cmp(k0, k1, dpf::cmp, static_cast<uint8_t>(x)),
            lcp_len(static_cast<uint8_t>(x), alpha));
    }

    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), dpf::detail::urandom_pad_rng> rng{
        dpf::uniform_sample<simde__m128i>, {}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    const uint8_t ends[] = {0x00, 0x91, 0xFF};
    auto g = dpf::geneval_cmp(x0, x1, std::begin(ends), std::end(ends), rng,
        dpf::lcp(uint64_t{1}));
    ASSERT_EQ(g.party0.size(), 3u);
    for (std::size_t i = 0; i < 3; ++i)
    {
        EXPECT_EQ((g.party0[i] + g.party1[i]) & g.mask,
            lcp_len(ends[i], alpha));
    }
}