#include <gtest/gtest.h>

#include "dpf.hpp"

#include <cstdint>
#include <initializer_list>
#include <random>
#include <vector>

namespace
{

uint64_t oracle(uint64_t x, uint64_t r, uint64_t p, uint64_t q,
    uint64_t nmask, uint64_t if_true, uint64_t if_false, uint64_t gmask)
{
    const uint64_t w = (x - r) & nmask;
    const bool inside = w >= p && w <= q;
    return (inside ? if_true : if_false) & gmask;
}

template <typename Input, typename Beta>
void expect_domain(Input r, Input p, Input q, Beta if_true, Beta if_false,
    uint64_t gmask)
{
    auto keys = dpf::make_dpf(r, dpf::ic(p, q, if_true, if_false));
    const uint64_t nmask = keys.first.input_mask;
    const uint64_t rb = static_cast<uint64_t>(r);
    const uint64_t pb = static_cast<uint64_t>(p);
    const uint64_t qb = static_cast<uint64_t>(q);
    for (uint64_t x = 0; x <= nmask; ++x)
    {
        const auto y0 = dpf::eval_point(dpf::ic, keys.first,
            static_cast<Input>(x));
        const auto y1 = dpf::eval_point(dpf::ic, keys.second,
            static_cast<Input>(x));
        const uint64_t got = static_cast<uint64_t>(dpf::reconstruct(y0, y1)) & gmask;
        const uint64_t want = oracle(x, rb, pb, qb, nmask,
            static_cast<uint64_t>(if_true), static_cast<uint64_t>(if_false), gmask);
        ASSERT_EQ(got, want) << "r=" << rb << " x=" << x
            << " p=" << pb << " q=" << qb;
    }
}

struct IcPad
{
    simde__m128i block() { return dpf::uniform_sample<simde__m128i>(); }
    uint8_t bit() { return static_cast<uint8_t>(dpf::uniform_sample<uint8_t>() & 1u); }
};

using IcRng = dpf::ds_randomness<decltype(&dpf::uniform_sample<simde__m128i>), IcPad>;

IcRng ic_rng()
{
    return {&dpf::uniform_sample<simde__m128i>, {}};
}

template <typename Input>
void expect_samples(Input r, Input p, Input q, uint32_t if_true, uint32_t if_false,
    std::initializer_list<Input> xs)
{
    auto keys = dpf::make_dpf(r, dpf::ic(p, q, if_true, if_false));
    const uint64_t nmask = keys.first.input_mask;
    const uint64_t rb = static_cast<uint64_t>(r);
    const uint64_t pb = static_cast<uint64_t>(p);
    const uint64_t qb = static_cast<uint64_t>(q);
    for (Input x : xs)
    {
        const uint64_t got = static_cast<uint64_t>(dpf::reconstruct(
            dpf::eval_point(dpf::ic, keys.first, x),
            dpf::eval_point(dpf::ic, keys.second, x))) & 0xffffffffu;
        const uint64_t want = oracle(static_cast<uint64_t>(x), rb, pb, qb, nmask,
            if_true, if_false, 0xffffffffu);
        EXPECT_EQ(got, want) << "r=" << rb << " x=" << static_cast<uint64_t>(x);
    }
}

}  // namespace

TEST(Ic, Uint8FullDomainCorners)
{
    const uint32_t betas[] = {1u, 7u, 255u};
    const uint32_t falses[] = {0u, 9u};
    const uint8_t intervals[][2] = {
        {0, 0}, {0, 255}, {5, 5}, {1, 20}, {200, 250}, {0, 1}, {254, 255}, {10, 40}
    };
    for (uint32_t beta : betas)
    {
        for (uint32_t f : falses)
        {
            for (const auto & iv : intervals)
            {
                for (int r = 0; r < 256; r += 17)
                {
                    expect_domain<uint8_t>(static_cast<uint8_t>(r), iv[0], iv[1],
                        beta, f, 0xffffffffu);
                }
            }
        }
    }
}

TEST(Ic, Uint8AllMasksOneInterval)
{
    expect_domain<uint8_t>(uint8_t{0}, uint8_t{10}, uint8_t{20},
        uint32_t{3}, uint32_t{0}, 0xffffffffu);
    expect_domain<uint8_t>(uint8_t{200}, uint8_t{10}, uint8_t{100},
        uint32_t{3}, uint32_t{1}, 0xffffffffu);
    expect_domain<uint8_t>(uint8_t{255}, uint8_t{0}, uint8_t{255},
        uint32_t{1}, uint32_t{0}, 0xffffffffu);
}

