#include <gtest/gtest.h>

#include "dpf.hpp"
#include "dpf/json.hpp"
#include "grotto/offset_horner.hpp"
#include "grotto/prefix_parity.hpp"

#include <array>
#include <cstdint>
#include <cstring>
#include <type_traits>
#include <vector>

namespace
{

simde__m128i g_roots[16];
int g_ri = 0;
simde__m128i take_root() { return g_roots[g_ri++]; }

struct PadA
{
    uint64_t n = 1;
    simde__m128i block()
    {
        auto v = simde_mm_set_epi64x(static_cast<long long>(n),
            static_cast<long long>(n * 9 + 3));
        n += 2;
        return v;
    }
    void fill(void * p, std::size_t nbytes)
    {
        auto * b = static_cast<unsigned char *>(p);
        for (std::size_t i = 0; i < nbytes; ++i)
            b[i] = static_cast<unsigned char>(n + i * 17);
        n += nbytes;
    }
    uint8_t bit() { return static_cast<uint8_t>(n++ & 1u); }
};

void reset_tape()
{
    g_ri = 0;
    for (int i = 0; i < 16; ++i)
        g_roots[i] = simde_mm_set_epi64x(0x11110000LL + i, 0x22220000LL + i * 3);
}

template <typename A, typename B>
auto recon(const A & a, const B & b)
{
    return dpf::reconstruct(a, b);
}

template <typename K0, typename K1, typename X>
uint64_t recon_cmp(const K0 & k0, const K1 & k1, X x)
{
    return recon(dpf::eval_point(dpf::cmp, k0, x),
        dpf::eval_point(dpf::cmp, k1, x)) & k0.cmp().mask;
}

template <std::size_t B, typename Spec>
void expect_u8_kind(uint8_t alpha, Spec spec, uint64_t below, uint64_t at,
    uint64_t above)
{
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::block_width<B>(std::move(spec)));
    using KT = std::decay_t<decltype(k0)>;
    EXPECT_GT(KT::cmp_block, 0u);
    EXPECT_EQ(KT::cmp_q, 2u);
    EXPECT_EQ(KT::cmp_h, 6u);
    EXPECT_EQ(k0.value_cw().size(), KT::cmp_checkpoints);
    EXPECT_EQ(k0.tail_cw().size(), KT::cmp_tail);
    EXPECT_EQ(KT::depth, KT::cmp_h);
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<uint8_t>(x);
        const uint64_t got = recon_cmp(k0, k1, q);
        const uint64_t want = q < alpha ? below : (q == alpha ? at : above);
        ASSERT_EQ(got, want) << "B=" << B << " x=" << x << " alpha=" << int(alpha);
    }
}

template <std::size_t B>
void exhaustive_lt(uint8_t alpha, uint64_t yt, uint64_t yf)
{
    expect_u8_kind<B>(alpha, dpf::lt(yt, yf), yt, yf, yf);
}

}  // namespace

TEST(BlockedDcf, ExhaustiveUint8AllWidths)
{
    const uint64_t yt = 9, yf = 0;
    for (uint8_t alpha : {uint8_t{0}, uint8_t{1}, uint8_t{7}, uint8_t{64},
            uint8_t{127}, uint8_t{200}, uint8_t{255}})
    {
        exhaustive_lt<1>(alpha, yt, yf);
        exhaustive_lt<2>(alpha, yt, yf);
        exhaustive_lt<3>(alpha, yt, yf);
        exhaustive_lt<4>(alpha, yt, yf);
    }
}

TEST(BlockedDcf, KindsIfFalseAndPayloadWidths)
{
    const uint8_t alpha = 40;
    expect_u8_kind<4>(alpha, dpf::lt(uint64_t{9}, uint64_t{2}), 9, 2, 2);
    expect_u8_kind<4>(alpha, dpf::leq(uint64_t{9}, uint64_t{2}), 9, 9, 2);
    expect_u8_kind<4>(alpha, dpf::gt(uint64_t{9}, uint64_t{2}), 2, 2, 9);
    expect_u8_kind<4>(alpha, dpf::geq(uint64_t{9}, uint64_t{2}), 2, 9, 9);
    expect_u8_kind<4>(alpha, dpf::lt(dpf::bit::one), 1, 0, 0);
    expect_u8_kind<2>(alpha, dpf::lt(uint16_t{7}, uint16_t{1}), 7, 1, 1);
}