TEST(Ic, MemoizerAgrees)
{
    const uint8_t r = 40, p = 7, q = 90;
    auto keys = dpf::make_dpf(r, dpf::ic(p, q, uint32_t{11}, uint32_t{2}));
    dpf::basic_path_memoizer<decltype(keys.first.key)> memo0;
    dpf::basic_path_memoizer<decltype(keys.second.key)> memo1;
    for (int x = 0; x < 256; ++x)
    {
        const auto a = dpf::eval_point(dpf::ic, keys.first, static_cast<uint8_t>(x), memo0);
        const auto b = dpf::eval_point(dpf::ic, keys.second, static_cast<uint8_t>(x), memo1);
        const auto c = dpf::eval_point(dpf::ic, keys.first, static_cast<uint8_t>(x));
        const auto d = dpf::eval_point(dpf::ic, keys.second, static_cast<uint8_t>(x));
        EXPECT_EQ(dpf::reconstruct(a, b), dpf::reconstruct(c, d));
    }
}

TEST(Ic, IntervalAndSequenceBuffers)
{
    const uint8_t r = 15, p = 4, q = 12;
    auto keys = dpf::make_dpf(r, dpf::ic(p, q, uint16_t{9}));
    auto buf0 = dpf::make_output_buffer(dpf::ic, keys.first, uint8_t{3}, uint8_t{18});
    auto buf1 = dpf::make_output_buffer(dpf::ic, keys.second, uint8_t{3}, uint8_t{18});
    dpf::basic_path_memoizer<decltype(keys.first.key)> memo;
    dpf::eval_interval(dpf::ic, keys.first, uint8_t{3}, uint8_t{18}, buf0, memo);
    dpf::eval_interval(dpf::ic, keys.second, uint8_t{3}, uint8_t{18}, buf1);
    for (std::size_t i = 0; i < buf0.size(); ++i)
    {
        const auto point = dpf::reconstruct(
            dpf::eval_point(dpf::ic, keys.first, static_cast<uint8_t>(3 + i)),
            dpf::eval_point(dpf::ic, keys.second, static_cast<uint8_t>(3 + i)));
        EXPECT_EQ(dpf::reconstruct(buf0[i], buf1[i]), point);
    }

    const uint8_t pts[] = {0, 9, 15, 255, 4};
    auto s0 = dpf::make_output_buffer(dpf::ic, keys.first, 5);
    auto s1 = dpf::make_output_buffer(dpf::ic, keys.second, 5);
    dpf::eval_sequence(dpf::ic, keys.first, std::begin(pts), std::end(pts), s0);
    dpf::eval_sequence(dpf::ic, keys.second, std::begin(pts), std::end(pts), s1);
    for (std::size_t i = 0; i < 5; ++i)
    {
        const auto point = dpf::reconstruct(
            dpf::eval_point(dpf::ic, keys.first, pts[i]),
            dpf::eval_point(dpf::ic, keys.second, pts[i]));
        EXPECT_EQ(dpf::reconstruct(s0[i], s1[i]), point);
    }
}

TEST(Ic, WildcardAssign)
{
    auto keys = dpf::make_dpf(uint8_t{33},
        dpf::ic(uint8_t{2}, uint8_t{8}, dpf::wildcard<uint32_t>));
    EXPECT_THROW(dpf::eval_point(dpf::ic, keys.first, uint8_t{0}), std::invalid_argument);
    dpf::assign_cmp(keys.first, keys.second, uint32_t{6}, uint32_t{1});
    expect_domain<uint8_t>(uint8_t{33}, uint8_t{2}, uint8_t{8},
        uint32_t{6}, uint32_t{1}, 0xffffffffu);
    // The keys just assigned are a different generation; check those directly.
    for (int x = 0; x < 256; ++x)
    {
        const uint64_t got = static_cast<uint64_t>(dpf::reconstruct(
            dpf::eval_point(dpf::ic, keys.first, static_cast<uint8_t>(x)),
            dpf::eval_point(dpf::ic, keys.second, static_cast<uint8_t>(x))));
        const uint64_t w = static_cast<uint64_t>(static_cast<uint8_t>(x - 33));
        const uint64_t want = (w >= 2 && w <= 8) ? 6u : 1u;
        EXPECT_EQ(got, want) << x;
    }
}