TEST(BlockedDcf, Block1MatchesFunctionNotBytes)
{
    const uint8_t alpha = 30;
    auto bare = dpf::make_dpf(alpha, dpf::lt(uint64_t{4}));
    auto blocked = dpf::make_dpf(alpha, dpf::block_width<1>(dpf::lt(uint64_t{4})));
    using Bare = std::decay_t<decltype(bare.first)>;
    using Blk = std::decay_t<decltype(blocked.first)>;
    EXPECT_NE(Bare::depth, Blk::depth);
    EXPECT_NE(bare.first.value_cw().size(), blocked.first.value_cw().size());
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<uint8_t>(x);
        EXPECT_EQ(recon_cmp(bare.first, bare.second, q),
            recon_cmp(blocked.first, blocked.second, q));
    }
}

TEST(BlockedDcf, DeeperOutputForcesFullTree)
{
    const uint32_t alpha = 0x01020304u;
    auto [k0, k1] = dpf::make_dpf(alpha, uint32_t{11},
        dpf::block_width<4>(dpf::lt_at<8>(uint64_t{5}, uint64_t{1})));
    using KT = std::decay_t<decltype(k0)>;
    EXPECT_EQ(KT::cmp_q, 0u);
    EXPECT_EQ(KT::cmp_h, 8u);
    EXPECT_GT(KT::depth, KT::cmp_h);
    EXPECT_EQ(recon(*dpf::eval_point(k0, alpha), *dpf::eval_point(k1, alpha)), 11u);
    EXPECT_EQ(recon(*dpf::eval_point(k0, alpha ^ 1u),
                  *dpf::eval_point(k1, alpha ^ 1u)), 0u);
    auto top = [](uint32_t v) { return static_cast<uint8_t>(v >> 24); };
    auto cmp = [&](uint32_t q) { return recon_cmp(k0, k1, q); };
    EXPECT_EQ(cmp((uint32_t{top(alpha)} - 1u) << 24), 5u);
    EXPECT_EQ(cmp(alpha), 1u);
    EXPECT_EQ(cmp((uint32_t{top(alpha)} + 1u) << 24), 1u);
}

TEST(BlockedDcf, ShallowerLeafKeepsTail)
{
    const uint16_t alpha = 0x1234;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<4>(uint8_t{3}),
        dpf::block_width<4>(dpf::lt(uint64_t{8})));
    using KT = std::decay_t<decltype(k0)>;
    EXPECT_EQ(KT::cmp_q, 2u);
    EXPECT_EQ(recon(*dpf::eval_point(dpf::out<0, 4>, k0, alpha),
                  *dpf::eval_point(dpf::out<0, 4>, k1, alpha)), 3u);
    EXPECT_EQ(recon_cmp(k0, k1, uint16_t{alpha - 1}), 8u);
    EXPECT_EQ(recon_cmp(k0, k1, alpha), 0u);
}

TEST(BlockedDcf, PointIntervalFullSequenceInnerProduct)
{
    const uint8_t alpha = 100;
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::block_width<4>(dpf::lt(uint64_t{6}, uint64_t{1})));
    using KT = std::decay_t<decltype(k0)>;
    auto path0 = dpf::make_basic_path_memoizer(k0);
    auto path1 = dpf::make_basic_path_memoizer(k1);
    EXPECT_EQ(recon_cmp(k0, k1, uint8_t{50}), 6u);
    EXPECT_EQ(dpf::reconstruct(
            dpf::eval_point(dpf::cmp, k0, uint8_t{50}, path0),
            dpf::eval_point(dpf::cmp, k1, uint8_t{50}, path1)) & k0.cmp().mask,
        6u);
    EXPECT_EQ(dpf::reconstruct(
            dpf::eval_point(dpf::cmp, k0, uint8_t{150}, path0),
            dpf::eval_point(dpf::cmp, k1, uint8_t{150}, path1)) & k0.cmp().mask,
        1u);

    auto narrow0 = dpf::make_output_buffer(dpf::cmp, k0, uint8_t{90}, uint8_t{110});
    auto narrow1 = dpf::make_output_buffer(dpf::cmp, k1, uint8_t{90}, uint8_t{110});
    dpf::eval_interval(dpf::cmp, k0, uint8_t{90}, uint8_t{110}, narrow0);
    dpf::eval_interval(dpf::cmp, k1, uint8_t{90}, uint8_t{110}, narrow1);
    for (uint8_t x = 90; x <= 110; ++x)
    {
        EXPECT_EQ(recon(narrow0[x - 90], narrow1[x - 90]) & k0.cmp().mask,
            recon_cmp(k0, k1, x));
    }

    constexpr std::size_t stop = KT::cmp_depth;
    dpf::detail::incr::cmp_full_interval_memo<KT, stop> memo0{21};
    dpf::detail::incr::cmp_full_interval_memo<KT, stop> memo1{21};
    auto again0 = dpf::make_output_buffer(dpf::cmp, k0, uint8_t{90}, uint8_t{110});
    auto again1 = dpf::make_output_buffer(dpf::cmp, k1, uint8_t{90}, uint8_t{110});
    dpf::eval_interval(dpf::cmp, k0, uint8_t{90}, uint8_t{110}, again0, memo0);
    dpf::eval_interval(dpf::cmp, k1, uint8_t{90}, uint8_t{110}, again1, memo1);
    EXPECT_EQ(recon(again0[0], again1[0]) & k0.cmp().mask, 6u);
    EXPECT_EQ(recon(again0[10], again1[10]) & k0.cmp().mask, 1u);

    auto full0 = dpf::eval_full(dpf::cmp, k0);
    auto full1 = dpf::eval_full(dpf::cmp, k1);
    ASSERT_EQ(full0.size(), 256u);
    for (int x = 0; x < 256; ++x)
    {
        EXPECT_EQ(recon(full0[x], full1[x]) & k0.cmp().mask,
            recon_cmp(k0, k1, static_cast<uint8_t>(x)));
    }

    std::array<uint8_t, 4> pts{{0, 99, 100, 255}};
    auto seq0 = dpf::make_output_buffer(dpf::cmp, k0, pts.size());
    auto seq1 = dpf::make_output_buffer(dpf::cmp, k1, pts.size());
    dpf::eval_sequence(dpf::cmp, k0, pts.begin(), pts.end(), seq0, path0);
    dpf::eval_sequence(dpf::cmp, k1, pts.begin(), pts.end(), seq1, path1);
    for (std::size_t i = 0; i < pts.size(); ++i)
    {
        EXPECT_EQ(recon(seq0[i], seq1[i]) & k0.cmp().mask,
            recon_cmp(k0, k1, pts[i]));
    }

    std::vector<uint64_t> w(21, 1);
    const uint64_t dot =
        dpf::eval_inner_product(dpf::cmp, k0, uint8_t{90}, uint8_t{110}, w)
        + dpf::eval_inner_product(dpf::cmp, k1, uint8_t{90}, uint8_t{110}, w);
    uint64_t want = 0;
    for (uint8_t x = 90; x <= 110; ++x)
        want = (want + recon_cmp(k0, k1, x)) & k0.cmp().mask;
    EXPECT_EQ(dot & k0.cmp().mask, want);
}

TEST(BlockedDcf, DealerMatchesDoernerShelat)
{
    const uint8_t alpha = 77;
    const uint8_t x0 = 3;
    const uint8_t x1 = static_cast<uint8_t>(alpha ^ x0);
    reset_tape();
    auto dealer = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::block_width<4>(dpf::lt(uint64_t{15}, uint64_t{2})));
    reset_tape();
    dpf::ds_randomness<simde__m128i (*)(), PadA> rng{take_root, {}};
    auto ds = dpf::make_dpf_doerner_shelat(x0, x1, rng,
        dpf::block_width<4>(dpf::lt(uint64_t{15}, uint64_t{2})));

    EXPECT_EQ(std::memcmp(dealer.first.correction_words().data(),
                  ds.first.correction_words().data(),
                  dealer.first.correction_words().size()
                      * sizeof(dealer.first.correction_words()[0])),
        0);
    EXPECT_EQ(dealer.first.correction_advice(), ds.first.correction_advice());
    EXPECT_EQ(dealer.first.value_cw(), ds.first.value_cw());
    EXPECT_EQ(dealer.first.tail_cw(), ds.first.tail_cw());
    EXPECT_EQ(dealer.first.cw_last(), ds.first.cw_last());
    EXPECT_EQ(dealer.first.cmp_addend().raw(), ds.first.cmp_addend().raw());
    EXPECT_EQ(dealer.second.cmp_addend().raw(), ds.second.cmp_addend().raw());
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<uint8_t>(x);
        EXPECT_EQ(recon_cmp(dealer.first, dealer.second, q),
            recon_cmp(ds.first, ds.second, q));
    }
}