TEST(Ic, DoernerShelatMatchesDealer)
{
    std::mt19937 rng{7};
    std::uniform_int_distribution<int> d(0, 255);
    for (int n = 0; n < 30; ++n)
    {
        const uint8_t r0 = static_cast<uint8_t>(d(rng));
        const uint8_t r1 = static_cast<uint8_t>(d(rng));
        const uint8_t p = static_cast<uint8_t>(d(rng));
        const uint8_t q = static_cast<uint8_t>(p + static_cast<uint8_t>(d(rng) % (256 - p)));
        const uint32_t beta = 1u + static_cast<uint32_t>(d(rng));
        const uint8_t r = static_cast<uint8_t>(r0 ^ r1);
        auto dealer = dpf::make_dpf(r, dpf::ic(p, q, beta));
        struct Pad
        {
            simde__m128i block() { return dpf::uniform_sample<simde__m128i>(); }
            uint8_t bit() { return static_cast<uint8_t>(dpf::uniform_sample<uint8_t>() & 1u); }
        };
        dpf::ds_randomness<decltype(&dpf::uniform_sample<simde__m128i>), Pad> rngs{
            &dpf::uniform_sample<simde__m128i>, {}};
        auto ds = dpf::make_dpf_doerner_shelat(r0, r1, rngs, dpf::ic(p, q, beta));
        for (int x = 0; x < 256; x += 5)
        {
            const auto dealer_y = dpf::reconstruct(
                dpf::eval_point(dpf::ic, dealer.first, static_cast<uint8_t>(x)),
                dpf::eval_point(dpf::ic, dealer.second, static_cast<uint8_t>(x)));
            const auto ds_y = dpf::reconstruct(
                dpf::eval_point(dpf::ic, ds.first, static_cast<uint8_t>(x)),
                dpf::eval_point(dpf::ic, ds.second, static_cast<uint8_t>(x)));
            EXPECT_EQ(dealer_y, ds_y) << "x=" << x;
        }
    }
}

TEST(Ic, Geneval)
{
    struct Pad
    {
        simde__m128i block() { return dpf::uniform_sample<simde__m128i>(); }
        uint8_t bit() { return static_cast<uint8_t>(dpf::uniform_sample<uint8_t>() & 1u); }
    };
    const uint8_t r0 = 9, r1 = 100, p = 3, q = 50;
    const uint32_t beta = 4;
    const uint8_t queries[] = {0, 3, 12, 49, 50, 51, 255};
    dpf::ds_randomness<decltype(&dpf::uniform_sample<simde__m128i>), Pad> rngs{
        &dpf::uniform_sample<simde__m128i>, {}};
    auto opened = dpf::geneval_ic(r0, r1, std::begin(queries), std::end(queries),
        rngs, dpf::ic(p, q, beta));
    ASSERT_EQ(opened.party0.size(), 7u);
    ASSERT_EQ(opened.live_levels, 8u);
    const uint8_t r = static_cast<uint8_t>(r0 ^ r1);
    for (std::size_t i = 0; i < 7; ++i)
    {
        const uint64_t got = (opened.party0[i] + opened.party1[i]) & 0xffffffffu;
        const uint64_t w = static_cast<uint64_t>(
            static_cast<uint8_t>(queries[i] - r));
        const uint64_t want = (w >= p && w <= q) ? beta : 0u;
        EXPECT_EQ(got, want) << i;
    }
}

TEST(Ic, RejectsWrappedBounds)
{
    EXPECT_THROW(dpf::make_dpf(uint8_t{1}, dpf::ic(uint8_t{9}, uint8_t{2}, uint32_t{1})),
        std::invalid_argument);
}

TEST(Ic, BitPayload)
{
    auto keys = dpf::make_dpf(uint8_t{4},
        dpf::ic(uint8_t{1}, uint8_t{3}, dpf::bit::one, dpf::bit::zero));
    for (int x = 0; x < 256; ++x)
    {
        const auto y = dpf::reconstruct(
            dpf::eval_point(dpf::ic, keys.first, static_cast<uint8_t>(x)),
            dpf::eval_point(dpf::ic, keys.second, static_cast<uint8_t>(x)));
        const uint64_t w = static_cast<uint64_t>(static_cast<uint8_t>(x - 4));
        EXPECT_EQ(static_cast<bool>(y), w >= 1 && w <= 3) << x;
    }
}

TEST(Ic, Uint16FullDomain)
{
    expect_domain<uint16_t>(uint16_t{0x0100}, uint16_t{20}, uint16_t{400},
        uint32_t{9}, uint32_t{3}, 0xffffffffu);
    expect_domain<uint16_t>(uint16_t{0xFF00}, uint16_t{0}, uint16_t{1},
        uint32_t{1}, uint32_t{0}, 0xffffffffu);
}