TEST(BlockedDcf, WildcardAssignMatchesConcrete)
{
    const uint8_t alpha = 19;
    const uint64_t yt = 42;
    reset_tape();
    auto wild = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::block_width<4>(dpf::lt(dpf::wildcard_value<uint64_t>{})));
    EXPECT_FALSE(wild.first.cmp_assigned());
    EXPECT_THROW(dpf::eval_point(dpf::cmp, wild.first, alpha), std::exception);
    reset_tape();
    auto concrete = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::block_width<4>(dpf::lt(yt)));
    dpf::assign_cmp(wild.first, wild.second, yt);
    EXPECT_TRUE(wild.first.cmp_assigned());
    EXPECT_EQ(wild.first.value_cw(), concrete.first.value_cw());
    EXPECT_EQ(wild.first.tail_cw(), concrete.first.tail_cw());
    EXPECT_EQ(wild.first.cw_last(), concrete.first.cw_last());
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<uint8_t>(x);
        EXPECT_EQ(recon_cmp(wild.first, wild.second, q),
            recon_cmp(concrete.first, concrete.second, q));
    }
}

TEST(BlockedDcf, TrivialDomainEdges)
{
    auto hi_gt = dpf::make_dpf(uint8_t{255},
        dpf::block_width<4>(dpf::gt(uint64_t{3})));
    EXPECT_EQ(hi_gt.first.cmp().trivial, dpf::cmp_trivial::always_false);
    EXPECT_EQ(recon_cmp(hi_gt.first, hi_gt.second, uint8_t{0}), 0u);
    EXPECT_EQ(recon_cmp(hi_gt.first, hi_gt.second, uint8_t{255}), 0u);

    auto hi_leq = dpf::make_dpf(uint8_t{255},
        dpf::block_width<4>(dpf::leq(uint64_t{3})));
    EXPECT_EQ(hi_leq.first.cmp().trivial, dpf::cmp_trivial::always_true);
    EXPECT_EQ(recon_cmp(hi_leq.first, hi_leq.second, uint8_t{0}), 3u);
    EXPECT_EQ(recon_cmp(hi_leq.first, hi_leq.second, uint8_t{255}), 3u);
}

TEST(BlockedDcf, JsonRoundTrip)
{
    const uint8_t alpha = 12;
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::block_width<4>(dpf::gt(uint64_t{7}, uint64_t{1})));
    using KT = typename std::decay_t<decltype(k0)>::key_type;
    const std::string s0 = dpf::json::to_json(k0.key());
    const std::string s1 = dpf::json::to_json(k1.key());
    EXPECT_NE(s0.find("block_width"), std::string::npos);
    EXPECT_NE(s0.find("tail_cw"), std::string::npos);
    auto r0 = dpf::json::from_json<KT>(s0);
    auto r1 = dpf::json::from_json<KT>(s1);
    EXPECT_EQ(r0.value_cw(), k0.value_cw());
    EXPECT_EQ(r0.tail_cw(), k0.tail_cw());
    EXPECT_EQ(r0.cmp().block_width, 4);
    EXPECT_EQ(r0.cmp().tail_bits, static_cast<int>(KT::cmp_q));
    const uint64_t mask = k0.cmp().mask;
    for (int x = 0; x < 256; ++x)
    {
        const auto q = static_cast<uint8_t>(x);
        const uint64_t got =
            (dpf::eval_point(dpf::cmp, r0, q) + dpf::eval_point(dpf::cmp, r1, q))
            & mask;
        EXPECT_EQ(got, recon_cmp(k0, k1, q));
    }
}

TEST(BlockedDcf, GenevalOpensCheckpointWords)
{
    const uint8_t alpha = 20;
    const uint8_t x0 = 1;
    const uint8_t x1 = static_cast<uint8_t>(alpha ^ x0);
    reset_tape();
    dpf::ds_randomness<simde__m128i (*)(), PadA> rng{take_root, {}};
    std::array<uint8_t, 3> ends{{0, 20, 21}};
    auto g = dpf::geneval_cmp(x0, x1, ends.begin(), ends.end(), rng,
        dpf::block_width<4>(dpf::lt(uint64_t{5})));
    using KT = std::decay_t<decltype(dpf::make_dpf(alpha,
        dpf::block_width<4>(dpf::lt(uint64_t{5}))).first)>;
    EXPECT_EQ(g.value_cw.size(), KT::cmp_checkpoints);
    EXPECT_EQ(g.tail_cw.size(), KT::cmp_tail);
    EXPECT_EQ(g.correction_words.size(), KT::depth);
    EXPECT_EQ((g.party0[0] + g.party1[0]) & g.mask, 5u);
    EXPECT_EQ((g.party0[1] + g.party1[1]) & g.mask, 0u);
    EXPECT_EQ((g.party0[2] + g.party1[2]) & g.mask, 0u);
}

TEST(BlockedDcf, GrottoPrefixSegmentAndHorner)
{
    const uint16_t alpha = 1000;
    std::array<uint16_t, 4> ends{{0, 500, 1000, 4000}};
    auto bare = dpf::make_dpf(alpha, dpf::gt(uint64_t{1}));
    auto blk = dpf::make_dpf(alpha, dpf::block_width<4>(dpf::gt(uint64_t{1})));
    const auto b0 = grotto::signed_prefix_parities(bare.first, ends);
    const auto b1 = grotto::signed_prefix_parities(bare.second, ends);
    const auto k0 = grotto::signed_prefix_parities(blk.first, ends);
    const auto k1 = grotto::signed_prefix_parities(blk.second, ends);
    const uint64_t maskp = bare.first.cmp().mask;
    for (std::size_t i = 0; i < ends.size(); ++i)
        EXPECT_EQ((b0[i] + b1[i]) & maskp, (k0[i] + k1[i]) & maskp);

    const auto s0 = grotto::signed_segment_parities(bare.first, ends);
    const auto s1 = grotto::signed_segment_parities(bare.second, ends);
    const auto t0 = grotto::signed_segment_parities(blk.first, ends);
    const auto t1 = grotto::signed_segment_parities(blk.second, ends);
    for (std::size_t i = 0; i < ends.size(); ++i)
        EXPECT_EQ((s0[i] + s1[i]) & maskp, (t0[i] + t1[i]) & maskp);

    const uint16_t center = 30;
    auto mat = grotto::make_offset_horner_keys<uint16_t, 1>(center);
    uint64_t payload[2];
    payload[0] = 1;
    payload[1] = center;
    using bare_pair = decltype(dpf::make_dpf(center, dpf::gt(uint64_t{0})));
    using pair = decltype(dpf::make_dpf(center,
        dpf::block_width<4>(dpf::gt(uint64_t{0}))));
    std::vector<bare_pair> bare_keys{
        dpf::make_dpf(center, dpf::gt(payload[0])),
        dpf::make_dpf(center, dpf::gt(payload[1]))};
    std::vector<pair> keys{
        dpf::make_dpf(center, dpf::block_width<4>(dpf::gt(payload[0]))),
        dpf::make_dpf(center, dpf::block_width<4>(dpf::gt(payload[1])))};
    std::vector<uint16_t> knots{0, 10, 40};
    std::vector<std::array<uint64_t, 2>> coeff{
        {1, 0},
        {2, 3},
        {4, 1}};
    const uint16_t eta = 0;
    const auto b0s = grotto::offset_horner_coefficient_share<0, 1>(
        bare_keys, mat.wrap_share, knots, coeff, eta);
    const auto b1s = grotto::offset_horner_coefficient_share<1, 1>(
        bare_keys, mat.wrap_share, knots, coeff, eta);
    const auto c0 = grotto::offset_horner_coefficient_share<0, 1>(
        keys, mat.wrap_share, knots, coeff, eta);
    const auto c1 = grotto::offset_horner_coefficient_share<1, 1>(
        keys, mat.wrap_share, knots, coeff, eta);
    for (std::size_t m = 0; m < 2; ++m)
        EXPECT_EQ(b0s[m] + b1s[m], c0[m] + c1[m]);
}