TEST(Ic, WideMasksSampleTheWrap)
{
    const uint32_t p32 = 10, q32 = 1000;
    const uint32_t r32 = 0xFFFFFFF0u;
    expect_samples<uint32_t>(r32, p32, q32, 7u, 2u, {
        0u, 1u, p32, q32, q32 + 1u, r32, r32 - 1u, r32 + 1u,
        0x80000000u, 0xffffffffu
    });

    const uint64_t p64 = 1, q64 = 3;
    const uint64_t r64 = ~uint64_t{0};
    expect_samples<uint64_t>(r64, p64, q64, 5u, 4u, {
        0ull, 1ull, 2ull, 3ull, 4ull, r64, r64 - 1ull,
        uint64_t{1} << 63
    });
    auto keys = dpf::make_dpf(r64, dpf::ic(p64, q64, uint32_t{5}, uint32_t{4}));
    EXPECT_EQ(keys.first.input_mask, ~uint64_t{0});
}

TEST(Ic, AdditiveDoernerShelatMatchesDealer)
{
    auto check = [](uint8_t r0, uint8_t r1, uint8_t p, uint8_t q,
                    uint32_t beta, uint32_t fals) {
        const uint8_t r = static_cast<uint8_t>(r0 + r1);
        auto dealer = dpf::make_dpf(r, dpf::ic(p, q, beta, fals));
        auto ds = dpf::make_dpf_doerner_shelat(dpf::arith_input, r0, r1, ic_rng(),
            dpf::ic(p, q, beta, fals));
        for (int x = 0; x < 256; ++x)
        {
            const auto dealer_y = dpf::reconstruct(
                dpf::eval_point(dpf::ic, dealer.first, static_cast<uint8_t>(x)),
                dpf::eval_point(dpf::ic, dealer.second, static_cast<uint8_t>(x)));
            const auto ds_y = dpf::reconstruct(
                dpf::eval_point(dpf::ic, ds.first, static_cast<uint8_t>(x)),
                dpf::eval_point(dpf::ic, ds.second, static_cast<uint8_t>(x)));
            EXPECT_EQ(dealer_y, ds_y) << "r0=" << int(r0) << " r1=" << int(r1)
                << " x=" << x;
        }
    };
    check(9, 100, 3, 50, 4, 0);
    check(200, 100, 0, 255, 8, 1);
    check(200, 200, 10, 20, 6, 2);
    check(0, 0, 1, 1, 1, 0);
    check(1, 0, 0, 0, 3, 9);
}

TEST(Ic, AdditiveGeneval)
{
    const uint8_t r0 = 250, r1 = 20, p = 4, q = 8;
    const uint32_t beta = 11, fals = 2;
    const uint8_t r = static_cast<uint8_t>(r0 + r1);
    const uint8_t queries[] = {0, 1, 4, 8, 9, 255};
    auto opened = dpf::geneval_ic(dpf::arith_input, r0, r1,
        std::begin(queries), std::end(queries), ic_rng(),
        dpf::ic(p, q, beta, fals));
    ASSERT_EQ(opened.party0.size(), 6u);
    for (std::size_t i = 0; i < 6; ++i)
    {
        const uint64_t got = (opened.party0[i] + opened.party1[i]) & 0xffffffffu;
        const uint64_t w = static_cast<uint64_t>(static_cast<uint8_t>(queries[i] - r));
        const uint64_t want = (w >= p && w <= q) ? beta : fals;
        EXPECT_EQ(got, want) << i;
    }

    const uint8_t none[] = {0};
    auto empty = dpf::geneval_ic(dpf::arith_input, r0, r1,
        std::begin(none), std::begin(none), ic_rng(), dpf::ic(p, q, beta));
    EXPECT_TRUE(empty.party0.empty());
    EXPECT_EQ(empty.live_levels, 0u);
}

TEST(Ic, IntervalRejectsDescendingEndpoints)
{
    auto keys = dpf::make_dpf(uint8_t{4}, dpf::ic(uint8_t{1}, uint8_t{6}, uint32_t{1}));
    auto buf = dpf::make_output_buffer(dpf::ic, keys.first, 1);
    EXPECT_THROW(dpf::eval_interval(dpf::ic, keys.first, uint8_t{9}, uint8_t{2}, buf),
        std::invalid_argument);
    EXPECT_THROW(dpf::make_output_buffer(dpf::ic, keys.first, uint8_t{9}, uint8_t{2}),
        std::invalid_argument);
}